rust_sasa/
utils.rs

1pub mod consts;
2pub mod io;
3use std::sync::LazyLock;
4
5use fnv::FnvHashMap;
6use pulp::Arch;
7use rayon::ThreadPoolBuilder;
8
9use crate::utils::consts::get_protor_radius;
10
11pub(crate) static ARCH: LazyLock<Arch> = LazyLock::new(Arch::new);
12
13pub(crate) fn simd_sum(values: &[f32]) -> f32 {
14    let mut total = 0f32;
15    ARCH.dispatch(|| {
16        for x in values {
17            total += x;
18        }
19    });
20    total
21}
22
23pub(crate) fn serialize_chain_id(s: &str) -> isize {
24    let mut result = 0;
25    for c in s.chars() {
26        if c.is_ascii_alphabetic() {
27            let position = c.to_ascii_uppercase() as isize - 64;
28            result = result * 10 + position;
29        }
30    }
31    result
32}
33
34/// Helper function to get atomic radius from custom config or default protor config
35pub fn get_radius(
36    residue_name: &str,
37    atom_name: &str,
38    radii_config: Option<&FnvHashMap<String, FnvHashMap<String, f32>>>,
39) -> Option<f32> {
40    // Check custom config first
41    if let Some(config) = radii_config {
42        if let Some(radius) = config
43            .get(residue_name)
44            .and_then(|inner| inner.get(atom_name))
45        {
46            return Some(*radius);
47        }
48    }
49    // Fall back to default protor config
50    get_protor_radius(residue_name, atom_name)
51}
52
53/// Configure the global rayon thread pool based on the threads argument
54///   - `-1`: Use all available CPU cores (default rayon behavior)
55///   - `1`: Single-threaded execution
56///   - `> 1`: Use specified number of threads
57///   - `0`: Invalid returns error
58pub fn configure_thread_pool(threads: isize) -> Result<(), std::io::Error> {
59    if threads == 0 {
60        return Err(std::io::Error::new(
61            std::io::ErrorKind::InvalidInput,
62            "Thread count must be -1 (all cores) or a positive number",
63        ));
64    }
65
66    // Configures global thread pool
67    if threads > 0 {
68        ThreadPoolBuilder::new()
69            .num_threads(threads as usize)
70            .build_global()
71            .map_err(|e| std::io::Error::other(format!("Failed to configure thread pool: {e}")))?;
72    }
73    // If threads == -1, use default rayon behavior (all cores)
74
75    Ok(())
76}