rust_sasa/
options.rs

1// Copyright (c) 2024 Maxwell Campbell. Licensed under the MIT License.
2use crate::structures::atomic::{ChainResult, ProteinResult, ResidueResult};
3use crate::utils::consts::{POLAR_AMINO_ACIDS, load_radii_from_file};
4use crate::utils::{combine_hash, get_radius, serialize_chain_id, simd_sum};
5use crate::{Atom, calculate_sasa_internal};
6use fnv::FnvHashMap;
7use pdbtbx::PDB;
8use snafu::OptionExt;
9use snafu::prelude::*;
10use std::marker::PhantomData;
11
12/// Options for configuring SASA (Solvent Accessible Surface Area) calculations.
13///
14/// This struct provides configuration options for SASA calculations at different levels
15/// of granularity (atom, residue, chain, or protein level). The type parameter `T`
16/// determines the output type and processing behavior.
17///
18/// # Type Parameters
19///
20/// * `T` - The processing level, which must implement [`SASAProcessor`]. Available levels:
21///   - [`AtomLevel`] - Returns SASA values for individual atoms
22///   - [`ResidueLevel`] - Returns SASA values aggregated by residue
23///   - [`ChainLevel`] - Returns SASA values aggregated by chain
24///   - [`ProteinLevel`] - Returns SASA values aggregated for the entire protein
25///
26/// # Fields
27///
28/// * `probe_radius` - Radius of the solvent probe sphere in Angstroms (default: 1.4)
29/// * `n_points` - Number of points on the sphere surface for sampling (default: 100)
30/// * `threads` - Number of threads to use for parallel processing (default: -1 for all cores)
31/// * `include_hydrogens` - Whether to include hydrogen atoms in calculations (default: false)
32/// * `radii_config` - Optional custom radii configuration (default: uses embedded protor.config)
33/// * `allow_vdw_fallback` - Allow fallback to PDBTBX van der Waals radii when radius is not found in radii file (default: false)
34/// * `include_hetatms` - Whether to include HETATM records (e.g. non-standard amino acids) in calculations (default: false)
35///
36/// # Examples
37///
38/// ```rust
39/// use rust_sasa::options::{SASAOptions, ResidueLevel};
40/// use pdbtbx::PDB;
41///
42/// // Create options with default settings
43/// let options = SASAOptions::<ResidueLevel>::new();
44///
45/// // Customize the configuration
46/// let custom_options = SASAOptions::<ResidueLevel>::new()
47///     .with_probe_radius(1.2)
48///     .with_n_points(200)
49///     .with_threads(-1)
50///     .with_include_hydrogens(false)
51///     .with_allow_vdw_fallback(true)
52///     .with_include_hetatms(false);
53///
54/// // Process a PDB structure
55/// # let pdb = PDB::new();
56/// let result = custom_options.process(&pdb)?;
57/// # Ok::<(), Box<dyn std::error::Error>>(())
58/// ```
59#[derive(Debug, Clone)]
60pub struct SASAOptions<T> {
61    probe_radius: f32,
62    n_points: usize,
63    threads: isize,
64    include_hydrogens: bool,
65    radii_config: Option<FnvHashMap<String, FnvHashMap<String, f32>>>,
66    allow_vdw_fallback: bool,
67    include_hetatms: bool,
68    read_radii_from_occupancy: bool,
69    _marker: PhantomData<T>,
70}
71
72// Zero-sized marker types for each level
73pub struct AtomLevel;
74pub struct ResidueLevel;
75pub struct ChainLevel;
76pub struct ProteinLevel;
77
78pub type AtomsMappingResult = Result<(Vec<Atom>, FnvHashMap<isize, Vec<usize>>), SASACalcError>;
79
80/// Macro to reduce duplication in atom building logic
81macro_rules! build_atom {
82    ($atoms:expr, $atom:expr, $element:expr, $residue_name:expr, $atom_name:expr, $parent_id:expr, $radii_config:expr, $allow_vdw_fallback:expr, $read_radii_from_occupancy:expr, $id:expr) => {{
83        let radius = if $read_radii_from_occupancy {
84            $atom.occupancy() as f32
85        } else {
86            match get_radius($residue_name, $atom_name, $radii_config) {
87                Some(r) => r,
88                None => {
89                    if $allow_vdw_fallback {
90                        $element
91                            .atomic_radius()
92                            .van_der_waals
93                            .context(VanDerWaalsMissingSnafu)? as f32
94                    } else {
95                        return Err(SASACalcError::RadiusMissing {
96                            residue_name: $residue_name.to_string(),
97                            atom_name: $atom_name.to_string(),
98                            element: $element.to_string(),
99                        });
100                    }
101                }
102            }
103        };
104
105        $atoms.push(Atom {
106            position: [
107                $atom.pos().0 as f32,
108                $atom.pos().1 as f32,
109                $atom.pos().2 as f32,
110            ],
111            radius,
112            id: $id as usize,
113            parent_id: $parent_id,
114        });
115    }};
116}
117
118// Trait that defines the processing behavior for each level
119pub trait SASAProcessor {
120    type Output;
121
122    fn process_atoms(
123        atoms: &[Atom],
124        atom_sasa: &[f32],
125        pdb: &PDB,
126        parent_to_atoms: &FnvHashMap<isize, Vec<usize>>,
127    ) -> Result<Self::Output, SASACalcError>;
128
129    fn build_atoms_and_mapping(
130        pdb: &PDB,
131        radii_config: Option<&FnvHashMap<String, FnvHashMap<String, f32>>>,
132        allow_vdw_fallback: bool,
133        include_hydrogens: bool,
134        include_hetatms: bool,
135        read_radii_from_occupancy: bool,
136    ) -> AtomsMappingResult;
137}
138
139impl SASAProcessor for AtomLevel {
140    type Output = Vec<f32>;
141
142    fn process_atoms(
143        _atoms: &[Atom],
144        atom_sasa: &[f32],
145        _pdb: &PDB,
146        _parent_to_atoms: &FnvHashMap<isize, Vec<usize>>,
147    ) -> Result<Self::Output, SASACalcError> {
148        Ok(atom_sasa.to_vec())
149    }
150
151    fn build_atoms_and_mapping(
152        pdb: &PDB,
153        radii_config: Option<&FnvHashMap<String, FnvHashMap<String, f32>>>,
154        allow_vdw_fallback: bool,
155        include_hydrogens: bool,
156        include_hetatms: bool,
157        read_radii_from_occupancy: bool,
158    ) -> Result<(Vec<Atom>, FnvHashMap<isize, Vec<usize>>), SASACalcError> {
159        let mut atoms = vec![];
160        for residue in pdb.residues() {
161            let residue_name = residue.name().context(FailedToGetResidueNameSnafu)?;
162            if let Some(conformer) = residue.conformers().next() {
163                for atom in conformer.atoms() {
164                    let element = atom.element().context(ElementMissingSnafu)?;
165                    let atom_name = atom.name();
166                    if element == &pdbtbx::Element::H && !include_hydrogens {
167                        continue;
168                    };
169                    if atom.hetero() && !include_hetatms {
170                        continue;
171                    }
172                    let conformer_alt = conformer.alternative_location().unwrap_or("");
173                    build_atom!(
174                        atoms,
175                        atom,
176                        element,
177                        residue_name,
178                        atom_name,
179                        None,
180                        radii_config,
181                        allow_vdw_fallback,
182                        read_radii_from_occupancy,
183                        combine_hash(&(conformer_alt, atom.serial_number()))
184                    );
185                }
186            }
187        }
188        Ok((atoms, FnvHashMap::default()))
189    }
190}
191
192impl SASAProcessor for ResidueLevel {
193    type Output = Vec<ResidueResult>;
194
195    fn process_atoms(
196        _atoms: &[Atom],
197        atom_sasa: &[f32],
198        pdb: &PDB,
199        parent_to_atoms: &FnvHashMap<isize, Vec<usize>>,
200    ) -> Result<Self::Output, SASACalcError> {
201        let mut residue_sasa = vec![];
202        for chain in pdb.chains() {
203            for residue in chain.residues() {
204                let residue_key = combine_hash(&(
205                    chain.id(),
206                    residue.serial_number(),
207                    residue.insertion_code().unwrap_or_default(),
208                ));
209                let residue_atom_index = parent_to_atoms
210                    .get(&residue_key)
211                    .context(AtomMapToLevelElementFailedSnafu)?;
212                let residue_atoms: Vec<_> = residue_atom_index
213                    .iter()
214                    .map(|&index| atom_sasa[index])
215                    .collect();
216                let sum = simd_sum(residue_atoms.as_slice());
217                let name = residue
218                    .name()
219                    .context(FailedToGetResidueNameSnafu)?
220                    .to_string();
221                residue_sasa.push(ResidueResult {
222                    serial_number: residue.serial_number(),
223                    value: sum,
224                    is_polar: POLAR_AMINO_ACIDS.contains(&name),
225                    chain_id: chain.id().to_string(),
226                    name,
227                })
228            }
229        }
230        Ok(residue_sasa)
231    }
232
233    fn build_atoms_and_mapping(
234        pdb: &PDB,
235        radii_config: Option<&FnvHashMap<String, FnvHashMap<String, f32>>>,
236        allow_vdw_fallback: bool,
237        include_hydrogens: bool,
238        include_hetatms: bool,
239        read_radii_from_occupancy: bool,
240    ) -> Result<(Vec<Atom>, FnvHashMap<isize, Vec<usize>>), SASACalcError> {
241        let mut atoms = vec![];
242        let mut parent_to_atoms = FnvHashMap::default();
243        let mut i = 0;
244        for chain in pdb.chains() {
245            let chain_id = chain.id();
246            for residue in chain.residues() {
247                let residue_name = residue.name().context(FailedToGetResidueNameSnafu)?;
248                let residue_key = combine_hash(&(
249                    chain_id,
250                    residue.serial_number(),
251                    residue.insertion_code().unwrap_or_default(),
252                ));
253                let mut temp = vec![];
254                if let Some(conformer) = residue.conformers().next() {
255                    for atom in conformer.atoms() {
256                        let element = atom.element().context(ElementMissingSnafu)?;
257                        let atom_name = atom.name();
258                        if element == &pdbtbx::Element::H && !include_hydrogens {
259                            continue;
260                        };
261                        if atom.hetero() && !include_hetatms {
262                            continue;
263                        }
264                        let conformer_alt = conformer.alternative_location().unwrap_or("");
265                        build_atom!(
266                            atoms,
267                            atom,
268                            element,
269                            residue_name,
270                            atom_name,
271                            Some(residue.serial_number()),
272                            radii_config,
273                            allow_vdw_fallback,
274                            read_radii_from_occupancy,
275                            combine_hash(&(conformer_alt, atom.serial_number()))
276                        );
277                        temp.push(i);
278                        i += 1;
279                    }
280                    parent_to_atoms.insert(residue_key, temp);
281                }
282            }
283        }
284        Ok((atoms, parent_to_atoms))
285    }
286}
287
288impl SASAProcessor for ChainLevel {
289    type Output = Vec<ChainResult>;
290
291    fn process_atoms(
292        _atoms: &[Atom],
293        atom_sasa: &[f32],
294        pdb: &PDB,
295        parent_to_atoms: &FnvHashMap<isize, Vec<usize>>,
296    ) -> Result<Self::Output, SASACalcError> {
297        let mut chain_sasa = vec![];
298        for chain in pdb.chains() {
299            let chain_id = serialize_chain_id(chain.id());
300            let chain_atom_index = parent_to_atoms
301                .get(&chain_id)
302                .context(AtomMapToLevelElementFailedSnafu)?;
303            let chain_atoms: Vec<_> = chain_atom_index
304                .iter()
305                .map(|&index| atom_sasa[index])
306                .collect();
307            let sum = simd_sum(chain_atoms.as_slice());
308            chain_sasa.push(ChainResult {
309                name: chain.id().to_string(),
310                value: sum,
311            })
312        }
313        Ok(chain_sasa)
314    }
315
316    fn build_atoms_and_mapping(
317        pdb: &PDB,
318        radii_config: Option<&FnvHashMap<String, FnvHashMap<String, f32>>>,
319        allow_vdw_fallback: bool,
320        include_hydrogens: bool,
321        include_hetatms: bool,
322        read_radii_from_occupancy: bool,
323    ) -> Result<(Vec<Atom>, FnvHashMap<isize, Vec<usize>>), SASACalcError> {
324        let mut atoms = vec![];
325        let mut parent_to_atoms = FnvHashMap::default();
326        let mut i = 0;
327        for chain in pdb.chains() {
328            let chain_id = serialize_chain_id(chain.id());
329            let mut temp = vec![];
330            for residue in chain.residues() {
331                let residue_name = residue.name().context(FailedToGetResidueNameSnafu)?;
332                if let Some(conformer) = residue.conformers().next() {
333                    for atom in conformer.atoms() {
334                        let element = atom.element().context(ElementMissingSnafu)?;
335                        let atom_name = atom.name();
336                        let conformer_alt = conformer.alternative_location().unwrap_or("");
337                        if element == &pdbtbx::Element::H && !include_hydrogens {
338                            continue;
339                        };
340                        if atom.hetero() && !include_hetatms {
341                            continue;
342                        }
343                        build_atom!(
344                            atoms,
345                            atom,
346                            element,
347                            residue_name,
348                            atom_name,
349                            Some(chain_id),
350                            radii_config,
351                            allow_vdw_fallback,
352                            read_radii_from_occupancy,
353                            combine_hash(&(conformer_alt, atom.serial_number()))
354                        );
355                        temp.push(i);
356                        i += 1
357                    }
358                }
359            }
360            parent_to_atoms.insert(chain_id, temp);
361        }
362        Ok((atoms, parent_to_atoms))
363    }
364}
365
366impl SASAProcessor for ProteinLevel {
367    type Output = ProteinResult;
368
369    fn process_atoms(
370        _atoms: &[Atom],
371        atom_sasa: &[f32],
372        pdb: &PDB,
373        parent_to_atoms: &FnvHashMap<isize, Vec<usize>>,
374    ) -> Result<Self::Output, SASACalcError> {
375        let mut polar_total: f32 = 0.0;
376        let mut non_polar_total: f32 = 0.0;
377        for chain in pdb.chains() {
378            for residue in chain.residues() {
379                let residue_key = combine_hash(&(
380                    chain.id(),
381                    residue.serial_number(),
382                    residue.insertion_code().unwrap_or_default(),
383                ));
384                let residue_atom_index = parent_to_atoms
385                    .get(&residue_key)
386                    .context(AtomMapToLevelElementFailedSnafu)?;
387                let residue_atoms: Vec<_> = residue_atom_index
388                    .iter()
389                    .map(|&index| atom_sasa[index])
390                    .collect();
391                let sum = simd_sum(residue_atoms.as_slice());
392                let name = residue
393                    .name()
394                    .context(FailedToGetResidueNameSnafu)?
395                    .to_string();
396                if POLAR_AMINO_ACIDS.contains(&name) {
397                    polar_total += sum
398                } else {
399                    non_polar_total += sum
400                }
401            }
402        }
403        let global_sum = simd_sum(atom_sasa);
404        Ok(ProteinResult {
405            global_total: global_sum,
406            polar_total,
407            non_polar_total,
408        })
409    }
410
411    fn build_atoms_and_mapping(
412        pdb: &PDB,
413        radii_config: Option<&FnvHashMap<String, FnvHashMap<String, f32>>>,
414        allow_vdw_fallback: bool,
415        include_hydrogens: bool,
416        include_hetatms: bool,
417        read_radii_from_occupancy: bool,
418    ) -> Result<(Vec<Atom>, FnvHashMap<isize, Vec<usize>>), SASACalcError> {
419        let mut atoms = vec![];
420        let mut parent_to_atoms = FnvHashMap::default();
421        let mut i = 0;
422        for chain in pdb.chains() {
423            let chain_id = chain.id();
424            for residue in chain.residues() {
425                let residue_name = residue.name().context(FailedToGetResidueNameSnafu)?;
426                let residue_key = combine_hash(&(
427                    chain_id,
428                    residue.serial_number(),
429                    residue.insertion_code().unwrap_or_default(),
430                ));
431                let mut temp = vec![];
432                if let Some(conformer) = residue.conformers().next() {
433                    for atom in conformer.atoms() {
434                        let element = atom.element().context(ElementMissingSnafu)?;
435                        let atom_name = atom.name();
436                        if element == &pdbtbx::Element::H && !include_hydrogens {
437                            continue;
438                        };
439                        if atom.hetero() && !include_hetatms {
440                            continue;
441                        }
442                        build_atom!(
443                            atoms,
444                            atom,
445                            element,
446                            residue_name,
447                            atom_name,
448                            Some(residue.serial_number()),
449                            radii_config,
450                            allow_vdw_fallback,
451                            read_radii_from_occupancy,
452                            combine_hash(&("", atom.serial_number()))
453                        );
454                        temp.push(i);
455                        i += 1;
456                    }
457                    parent_to_atoms.insert(residue_key, temp);
458                }
459            }
460        }
461        Ok((atoms, parent_to_atoms))
462    }
463}
464
465#[derive(Debug, Snafu)]
466pub enum SASACalcError {
467    #[snafu(display("Element missing for atom"))]
468    ElementMissing,
469
470    #[snafu(display("Van der Waals radius missing for element"))]
471    VanDerWaalsMissing,
472
473    #[snafu(display(
474        "Radius not found for residue '{}' atom '{}' of type '{}'. This error can can be ignored, if you are using the CLI pass --allow-vdw-fallback or use with_allow_vdw_fallback if you are using the API.",
475        residue_name,
476        atom_name,
477        element
478    ))]
479    RadiusMissing {
480        residue_name: String,
481        atom_name: String,
482        element: String,
483    },
484
485    #[snafu(display("Failed to map atoms back to level element"))]
486    AtomMapToLevelElementFailed,
487
488    #[snafu(display("Failed to get residue name"))]
489    FailedToGetResidueName,
490
491    #[snafu(display("Failed to load radii file: {source}"))]
492    RadiiFileLoad { source: std::io::Error },
493}
494
495impl<T> SASAOptions<T> {
496    /// Create a new SASAOptions with the specified level type
497    pub fn new() -> SASAOptions<T> {
498        SASAOptions {
499            probe_radius: 1.4,
500            n_points: 100,
501            threads: -1,
502            include_hydrogens: false,
503            radii_config: None,
504            allow_vdw_fallback: false,
505            include_hetatms: false,
506            read_radii_from_occupancy: false,
507            _marker: PhantomData,
508        }
509    }
510
511    /// Set the probe radius (default: 1.4 Angstroms)
512    pub fn with_probe_radius(mut self, radius: f32) -> Self {
513        self.probe_radius = radius;
514        self
515    }
516
517    /// Include or exclude HETATM records in protein.
518    pub fn with_include_hetatms(mut self, include_hetatms: bool) -> Self {
519        self.include_hetatms = include_hetatms;
520        self
521    }
522
523    /// Set the number of points on the sphere for sampling (default: 100)
524    pub fn with_n_points(mut self, points: usize) -> Self {
525        self.n_points = points;
526        self
527    }
528
529    /// Set whether radii should be read from input protein occupancy values. (default: false)
530    pub fn with_read_radii_from_occupancy(mut self, read_radii_from_occupancy: bool) -> Self {
531        self.read_radii_from_occupancy = read_radii_from_occupancy;
532        self
533    }
534
535    /// Configure the number of threads to use for parallel processing
536    ///   - `-1`: Use all available CPU cores (default)
537    ///   - `1`: Single-threaded execution (disables parallelism)
538    ///   - `> 1`: Use specified number of threads
539    pub fn with_threads(mut self, threads: isize) -> Self {
540        self.threads = threads;
541        self
542    }
543
544    /// Include or exclude hydrogen atoms in calculations (default: false)
545    pub fn with_include_hydrogens(mut self, include_hydrogens: bool) -> Self {
546        self.include_hydrogens = include_hydrogens;
547        self
548    }
549
550    /// Load custom radii configuration from a file (default: uses embedded protor.config)
551    pub fn with_radii_file(mut self, path: &str) -> Result<Self, std::io::Error> {
552        self.radii_config = Some(load_radii_from_file(path)?);
553        Ok(self)
554    }
555
556    /// Allow fallback to PDBTBX van der Waals radii when radius is not found in radii config file (default: false)
557    pub fn with_allow_vdw_fallback(mut self, allow: bool) -> Self {
558        self.allow_vdw_fallback = allow;
559        self
560    }
561}
562
563// Convenience constructors for each level
564impl SASAOptions<AtomLevel> {
565    pub fn atom_level() -> Self {
566        Self::new()
567    }
568}
569
570impl SASAOptions<ResidueLevel> {
571    pub fn residue_level() -> Self {
572        Self::new()
573    }
574}
575
576impl SASAOptions<ChainLevel> {
577    pub fn chain_level() -> Self {
578        Self::new()
579    }
580}
581
582impl SASAOptions<ProteinLevel> {
583    pub fn protein_level() -> Self {
584        Self::new()
585    }
586}
587
588impl<T> Default for SASAOptions<T> {
589    fn default() -> Self {
590        Self::new()
591    }
592}
593
594impl<T: SASAProcessor> SASAOptions<T> {
595    /// This function calculates the SASA for a given protein. The output type is determined by the level type parameter.
596    /// Probe radius and n_points can be customized, defaulting to 1.4 and 100 respectively.
597    /// If you want more fine-grained control you may want to use [calculate_sasa_internal] instead.
598    /// ## Example
599    /// ```
600    /// use pdbtbx::StrictnessLevel;
601    /// use rust_sasa::options::{SASAOptions, ResidueLevel};
602    /// let (mut pdb, _errors) = pdbtbx::open("./tests/data/pdbs/example.cif").unwrap();
603    /// let result = SASAOptions::<ResidueLevel>::new().process(&pdb);
604    /// ```
605    pub fn process(&self, pdb: &PDB) -> Result<T::Output, SASACalcError> {
606        let (atoms, parent_to_atoms) = T::build_atoms_and_mapping(
607            pdb,
608            self.radii_config.as_ref(),
609            self.allow_vdw_fallback,
610            self.include_hydrogens,
611            self.include_hetatms,
612            self.read_radii_from_occupancy,
613        )?;
614        let atom_sasa =
615            calculate_sasa_internal(&atoms, self.probe_radius, self.n_points, self.threads);
616        T::process_atoms(&atoms, &atom_sasa, pdb, &parent_to_atoms)
617    }
618}
rust_sasa/options.rs

rust_sasa/
options.rs
