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
6/// Bond connectivity stored as SoA (struct-of-arrays).
7///
8/// `atom_a` and `atom_b` are atom indices; `order` is bond order (1 = single, 2 = double, etc.).
9///
10/// After construction the arrays are sorted by `atom_a` so that
11/// `atom_bond_starts[i]..atom_bond_starts[i+1]` gives the bond range for atom `i`.
12#[derive(Clone, Debug)]
13pub struct Bonds {
14    /// First atom index for each bond (sorted).
15    pub atom_a: Vec<u32>,
16    /// Second atom index for each bond.
17    pub atom_b: Vec<u32>,
18    /// Bond order for each bond (1 = single, 2 = double, 3 = triple, etc.).
19    pub order: Vec<u8>,
20    /// CSR-style start index into bonds arrays, indexed by atom.
21    ///
22    /// `atom_bond_starts.len() == n_atoms + 1`; atom `i`'s bonds span
23    /// `atom_bond_starts[i]..atom_bond_starts[i+1]`.
24    pub atom_bond_starts: Vec<u32>,
25}
26
27impl Bonds {
28    /// Construct from parallel atom_a/atom_b/order vectors.
29    ///
30    /// Sorts bonds by `atom_a` and builds the `atom_bond_starts` CSR index.
31    ///
32    /// # Panics
33    /// Panics if `atom_a`, `atom_b`, and `order` have different lengths,
34    /// or if any atom index is >= `n_atoms`.
35    pub fn from_unsorted(
36        atom_a: Vec<u32>,
37        atom_b: Vec<u32>,
38        order: Vec<u8>,
39        n_atoms: usize,
40    ) -> Self {
41        assert_eq!(atom_a.len(), atom_b.len(), "atom_a and atom_b must have the same length");
42        assert_eq!(atom_a.len(), order.len(), "atom_a and order must have the same length");
43
44        let n_bonds = atom_a.len();
45
46        // Build sort permutation by atom_a
47        let mut indices: Vec<usize> = (0..n_bonds).collect();
48        indices.sort_by_key(|&i| atom_a[i]);
49
50        let sorted_a: Vec<u32> = indices.iter().map(|&i| atom_a[i]).collect();
51        let sorted_b: Vec<u32> = indices.iter().map(|&i| atom_b[i]).collect();
52        let sorted_order: Vec<u8> = indices.iter().map(|&i| order[i]).collect();
53
54        // Build CSR atom_bond_starts: atom_bond_starts[i] = first bond index for atom i
55        let mut atom_bond_starts = vec![0u32; n_atoms + 1];
56        for &a in &sorted_a {
57            let a = a as usize;
58            assert!(a < n_atoms, "atom index {} out of range (n_atoms={})", a, n_atoms);
59            atom_bond_starts[a + 1] += 1;
60        }
61        // prefix-sum
62        for i in 1..=n_atoms {
63            atom_bond_starts[i] += atom_bond_starts[i - 1];
64        }
65
66        Self {
67            atom_a: sorted_a,
68            atom_b: sorted_b,
69            order: sorted_order,
70            atom_bond_starts,
71        }
72    }
73
74    /// Returns an iterator over `(atom_b, order)` pairs for all bonds from `atom_idx`.
75    ///
76    /// Only bonds where `atom_idx` appears as `atom_a` are returned. For undirected
77    /// traversal, callers should also check the reverse (atom_b side) or store both
78    /// directions when constructing.
79    pub fn bonds_for_atom(&self, atom_idx: usize) -> impl Iterator<Item = (u32, u8)> + '_ {
80        let start = self.atom_bond_starts[atom_idx] as usize;
81        let end = self.atom_bond_starts[atom_idx + 1] as usize;
82        self.atom_b[start..end]
83            .iter()
84            .zip(self.order[start..end].iter())
85            .map(|(&b, &o)| (b, o))
86    }
87
88    /// Number of bonds.
89    pub fn len(&self) -> usize {
90        self.atom_a.len()
91    }
92
93    /// Returns `true` if there are no bonds.
94    pub fn is_empty(&self) -> bool {
95        self.atom_a.is_empty()
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102
103    #[test]
104    fn test_bonds_from_unsorted_sorted_correctly() {
105        // Provide bonds out of order: (2->3), (0->1), (1->2)
106        let atom_a = vec![2u32, 0, 1];
107        let atom_b = vec![3u32, 1, 2];
108        let order = vec![1u8, 2, 1];
109
110        let bonds = Bonds::from_unsorted(atom_a, atom_b, order, 4);
111
112        // After sorting by atom_a the order should be: (0->1,2), (1->2,1), (2->3,1)
113        assert_eq!(bonds.atom_a, vec![0, 1, 2]);
114        assert_eq!(bonds.atom_b, vec![1, 2, 3]);
115        assert_eq!(bonds.order, vec![2, 1, 1]);
116
117        // CSR index: atom_bond_starts should be [0, 1, 2, 3, 3]
118        assert_eq!(bonds.atom_bond_starts, vec![0, 1, 2, 3, 3]);
119        assert_eq!(bonds.len(), 3);
120    }
121
122    #[test]
123    fn test_bonds_for_atom() {
124        let atom_a = vec![0u32, 0, 1];
125        let atom_b = vec![1u32, 2, 2];
126        let order = vec![1u8, 2, 1];
127
128        let bonds = Bonds::from_unsorted(atom_a, atom_b, order, 3);
129
130        // Atom 0 has bonds to 1 (order=1) and 2 (order=2)
131        let bonds_0: Vec<(u32, u8)> = bonds.bonds_for_atom(0).collect();
132        assert_eq!(bonds_0.len(), 2);
133        assert!(bonds_0.contains(&(1, 1)));
134        assert!(bonds_0.contains(&(2, 2)));
135
136        // Atom 1 has a bond to 2 (order=1)
137        let bonds_1: Vec<(u32, u8)> = bonds.bonds_for_atom(1).collect();
138        assert_eq!(bonds_1, vec![(2, 1)]);
139
140        // Atom 2 has no outgoing bonds (only incoming)
141        let bonds_2: Vec<(u32, u8)> = bonds.bonds_for_atom(2).collect();
142        assert!(bonds_2.is_empty());
143    }
144
145    #[test]
146    fn test_bonds_empty() {
147        let bonds = Bonds::from_unsorted(vec![], vec![], vec![], 5);
148        assert!(bonds.is_empty());
149        assert_eq!(bonds.len(), 0);
150        assert_eq!(bonds.atom_bond_starts, vec![0, 0, 0, 0, 0, 0]);
151    }
152}
ferritin_core/model/bonds.rs

ferritin_core/model/
bonds.rs
