cosmolkit-core 0.2.1

Redesigned COSMolKit core with value-style molecule state and explicit topology operation contracts
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use crate::{AtomId, Bond, BondId};

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum AdjacencyError {
    #[error("bond {bond} {endpoint} atom index {atom} is out of range for {atom_count} atoms")]
    BondAtomOutOfRange {
        bond: BondId,
        endpoint: &'static str,
        atom: AtomId,
        atom_count: usize,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NeighborRef {
    pub atom_index: usize,
    pub bond: BondId,
}

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct AdjacencyList {
    offsets: Vec<usize>,
    entries: Vec<NeighborRef>,
}

impl AdjacencyList {
    pub fn try_from_topology(atom_count: usize, bonds: &[Bond]) -> Result<Self, AdjacencyError> {
        let mut degrees = vec![0usize; atom_count];
        for bond in bonds {
            let begin = bond.begin();
            let end = bond.end();
            if begin.index() >= atom_count {
                return Err(AdjacencyError::BondAtomOutOfRange {
                    bond: bond.id(),
                    endpoint: "begin",
                    atom: begin,
                    atom_count,
                });
            }
            if end.index() >= atom_count {
                return Err(AdjacencyError::BondAtomOutOfRange {
                    bond: bond.id(),
                    endpoint: "end",
                    atom: end,
                    atom_count,
                });
            }
            degrees[begin.index()] += 1;
            degrees[end.index()] += 1;
        }

        let mut offsets = Vec::with_capacity(atom_count + 1);
        offsets.push(0);
        for degree in degrees {
            offsets.push(offsets.last().copied().expect("offsets is nonempty") + degree);
        }

        let mut entries = vec![
            NeighborRef {
                atom_index: 0,
                bond: BondId::new(0),
            };
            offsets.last().copied().unwrap_or(0)
        ];
        let mut cursor = offsets[..atom_count].to_vec();
        for bond in bonds {
            let begin = bond.begin().index();
            let end = bond.end().index();
            let begin_slot = cursor[begin];
            entries[begin_slot] = NeighborRef {
                atom_index: end,
                bond: bond.id(),
            };
            cursor[begin] += 1;

            let end_slot = cursor[end];
            entries[end_slot] = NeighborRef {
                atom_index: begin,
                bond: bond.id(),
            };
            cursor[end] += 1;
        }

        Ok(Self { offsets, entries })
    }

    #[must_use]
    pub fn from_topology(atom_count: usize, bonds: &[Bond]) -> Self {
        Self::try_from_topology(atom_count, bonds)
            .expect("topology must be valid before building adjacency")
    }

    #[must_use]
    pub fn neighbors_of(&self, atom_index: usize) -> &[NeighborRef] {
        let Some(window) = self.offsets.get(atom_index..atom_index.saturating_add(2)) else {
            return &[];
        };
        if window.len() != 2 {
            return &[];
        }
        &self.entries[window[0]..window[1]]
    }
}

#[cfg(test)]
mod tests {
    use crate::{AdjacencyList, AtomSpec, BondOrder, BondSpec, Element, MoleculeBuilder};

    #[test]
    fn adjacency_list_stores_incident_bonds_in_csr_layout() {
        let mut builder = MoleculeBuilder::new();
        let left = builder.add_atom(AtomSpec::new(Element::C));
        let middle = builder.add_atom(AtomSpec::new(Element::C));
        let right = builder.add_atom(AtomSpec::new(Element::O));
        builder
            .add_bond(BondSpec::new(left, middle, BondOrder::Single))
            .unwrap();
        builder
            .add_bond(BondSpec::new(middle, right, BondOrder::Double))
            .unwrap();
        let molecule = builder.build().unwrap();

        let adjacency =
            AdjacencyList::try_from_topology(molecule.num_atoms(), molecule.bonds()).unwrap();

        assert_eq!(adjacency.neighbors_of(left.index()).len(), 1);
        assert_eq!(adjacency.neighbors_of(middle.index()).len(), 2);
        assert_eq!(adjacency.neighbors_of(right.index()).len(), 1);
        assert!(adjacency.neighbors_of(usize::MAX).is_empty());
    }
}