ferritin_core/model/
bonds.rs

1//! Bond connectivity stored in struct-of-arrays (SoA) layout with CSR indexing.
2//!
3//! After construction via [`Bonds::from_unsorted`], bonds are sorted by `atom_a`
4//! and `atom_bond_starts` provides O(1) access to all bonds for a given atom.
5
6use super::tables::{AtomsTable, ResidueGroup, ResiduesTable};
7use crate::data::Segmentation;
8use crate::info::constants::get_bonds_canonical20;
9
10/// Bond connectivity stored as SoA (struct-of-arrays).
11///
12/// `atom_a` and `atom_b` are atom indices; `order` is bond order (1 = single, 2 = double, etc.).
13///
14/// After construction the arrays are sorted by `atom_a` so that
15/// `atom_bond_starts[i]..atom_bond_starts[i+1]` gives the bond range for atom `i`.
16#[derive(Clone, Debug)]
17pub struct Bonds {
18    /// First atom index for each bond (sorted).
19    pub atom_a: Vec<u32>,
20    /// Second atom index for each bond.
21    pub atom_b: Vec<u32>,
22    /// Bond order for each bond (1 = single, 2 = double, 3 = triple, etc.).
23    pub order: Vec<u8>,
24    /// CSR-style start index into bonds arrays, indexed by atom.
25    ///
26    /// `atom_bond_starts.len() == n_atoms + 1`; atom `i`'s bonds span
27    /// `atom_bond_starts[i]..atom_bond_starts[i+1]`.
28    pub atom_bond_starts: Vec<u32>,
29}
30
31impl Bonds {
32    /// Construct from parallel atom_a/atom_b/order vectors.
33    ///
34    /// Sorts bonds by `atom_a` and builds the `atom_bond_starts` CSR index.
35    ///
36    /// # Panics
37    /// Panics if `atom_a`, `atom_b`, and `order` have different lengths,
38    /// or if any atom index is >= `n_atoms`.
39    pub fn from_unsorted(
40        atom_a: Vec<u32>,
41        atom_b: Vec<u32>,
42        order: Vec<u8>,
43        n_atoms: usize,
44    ) -> Self {
45        assert_eq!(
46            atom_a.len(),
47            atom_b.len(),
48            "atom_a and atom_b must have the same length"
49        );
50        assert_eq!(
51            atom_a.len(),
52            order.len(),
53            "atom_a and order must have the same length"
54        );
55
56        let n_bonds = atom_a.len();
57
58        // Build sort permutation by atom_a
59        let mut indices: Vec<usize> = (0..n_bonds).collect();
60        indices.sort_by_key(|&i| atom_a[i]);
61
62        let sorted_a: Vec<u32> = indices.iter().map(|&i| atom_a[i]).collect();
63        let sorted_b: Vec<u32> = indices.iter().map(|&i| atom_b[i]).collect();
64        let sorted_order: Vec<u8> = indices.iter().map(|&i| order[i]).collect();
65
66        // Build CSR atom_bond_starts: atom_bond_starts[i] = first bond index for atom i
67        let mut atom_bond_starts = vec![0u32; n_atoms + 1];
68        for &a in &sorted_a {
69            let a = a as usize;
70            assert!(
71                a < n_atoms,
72                "atom index {} out of range (n_atoms={})",
73                a,
74                n_atoms
75            );
76            atom_bond_starts[a + 1] += 1;
77        }
78        // prefix-sum
79        for i in 1..=n_atoms {
80            atom_bond_starts[i] += atom_bond_starts[i - 1];
81        }
82
83        Self {
84            atom_a: sorted_a,
85            atom_b: sorted_b,
86            order: sorted_order,
87            atom_bond_starts,
88        }
89    }
90
91    /// Returns an iterator over `(atom_b, order)` pairs for all bonds from `atom_idx`.
92    ///
93    /// Only bonds where `atom_idx` appears as `atom_a` are returned. For undirected
94    /// traversal, callers should also check the reverse (atom_b side) or store both
95    /// directions when constructing.
96    pub fn bonds_for_atom(&self, atom_idx: usize) -> impl Iterator<Item = (u32, u8)> + '_ {
97        let start = self.atom_bond_starts[atom_idx] as usize;
98        let end = self.atom_bond_starts[atom_idx + 1] as usize;
99        self.atom_b[start..end]
100            .iter()
101            .zip(self.order[start..end].iter())
102            .map(|(&b, &o)| (b, o))
103    }
104
105    /// Number of bonds.
106    pub fn len(&self) -> usize {
107        self.atom_a.len()
108    }
109
110    /// Returns `true` if there are no bonds.
111    pub fn is_empty(&self) -> bool {
112        self.atom_a.is_empty()
113    }
114}
115
116/// Infer bond connectivity from canonical-residue templates plus backbone C-N
117/// peptide bonds between consecutive polymer residues on the same chain.
118///
119/// CIF/mmCIF files rarely carry explicit connectivity records, so without this
120/// a freshly-parsed [`super::model::Model`] has no bonds at all (see
121/// [`AtomicHierarchy`](super::hierarchy::AtomicHierarchy)::bonds). Mirrors
122/// `AtomCollection::connect_via_residue_names`.
123pub(crate) fn infer_bonds_from_residue_templates(
124    atoms: &AtomsTable,
125    residues: &ResiduesTable,
126    atom_to_residue: &Segmentation,
127    residue_to_chain: &Segmentation,
128) -> Bonds {
129    let n_atoms = atoms.len();
130    let n_residues = residues.len();
131    let aa_bond_info = get_bonds_canonical20();
132
133    let mut atom_a = Vec::new();
134    let mut atom_b = Vec::new();
135    let mut order = Vec::new();
136
137    for res_idx in 0..n_residues {
138        let range = atom_to_residue.segment(res_idx);
139        if let Some(bond_dict_for_res) = aa_bond_info.get(residues.comp_id[res_idx].as_str()) {
140            for &(name1, name2, bond_order) in bond_dict_for_res {
141                let indices1: Vec<usize> = range
142                    .clone()
143                    .filter(|&i| atoms.atom_name[i] == name1)
144                    .collect();
145                let indices2: Vec<usize> = range
146                    .clone()
147                    .filter(|&i| atoms.atom_name[i] == name2)
148                    .collect();
149                for &i in &indices1 {
150                    for &j in &indices2 {
151                        atom_a.push(i as u32);
152                        atom_b.push(j as u32);
153                        order.push(bond_order as u8);
154                    }
155                }
156            }
157        }
158    }
159
160    // Backbone C->N peptide bonds between consecutive polymer residues on the same chain.
161    for res_idx in 0..n_residues.saturating_sub(1) {
162        let next_idx = res_idx + 1;
163        if residues.group[res_idx] != ResidueGroup::Polymer
164            || residues.group[next_idx] != ResidueGroup::Polymer
165        {
166            continue;
167        }
168        if residue_to_chain.segment_of(res_idx) != residue_to_chain.segment_of(next_idx) {
169            continue;
170        }
171        let curr_range = atom_to_residue.segment(res_idx);
172        let next_range = atom_to_residue.segment(next_idx);
173        let c_idx = curr_range.clone().find(|&i| atoms.atom_name[i] == "C");
174        let n_idx = next_range.clone().find(|&i| atoms.atom_name[i] == "N");
175        if let (Some(c), Some(n)) = (c_idx, n_idx) {
176            atom_a.push(c as u32);
177            atom_b.push(n as u32);
178            order.push(1);
179        }
180    }
181
182    Bonds::from_unsorted(atom_a, atom_b, order, n_atoms)
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn test_bonds_from_unsorted_sorted_correctly() {
191        // Provide bonds out of order: (2->3), (0->1), (1->2)
192        let atom_a = vec![2u32, 0, 1];
193        let atom_b = vec![3u32, 1, 2];
194        let order = vec![1u8, 2, 1];
195
196        let bonds = Bonds::from_unsorted(atom_a, atom_b, order, 4);
197
198        // After sorting by atom_a the order should be: (0->1,2), (1->2,1), (2->3,1)
199        assert_eq!(bonds.atom_a, vec![0, 1, 2]);
200        assert_eq!(bonds.atom_b, vec![1, 2, 3]);
201        assert_eq!(bonds.order, vec![2, 1, 1]);
202
203        // CSR index: atom_bond_starts should be [0, 1, 2, 3, 3]
204        assert_eq!(bonds.atom_bond_starts, vec![0, 1, 2, 3, 3]);
205        assert_eq!(bonds.len(), 3);
206    }
207
208    #[test]
209    fn test_bonds_for_atom() {
210        let atom_a = vec![0u32, 0, 1];
211        let atom_b = vec![1u32, 2, 2];
212        let order = vec![1u8, 2, 1];
213
214        let bonds = Bonds::from_unsorted(atom_a, atom_b, order, 3);
215
216        // Atom 0 has bonds to 1 (order=1) and 2 (order=2)
217        let bonds_0: Vec<(u32, u8)> = bonds.bonds_for_atom(0).collect();
218        assert_eq!(bonds_0.len(), 2);
219        assert!(bonds_0.contains(&(1, 1)));
220        assert!(bonds_0.contains(&(2, 2)));
221
222        // Atom 1 has a bond to 2 (order=1)
223        let bonds_1: Vec<(u32, u8)> = bonds.bonds_for_atom(1).collect();
224        assert_eq!(bonds_1, vec![(2, 1)]);
225
226        // Atom 2 has no outgoing bonds (only incoming)
227        let bonds_2: Vec<(u32, u8)> = bonds.bonds_for_atom(2).collect();
228        assert!(bonds_2.is_empty());
229    }
230
231    #[test]
232    fn test_bonds_empty() {
233        let bonds = Bonds::from_unsorted(vec![], vec![], vec![], 5);
234        assert!(bonds.is_empty());
235        assert_eq!(bonds.len(), 0);
236        assert_eq!(bonds.atom_bond_starts, vec![0, 0, 0, 0, 0, 0]);
237    }
238}
ferritin_core/model/bonds.rs

ferritin_core/model/
bonds.rs
