molframe_core/chunk/
parent.rs

1//! Finding the residue an atom belongs to.
2//!
3//! The downward direction is free: a residue's atoms are a range, so it is two
4//! loads and an addition. The upward direction is the awkward one, and it admits
5//! a trade the other does not.
6//!
7//! Storing a residue number beside every atom answers it in one load and costs
8//! four bytes per atom — sixteen megabytes on a four-million-atom virus capsid,
9//! which is cache the geometric kernels do not get. Storing nothing answers it
10//! by binary search over the residue offsets, which is fine occasionally and
11//! poor in a loop.
12//!
13//! The default is neither: record the residue at the start of every block of a
14//! hundred and twenty-eight atoms and scan forward within the block. That is
15//! constant work per lookup at a thirty-second of the memory, because a block
16//! spans a handful of residues and the scan never leaves one cache line.
17
18use crate::index::ResidueIndex;
19use crate::topology::ResidueTable;
20use std::mem::size_of_val;
21
22/// How an atom's residue is recovered.
23#[derive(Clone, PartialEq, Eq, Debug, Default)]
24pub enum ParentMapping {
25    /// Nothing stored; search the residue offsets.
26    #[default]
27    OffsetsOnly,
28    /// One residue position per atom.
29    Explicit(Vec<u32>),
30    /// One residue position per block of atoms, refined by a bounded scan.
31    BlockIndexed {
32        /// Atoms per block.
33        block: u16,
34        /// The residue the first atom of each block belongs to.
35        parents: Vec<u32>,
36    },
37}
38
39impl ParentMapping {
40    /// Atoms per block in the default mapping.
41    ///
42    /// A hundred and twenty-eight four-byte entries is one entry per five
43    /// hundred and twelve bytes of atom data, and the scan it implies stays
44    /// inside a few residues.
45    pub const DEFAULT_BLOCK: u16 = 128;
46
47    /// Builds the default mapping for a chunk whose atoms start at `first_atom`.
48    ///
49    /// `parent_of` is the residue of each atom in chunk order.
50    #[must_use]
51    pub fn block_indexed(parents_in_order: &[u32]) -> Self {
52        let block = Self::DEFAULT_BLOCK;
53        let parents = parents_in_order
54            .iter()
55            .step_by(usize::from(block))
56            .copied()
57            .collect();
58
59        Self::BlockIndexed { block, parents }
60    }
61
62    /// Builds a mapping that answers in one load, for a chunk whose atoms are
63    /// looked up individually often enough to justify the memory.
64    #[must_use]
65    pub fn explicit(parents_in_order: &[u32]) -> Self {
66        Self::Explicit(parents_in_order.to_vec())
67    }
68
69    /// The bytes this mapping costs for `atoms` atoms.
70    #[must_use]
71    pub fn bytes(&self) -> usize {
72        match self {
73            Self::OffsetsOnly => 0,
74            Self::Explicit(parents) | Self::BlockIndexed { parents, .. } => {
75                size_of_val(parents.as_slice())
76            }
77        }
78    }
79
80    /// The residue of the atom at `local` within a chunk whose atoms begin at
81    /// `first_atom` in the structure's flat order.
82    ///
83    /// `residues` is consulted only where the mapping cannot answer alone.
84    #[must_use]
85    pub fn resolve(
86        &self,
87        local: u32,
88        first_atom: u32,
89        residues: &ResidueTable,
90    ) -> Option<ResidueIndex> {
91        match self {
92            Self::OffsetsOnly => Self::resolve_from_offsets(local, first_atom, residues),
93            Self::Explicit(parents) => Self::resolve_explicit(local, parents),
94            Self::BlockIndexed { block, parents } => {
95                Self::resolve_block_indexed(local, first_atom, *block, parents, residues)
96            }
97        }
98    }
99
100    fn resolve_from_offsets(
101        local: u32,
102        first_atom: u32,
103        residues: &ResidueTable,
104    ) -> Option<ResidueIndex> {
105        first_atom
106            .checked_add(local)
107            .and_then(|atom| residues.containing(atom))
108    }
109
110    fn resolve_explicit(local: u32, parents: &[u32]) -> Option<ResidueIndex> {
111        usize::try_from(local)
112            .ok()
113            .and_then(|index| parents.get(index))
114            .copied()
115            .map(ResidueIndex::new)
116    }
117
118    fn resolve_block_indexed(
119        local: u32,
120        first_atom: u32,
121        block: u16,
122        parents: &[u32],
123        residues: &ResidueTable,
124    ) -> Option<ResidueIndex> {
125        let start = Self::block_parent(local, block, parents)?;
126
127        let atom = first_atom.checked_add(local)?;
128
129        Self::scan_forward(start, atom, residues)
130    }
131
132    fn block_parent(local: u32, block: u16, parents: &[u32]) -> Option<u32> {
133        let block = u32::from(block).max(1);
134        let block_index = local / block;
135
136        usize::try_from(block_index)
137            .ok()
138            .and_then(|index| parents.get(index))
139            .copied()
140    }
141
142    /// Walks forward from a known residue until the one containing `atom`.
143    ///
144    /// Bounded by the number of residues a block spans, which is small by
145    /// construction. Returns nothing rather than looping if the tables are
146    /// inconsistent with each other.
147    fn scan_forward(start: u32, atom: u32, residues: &ResidueTable) -> Option<ResidueIndex> {
148        let start = usize::try_from(start).ok()?;
149
150        (start..residues.len())
151            .map_while(|candidate| {
152                let candidate = u32::try_from(candidate).ok()?;
153
154                let residue = ResidueIndex::new(candidate);
155                let range = residues.atoms(residue)?;
156
157                (atom >= range.start).then_some((residue, range))
158            })
159            .find_map(|(residue, range)| range.contains(&atom).then_some(residue))
160    }
161}
162
163#[cfg(test)]
164#[path = "parent_tests.rs"]
165mod tests;
molframe_core/chunk/parent.rs

molframe_core/chunk/
parent.rs
