Skip to main content

molframe_core/anisotropy/
model.rs

1//! Anisotropic displacement ellipsoids stored as a sparse per-atom table.
2//!
3//! Most structures carry no anisotropic ellipsoids, so the table is optional
4//! in the same sense connectivity is: `available` records whether the source
5//! resolved the category at all, distinct from an empty but known set. Rows
6//! are kept sorted by atom position, so building costs O(rows log rows) once
7//! and every per-atom lookup costs O(log rows); a structure without the
8//! category pays nothing beyond a shared empty table.
9
10use crate::index::{AnisotropyIndex, AtomIndex};
11use std::sync::Arc;
12
13/// One anisotropic displacement ellipsoid before it is packed into columns.
14#[derive(Clone, Copy, Debug, PartialEq)]
15pub struct AnisotropicDisplacement {
16    /// The atom row this ellipsoid belongs to.
17    pub atom: AtomIndex,
18    /// Displacement tensor in ångström squared, stored as
19    /// `[U11, U22, U33, U12, U13, U23]`.
20    pub u: [f32; 6],
21}
22
23/// A compact, immutable anisotropic displacement table sorted by atom.
24#[derive(Clone, Debug)]
25pub struct AnisotropyTable {
26    atom: Arc<Vec<AtomIndex>>,
27    u: Arc<Vec<[f32; 6]>>,
28    available: bool,
29}
30
31impl Default for AnisotropyTable {
32    fn default() -> Self {
33        Self {
34            atom: Arc::new(Vec::new()),
35            u: Arc::new(Vec::new()),
36            available: false,
37        }
38    }
39}
40
41impl AnisotropyTable {
42    /// Number of ellipsoid rows.
43    #[must_use]
44    pub fn len(&self) -> usize {
45        self.atom.len()
46    }
47
48    /// Whether no ellipsoid is recorded.
49    #[must_use]
50    pub fn is_empty(&self) -> bool {
51        self.atom.is_empty()
52    }
53
54    /// Whether anisotropic displacement was resolved, including a known empty set.
55    #[must_use]
56    pub const fn is_available(&self) -> bool {
57        self.available
58    }
59
60    /// Returns one ellipsoid by position.
61    #[must_use]
62    pub fn get(&self, row: AnisotropyIndex) -> Option<AnisotropicDisplacement> {
63        let position = row.as_usize();
64        Some(AnisotropicDisplacement {
65            atom: *self.atom.get(position)?,
66            u: *self.u.get(position)?,
67        })
68    }
69
70    /// Returns the tensor of one atom, or `None` when it is isotropic.
71    #[must_use]
72    pub fn for_atom(&self, atom: AtomIndex) -> Option<[f32; 6]> {
73        let position = self.atom.binary_search(&atom).ok()?;
74        Some(self.u[position])
75    }
76
77    /// Walks records in stable atom order.
78    #[must_use]
79    pub fn iter(&self) -> impl ExactSizeIterator<Item = AnisotropicDisplacement> + '_ {
80        self.atom
81            .iter()
82            .zip(self.u.iter())
83            .map(|(atom, u)| AnisotropicDisplacement { atom: *atom, u: *u })
84    }
85}
86
87/// Deterministic anisotropy-table construction with atom deduplication.
88///
89/// Records accumulate in insertion order and are sorted and deduplicated once
90/// at `finish`. The first record pushed for an atom wins, so a later duplicate
91/// never replaces an earlier assignment.
92#[derive(Debug, Default)]
93pub struct AnisotropyTableBuilder {
94    records: Vec<AnisotropicDisplacement>,
95}
96
97impl AnisotropyTableBuilder {
98    /// Creates an empty builder.
99    #[must_use]
100    pub const fn new() -> Self {
101        Self {
102            records: Vec::new(),
103        }
104    }
105
106    /// Adds one ellipsoid. The first assignment wins on duplicates.
107    pub fn push(&mut self, record: AnisotropicDisplacement) {
108        self.records.push(record);
109    }
110
111    /// Packs records into immutable columns.
112    #[must_use]
113    pub fn finish(self) -> AnisotropyTable {
114        self.finish_with_availability(true)
115    }
116
117    /// Packs records while preserving whether the category was resolved.
118    ///
119    /// A merge can retain explicit ellipsoids from one input while marking the
120    /// combined table unavailable because another input carried no data.
121    #[must_use]
122    pub fn finish_with_availability(mut self, available: bool) -> AnisotropyTable {
123        // The sort is stable and `dedup_by_key` keeps the first of each run, so
124        // the earliest record for an atom wins exactly as an insert-if-absent
125        // would, and the output stays in ascending atom order.
126        self.records.sort_by_key(|record| record.atom.get());
127        self.records.dedup_by_key(|record| record.atom.get());
128
129        let mut atom = Vec::with_capacity(self.records.len());
130        let mut u = Vec::with_capacity(self.records.len());
131        for record in self.records {
132            atom.push(record.atom);
133            u.push(record.u);
134        }
135        AnisotropyTable {
136            atom: Arc::new(atom),
137            u: Arc::new(u),
138            available,
139        }
140    }
141}
142
143#[cfg(test)]
144#[path = "model_tests.rs"]
145mod tests;