sdfrust/
graph.rs

1//! Graph adjacency infrastructure for molecular structures.
2//!
3//! Provides pre-computed adjacency lists and degree vectors for efficient
4//! graph-based algorithms (ring perception, featurization, fingerprints).
5//!
6//! # Example
7//!
8//! ```rust
9//! use sdfrust::{Molecule, Atom, Bond, BondOrder};
10//! use sdfrust::graph::AdjacencyList;
11//!
12//! let mut mol = Molecule::new("methane");
13//! mol.atoms.push(Atom::new(0, "C", 0.0, 0.0, 0.0));
14//! mol.atoms.push(Atom::new(1, "H", 1.0, 0.0, 0.0));
15//! mol.atoms.push(Atom::new(2, "H", -1.0, 0.0, 0.0));
16//! mol.bonds.push(Bond::new(0, 1, BondOrder::Single));
17//! mol.bonds.push(Bond::new(0, 2, BondOrder::Single));
18//!
19//! let adj = AdjacencyList::from_molecule(&mol);
20//! assert_eq!(adj.degree(0), 2);
21//! assert_eq!(adj.neighbors(0), &[(1, 0), (2, 1)]);
22//! ```
23
24use crate::molecule::Molecule;
25
26/// Pre-computed adjacency list for a molecule's bond graph.
27///
28/// Stores neighbor lists and degree vectors for O(1) access.
29/// Each neighbor entry is `(atom_index, bond_index)`.
30#[derive(Debug, Clone)]
31pub struct AdjacencyList {
32    /// For each atom, a list of (neighbor_atom_index, bond_index) pairs.
33    adj: Vec<Vec<(usize, usize)>>,
34    /// Degree (total bond count) for each atom.
35    degrees: Vec<usize>,
36    /// Heavy atom degree (bonds to non-H atoms) for each atom.
37    heavy_degrees: Vec<usize>,
38    /// Number of atoms.
39    num_atoms: usize,
40}
41
42impl AdjacencyList {
43    /// Build an adjacency list from a molecule.
44    ///
45    /// # Example
46    ///
47    /// ```rust
48    /// use sdfrust::{Molecule, Atom, Bond, BondOrder};
49    /// use sdfrust::graph::AdjacencyList;
50    ///
51    /// let mut mol = Molecule::new("water");
52    /// mol.atoms.push(Atom::new(0, "O", 0.0, 0.0, 0.0));
53    /// mol.atoms.push(Atom::new(1, "H", 1.0, 0.0, 0.0));
54    /// mol.atoms.push(Atom::new(2, "H", -0.3, 0.9, 0.0));
55    /// mol.bonds.push(Bond::new(0, 1, BondOrder::Single));
56    /// mol.bonds.push(Bond::new(0, 2, BondOrder::Single));
57    ///
58    /// let adj = AdjacencyList::from_molecule(&mol);
59    /// assert_eq!(adj.num_atoms(), 3);
60    /// assert_eq!(adj.degree(0), 2);
61    /// assert_eq!(adj.degree(1), 1);
62    /// ```
63    pub fn from_molecule(mol: &Molecule) -> Self {
64        let n = mol.atoms.len();
65        let mut adj: Vec<Vec<(usize, usize)>> = vec![vec![]; n];
66        let mut degrees = vec![0usize; n];
67        let mut heavy_degrees = vec![0usize; n];
68
69        for (bond_idx, bond) in mol.bonds.iter().enumerate() {
70            if bond.atom1 < n && bond.atom2 < n {
71                adj[bond.atom1].push((bond.atom2, bond_idx));
72                adj[bond.atom2].push((bond.atom1, bond_idx));
73                degrees[bond.atom1] += 1;
74                degrees[bond.atom2] += 1;
75
76                let elem1 = mol.atoms[bond.atom1].element.as_str();
77                let elem2 = mol.atoms[bond.atom2].element.as_str();
78                if !is_hydrogen(elem2) {
79                    heavy_degrees[bond.atom1] += 1;
80                }
81                if !is_hydrogen(elem1) {
82                    heavy_degrees[bond.atom2] += 1;
83                }
84            }
85        }
86
87        Self {
88            adj,
89            degrees,
90            heavy_degrees,
91            num_atoms: n,
92        }
93    }
94
95    /// Returns the number of atoms in the graph.
96    pub fn num_atoms(&self) -> usize {
97        self.num_atoms
98    }
99
100    /// Returns the neighbors of atom `idx` as `(neighbor_atom, bond_index)` pairs.
101    ///
102    /// Returns an empty slice if the index is out of bounds.
103    pub fn neighbors(&self, idx: usize) -> &[(usize, usize)] {
104        self.adj.get(idx).map(|v| v.as_slice()).unwrap_or(&[])
105    }
106
107    /// Returns just the neighbor atom indices for atom `idx`.
108    pub fn neighbor_atoms(&self, idx: usize) -> Vec<usize> {
109        self.neighbors(idx).iter().map(|(a, _)| *a).collect()
110    }
111
112    /// Returns the degree (number of bonds) for atom `idx`.
113    ///
114    /// Returns 0 if the index is out of bounds.
115    pub fn degree(&self, idx: usize) -> usize {
116        self.degrees.get(idx).copied().unwrap_or(0)
117    }
118
119    /// Returns the heavy atom degree (bonds to non-H atoms) for atom `idx`.
120    ///
121    /// Returns 0 if the index is out of bounds.
122    pub fn heavy_degree(&self, idx: usize) -> usize {
123        self.heavy_degrees.get(idx).copied().unwrap_or(0)
124    }
125
126    /// Returns a slice of all degrees.
127    pub fn degrees(&self) -> &[usize] {
128        &self.degrees
129    }
130
131    /// Returns a slice of all heavy atom degrees.
132    pub fn heavy_degrees(&self) -> &[usize] {
133        &self.heavy_degrees
134    }
135
136    /// Returns the bond indices for bonds involving atom `idx`.
137    pub fn bond_indices(&self, idx: usize) -> Vec<usize> {
138        self.neighbors(idx).iter().map(|(_, b)| *b).collect()
139    }
140}
141
142/// Check if an element symbol is hydrogen (H, D, or T).
143pub(crate) fn is_hydrogen(element: &str) -> bool {
144    let e = element.trim();
145    e.eq_ignore_ascii_case("H") || e.eq_ignore_ascii_case("D") || e.eq_ignore_ascii_case("T")
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use crate::atom::Atom;
152    use crate::bond::{Bond, BondOrder};
153
154    #[test]
155    fn test_adjacency_list_methane() {
156        let mut mol = Molecule::new("methane");
157        mol.atoms.push(Atom::new(0, "C", 0.0, 0.0, 0.0));
158        mol.atoms.push(Atom::new(1, "H", 1.0, 0.0, 0.0));
159        mol.atoms.push(Atom::new(2, "H", -1.0, 0.0, 0.0));
160        mol.atoms.push(Atom::new(3, "H", 0.0, 1.0, 0.0));
161        mol.atoms.push(Atom::new(4, "H", 0.0, -1.0, 0.0));
162        mol.bonds.push(Bond::new(0, 1, BondOrder::Single));
163        mol.bonds.push(Bond::new(0, 2, BondOrder::Single));
164        mol.bonds.push(Bond::new(0, 3, BondOrder::Single));
165        mol.bonds.push(Bond::new(0, 4, BondOrder::Single));
166
167        let adj = AdjacencyList::from_molecule(&mol);
168        assert_eq!(adj.num_atoms(), 5);
169        assert_eq!(adj.degree(0), 4);
170        assert_eq!(adj.heavy_degree(0), 0); // All neighbors are H
171        assert_eq!(adj.degree(1), 1);
172        assert_eq!(adj.heavy_degree(1), 1); // C is a heavy atom
173    }
174
175    #[test]
176    fn test_adjacency_list_benzene() {
177        let mut mol = Molecule::new("benzene");
178        for i in 0..6 {
179            mol.atoms.push(Atom::new(i, "C", 0.0, 0.0, 0.0));
180        }
181        for i in 0..6 {
182            mol.bonds
183                .push(Bond::new(i, (i + 1) % 6, BondOrder::Aromatic));
184        }
185
186        let adj = AdjacencyList::from_molecule(&mol);
187        for i in 0..6 {
188            assert_eq!(adj.degree(i), 2);
189            assert_eq!(adj.heavy_degree(i), 2);
190        }
191    }
192
193    #[test]
194    fn test_adjacency_list_empty() {
195        let mol = Molecule::new("empty");
196        let adj = AdjacencyList::from_molecule(&mol);
197        assert_eq!(adj.num_atoms(), 0);
198    }
199
200    #[test]
201    fn test_neighbor_atoms() {
202        let mut mol = Molecule::new("test");
203        mol.atoms.push(Atom::new(0, "C", 0.0, 0.0, 0.0));
204        mol.atoms.push(Atom::new(1, "C", 1.0, 0.0, 0.0));
205        mol.atoms.push(Atom::new(2, "O", 2.0, 0.0, 0.0));
206        mol.bonds.push(Bond::new(0, 1, BondOrder::Single));
207        mol.bonds.push(Bond::new(1, 2, BondOrder::Double));
208
209        let adj = AdjacencyList::from_molecule(&mol);
210        let neighbors = adj.neighbor_atoms(1);
211        assert_eq!(neighbors.len(), 2);
212        assert!(neighbors.contains(&0));
213        assert!(neighbors.contains(&2));
214    }
215
216    #[test]
217    fn test_bond_indices() {
218        let mut mol = Molecule::new("test");
219        mol.atoms.push(Atom::new(0, "C", 0.0, 0.0, 0.0));
220        mol.atoms.push(Atom::new(1, "C", 1.0, 0.0, 0.0));
221        mol.atoms.push(Atom::new(2, "O", 2.0, 0.0, 0.0));
222        mol.bonds.push(Bond::new(0, 1, BondOrder::Single));
223        mol.bonds.push(Bond::new(1, 2, BondOrder::Double));
224
225        let adj = AdjacencyList::from_molecule(&mol);
226        let bonds = adj.bond_indices(1);
227        assert_eq!(bonds.len(), 2);
228        assert!(bonds.contains(&0));
229        assert!(bonds.contains(&1));
230    }
231
232    #[test]
233    fn test_is_hydrogen() {
234        assert!(is_hydrogen("H"));
235        assert!(is_hydrogen("D"));
236        assert!(is_hydrogen("T"));
237        assert!(is_hydrogen(" H "));
238        assert!(!is_hydrogen("C"));
239        assert!(!is_hydrogen("He"));
240    }
241
242    #[test]
243    fn test_out_of_bounds() {
244        let mol = Molecule::new("empty");
245        let adj = AdjacencyList::from_molecule(&mol);
246        assert_eq!(adj.degree(999), 0);
247        assert_eq!(adj.heavy_degree(999), 0);
248        assert!(adj.neighbors(999).is_empty());
249    }
250}
sdfrust/graph.rs

sdfrust/
graph.rs
