ferritin_core/model/
hierarchy.rs

1//! Atomic hierarchy: topology layer shared across trajectory frames.
2//!
3//! [`AtomicHierarchy`] holds all structural data that does not change between
4//! trajectory frames (connectivity, sequence, chain/residue assignments).
5//! It is wrapped in [`std::sync::Arc`] so multiple [`super::model::Model`]s
6//! (frames) can share a single topology without cloning.
7
8use super::bonds::Bonds;
9use super::error::ModelError;
10use super::tables::{AtomsTable, ChainsTable, ResiduesTable};
11use crate::data::Segmentation;
12use std::ops::Range;
13
14/// Topology layer: all structural data that is constant across trajectory frames.
15///
16/// Wrap in `Arc<AtomicHierarchy>` so multiple `Model`s (frames) share one topology
17/// without duplication.
18///
19/// # Hierarchy
20///
21/// ```text
22/// Chain  ──┐
23///           ├── Residue ──┐
24///                          └── Atom
25/// ```
26///
27/// The `atom_to_residue` and `residue_to_chain` segmentations provide O(1)
28/// range queries and O(log n) reverse lookups.
29#[derive(Clone, Debug)]
30pub struct AtomicHierarchy {
31    /// Per-atom topology data.
32    pub atoms: AtomsTable,
33    /// Per-residue topology data.
34    pub residues: ResiduesTable,
35    /// Per-chain topology data.
36    pub chains: ChainsTable,
37    /// Segmentation mapping atom index → residue index.
38    ///
39    /// `atom_to_residue.segment(res_idx)` gives the range of atom indices
40    /// belonging to residue `res_idx`.
41    pub atom_to_residue: Segmentation,
42    /// Segmentation mapping residue index → chain index.
43    ///
44    /// `residue_to_chain.segment(chain_idx)` gives the range of residue indices
45    /// belonging to chain `chain_idx`.
46    pub residue_to_chain: Segmentation,
47    /// Bond connectivity.
48    pub bonds: Bonds,
49}
50
51impl AtomicHierarchy {
52    /// Validate table lengths and hierarchy segmentations.
53    pub fn validate(&self) -> Result<(), ModelError> {
54        self.atoms.validate()?;
55        self.residues.validate()?;
56        self.chains.validate()?;
57        if self.atom_to_residue.count() != self.residues.len()
58            || self.atom_to_residue.element_count() != self.atoms.len()
59        {
60            return Err(ModelError::new(format!(
61                "atom_to_residue describes {} residues and {} atoms, expected {} residues and {} atoms",
62                self.atom_to_residue.count(),
63                self.atom_to_residue.element_count(),
64                self.residues.len(),
65                self.atoms.len()
66            )));
67        }
68        if self.residue_to_chain.count() != self.chains.len()
69            || self.residue_to_chain.element_count() != self.residues.len()
70        {
71            return Err(ModelError::new(format!(
72                "residue_to_chain describes {} chains and {} residues, expected {} chains and {} residues",
73                self.residue_to_chain.count(),
74                self.residue_to_chain.element_count(),
75                self.chains.len(),
76                self.residues.len()
77            )));
78        }
79        if self.bonds.atom_bond_starts.len() != self.atoms.len() + 1 {
80            return Err(ModelError::new(format!(
81                "bond CSR index has length {}, expected {}",
82                self.bonds.atom_bond_starts.len(),
83                self.atoms.len() + 1
84            )));
85        }
86        if self.bonds.atom_a.len() != self.bonds.atom_b.len()
87            || self.bonds.atom_a.len() != self.bonds.order.len()
88        {
89            return Err(ModelError::new("bond columns have inconsistent lengths"));
90        }
91        if self
92            .bonds
93            .atom_a
94            .iter()
95            .chain(&self.bonds.atom_b)
96            .any(|&atom| atom as usize >= self.atoms.len())
97        {
98            return Err(ModelError::new("bond endpoint is out of atom range"));
99        }
100        if self.bonds.atom_bond_starts.windows(2).any(|w| w[0] > w[1])
101            || self.bonds.atom_bond_starts.last().copied().unwrap_or(0) as usize
102                != self.bonds.atom_a.len()
103        {
104            return Err(ModelError::new(
105                "bond CSR index is inconsistent with bond columns",
106            ));
107        }
108        Ok(())
109    }
110    /// Total number of atoms.
111    pub fn n_atoms(&self) -> usize {
112        self.atoms.len()
113    }
114
115    /// Total number of residues.
116    pub fn n_residues(&self) -> usize {
117        self.residues.len()
118    }
119
120    /// Total number of chains.
121    pub fn n_chains(&self) -> usize {
122        self.chains.len()
123    }
124
125    /// Returns the range of atom indices belonging to residue `res_idx`.
126    ///
127    /// # Panics
128    /// Panics if `res_idx >= n_residues()`.
129    pub fn atoms_in_residue(&self, res_idx: usize) -> Range<usize> {
130        self.atom_to_residue.segment(res_idx)
131    }
132
133    /// Returns the range of residue indices belonging to chain `chain_idx`.
134    ///
135    /// # Panics
136    /// Panics if `chain_idx >= n_chains()`.
137    pub fn residues_in_chain(&self, chain_idx: usize) -> Range<usize> {
138        self.residue_to_chain.segment(chain_idx)
139    }
140
141    /// Returns the residue index that contains atom `atom_idx`.
142    ///
143    /// O(log n) binary search over residue boundaries.
144    ///
145    /// # Panics
146    /// Panics if `atom_idx >= n_atoms()`.
147    pub fn residue_of_atom(&self, atom_idx: usize) -> usize {
148        self.atom_to_residue.segment_of(atom_idx)
149    }
150
151    /// Returns the chain index that contains residue `res_idx`.
152    ///
153    /// O(log n) binary search over chain boundaries.
154    ///
155    /// # Panics
156    /// Panics if `res_idx >= n_residues()`.
157    pub fn chain_of_residue(&self, res_idx: usize) -> usize {
158        self.residue_to_chain.segment_of(res_idx)
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use crate::info::elements::Element;
166    use crate::model::tables::ResidueGroup;
167
168    /// Build a test hierarchy: 3 chains, 10 residues (4+3+3), 30 atoms (3 per residue).
169    fn make_test_hierarchy() -> AtomicHierarchy {
170        // Chain 0: residues 0..4, Chain 1: residues 4..7, Chain 2: residues 7..10
171        // Each residue has 3 atoms.
172        let n_residues = 10;
173        let n_atoms = 30; // 3 per residue
174        let _n_chains = 3;
175
176        let atoms = AtomsTable {
177            atom_name: (0..n_atoms)
178                .map(|i| {
179                    match i % 3 {
180                        0 => "N",
181                        1 => "CA",
182                        _ => "C",
183                    }
184                    .to_string()
185                })
186                .collect(),
187            auth_atom_name: (0..n_atoms)
188                .map(|i| {
189                    match i % 3 {
190                        0 => "N",
191                        1 => "CA",
192                        _ => "C",
193                    }
194                    .to_string()
195                })
196                .collect(),
197            element: (0..n_atoms)
198                .map(|i| if i % 3 == 0 { Element::N } else { Element::C })
199                .collect(),
200            alt_loc: vec![None; n_atoms],
201            formal_charge: vec![None; n_atoms],
202        };
203
204        let residues = ResiduesTable {
205            comp_id: (0..n_residues).map(|i| format!("RES{}", i)).collect(),
206            auth_comp_id: (0..n_residues).map(|i| format!("RES{}", i)).collect(),
207            label_seq_id: (0..n_residues as i32).collect(),
208            auth_seq_id: (1..=n_residues as i32).collect(),
209            ins_code: vec![None; n_residues],
210            group: vec![ResidueGroup::Polymer; n_residues],
211        };
212
213        let chains = ChainsTable {
214            label_asym_id: vec!["A".into(), "B".into(), "C".into()],
215            auth_asym_id: vec!["A".into(), "B".into(), "C".into()],
216            entity_id: vec!["1".into(), "2".into(), "3".into()],
217        };
218
219        // atom_to_residue: each residue has 3 atoms => offsets [0,3,6,...,30]
220        let atom_offsets: Vec<u32> = (0..=n_residues as u32).map(|i| i * 3).collect();
221        let atom_to_residue = Segmentation::from_offsets(atom_offsets);
222
223        // residue_to_chain: chain 0 has 4 residues, chain 1 has 3, chain 2 has 3
224        let residue_offsets: Vec<u32> = vec![0, 4, 7, 10];
225        let residue_to_chain = Segmentation::from_offsets(residue_offsets);
226
227        let bonds = Bonds::from_unsorted(vec![], vec![], vec![], n_atoms);
228
229        AtomicHierarchy {
230            atoms,
231            residues,
232            chains,
233            atom_to_residue,
234            residue_to_chain,
235            bonds,
236        }
237    }
238
239    #[test]
240    fn test_hierarchy_basic() {
241        let h = make_test_hierarchy();
242
243        assert_eq!(h.n_atoms(), 30);
244        assert_eq!(h.n_residues(), 10);
245        assert_eq!(h.n_chains(), 3);
246
247        // residue 0 -> atoms 0..3
248        assert_eq!(h.atoms_in_residue(0), 0..3);
249        // residue 4 -> atoms 12..15
250        assert_eq!(h.atoms_in_residue(4), 12..15);
251        // residue 9 -> atoms 27..30
252        assert_eq!(h.atoms_in_residue(9), 27..30);
253
254        // chain 0 -> residues 0..4
255        assert_eq!(h.residues_in_chain(0), 0..4);
256        // chain 1 -> residues 4..7
257        assert_eq!(h.residues_in_chain(1), 4..7);
258        // chain 2 -> residues 7..10
259        assert_eq!(h.residues_in_chain(2), 7..10);
260    }
261
262    #[test]
263    fn test_hierarchy_residue_of_atom() {
264        let h = make_test_hierarchy();
265
266        // atoms 0,1,2 -> residue 0
267        assert_eq!(h.residue_of_atom(0), 0);
268        assert_eq!(h.residue_of_atom(1), 0);
269        assert_eq!(h.residue_of_atom(2), 0);
270        // atoms 3,4,5 -> residue 1
271        assert_eq!(h.residue_of_atom(3), 1);
272        assert_eq!(h.residue_of_atom(5), 1);
273        // atom 29 -> residue 9
274        assert_eq!(h.residue_of_atom(29), 9);
275    }
276
277    #[test]
278    fn test_hierarchy_chain_of_residue() {
279        let h = make_test_hierarchy();
280
281        // residues 0..4 -> chain 0
282        for r in 0..4 {
283            assert_eq!(
284                h.chain_of_residue(r),
285                0,
286                "residue {} should be in chain 0",
287                r
288            );
289        }
290        // residues 4..7 -> chain 1
291        for r in 4..7 {
292            assert_eq!(
293                h.chain_of_residue(r),
294                1,
295                "residue {} should be in chain 1",
296                r
297            );
298        }
299        // residues 7..10 -> chain 2
300        for r in 7..10 {
301            assert_eq!(
302                h.chain_of_residue(r),
303                2,
304                "residue {} should be in chain 2",
305                r
306            );
307        }
308    }
309}
ferritin_core/model/hierarchy.rs

ferritin_core/model/
hierarchy.rs
