ferritin_core/unit/
unit.rs

1//! Zero-copy view into a subset of a [`Model`].
2
3use crate::data::OrderedSet;
4use crate::model::Model;
5
6/// Zero-copy view into a subset of a Model.
7///
8/// `Unit` holds a reference to a `Model` and an `OrderedSet` of selected atom indices.
9/// All operations are lazy — iterators do not allocate, and set operations produce
10/// new `Unit`s without copying coordinate data.
11#[derive(Clone)]
12pub struct Unit<'a> {
13    model: &'a Model,
14    atoms: OrderedSet,
15}
16
17impl<'a> Unit<'a> {
18    /// Create a `Unit` from a boolean mask.
19    ///
20    /// The mask must have the same length as `model.n_atoms()`.
21    /// Only atoms where `mask[i] == true` are included.
22    ///
23    /// # Panics
24    /// Panics if `mask.len() != model.n_atoms()`.
25    pub fn from_mask(model: &'a Model, mask: &[bool]) -> Self {
26        assert_eq!(
27            mask.len(),
28            model.n_atoms(),
29            "mask length must equal number of atoms"
30        );
31        let indices: Vec<u32> = mask
32            .iter()
33            .enumerate()
34            .filter_map(|(i, &selected)| if selected { Some(i as u32) } else { None })
35            .collect();
36
37        let atoms = if indices.is_empty() {
38            OrderedSet::interval(0, 0)
39        } else {
40            OrderedSet::from_sorted(indices)
41        };
42
43        Self { model, atoms }
44    }
45
46    /// Create a `Unit` containing all atoms in the model.
47    pub fn all(model: &'a Model) -> Self {
48        let n = model.n_atoms() as u32;
49        Self {
50            model,
51            atoms: OrderedSet::interval(0, n),
52        }
53    }
54
55    /// Create a `Unit` from pre-computed indices.
56    pub fn from_indices(model: &'a Model, indices: OrderedSet) -> Self {
57        Self {
58            model,
59            atoms: indices,
60        }
61    }
62
63    /// Number of selected atoms.
64    pub fn len(&self) -> usize {
65        self.atoms.len()
66    }
67
68    /// Returns `true` if no atoms are selected.
69    pub fn is_empty(&self) -> bool {
70        self.atoms.is_empty()
71    }
72
73    /// Reference to the underlying model.
74    pub fn model(&self) -> &Model {
75        self.model
76    }
77
78    /// Lazy iterator over selected atom indices.
79    pub fn atom_indices(&self) -> impl Iterator<Item = u32> + '_ {
80        self.atoms.iter()
81    }
82
83    /// Lazy iterator over coordinates of selected atoms.
84    pub fn coords(&self) -> impl Iterator<Item = [f32; 3]> + '_ {
85        self.atoms.iter().map(|i| self.model.coord(i as usize))
86    }
87
88    /// Union of two units (must reference the same model).
89    ///
90    /// # Panics
91    /// Panics if the units reference different models.
92    pub fn union(&self, other: &Unit<'a>) -> Self {
93        assert!(
94            std::ptr::eq(self.model, other.model),
95            "cannot union Units from different Models"
96        );
97        Self {
98            model: self.model,
99            atoms: self.atoms.union(&other.atoms),
100        }
101    }
102
103    /// Intersection of two units (must reference the same model).
104    ///
105    /// # Panics
106    /// Panics if the units reference different models.
107    pub fn intersection(&self, other: &Unit<'a>) -> Self {
108        assert!(
109            std::ptr::eq(self.model, other.model),
110            "cannot intersect Units from different Models"
111        );
112        Self {
113            model: self.model,
114            atoms: self.atoms.intersection(&other.atoms),
115        }
116    }
117
118    /// Difference of two units: atoms in `self` but not in `other`.
119    ///
120    /// # Panics
121    /// Panics if the units reference different models.
122    pub fn difference(&self, other: &Unit<'a>) -> Self {
123        assert!(
124            std::ptr::eq(self.model, other.model),
125            "cannot difference Units from different Models"
126        );
127        Self {
128            model: self.model,
129            atoms: self.atoms.difference(&other.atoms),
130        }
131    }
132
133    /// Select atoms belonging to a specific chain.
134    ///
135    /// Returns `None` if the chain doesn't exist.
136    pub fn chain(model: &'a Model, chain_id: &str) -> Option<Self> {
137        let chain_idx = model
138            .hierarchy
139            .chains
140            .label_asym_id
141            .iter()
142            .position(|c| c == chain_id)?;
143
144        let res_range = model.hierarchy.residue_to_chain.segment(chain_idx);
145        if res_range.is_empty() {
146            return Some(Self {
147                model,
148                atoms: OrderedSet::interval(0, 0),
149            });
150        }
151        let atom_start = model.hierarchy.atom_to_residue.segment(res_range.start).start;
152        let atom_end = model.hierarchy.atom_to_residue.segment(res_range.end - 1).end;
153
154        Some(Self {
155            model,
156            atoms: OrderedSet::interval(atom_start as u32, atom_end as u32),
157        })
158    }
159
160    /// Select backbone atoms (N, CA, C, O).
161    pub fn backbone(model: &'a Model) -> Self {
162        let backbone_names = ["N", "CA", "C", "O"];
163        let indices: Vec<u32> = model
164            .hierarchy
165            .atoms
166            .atom_name
167            .iter()
168            .enumerate()
169            .filter_map(|(i, name)| {
170                if backbone_names.contains(&name.as_str()) {
171                    Some(i as u32)
172                } else {
173                    None
174                }
175            })
176            .collect();
177
178        let atoms = if indices.is_empty() {
179            OrderedSet::interval(0, 0)
180        } else {
181            OrderedSet::from_sorted(indices)
182        };
183
184        Self { model, atoms }
185    }
186}
187
188impl Model {
189    /// Create a filtered view of this model using a boolean mask.
190    ///
191    /// # Panics
192    /// Panics if `mask.len() != self.n_atoms()`.
193    pub fn filter(&self, mask: &[bool]) -> Unit<'_> {
194        Unit::from_mask(self, mask)
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use crate::data::Segmentation;
202    use crate::model::bonds::Bonds;
203    use crate::model::conformation::AtomicConformation;
204    use crate::model::hierarchy::AtomicHierarchy;
205    use crate::model::tables::{AtomsTable, ChainsTable, ResidueGroup, ResiduesTable};
206    use std::sync::Arc;
207
208    fn make_test_model(n_atoms: usize) -> Model {
209        let atoms = AtomsTable {
210            atom_name: (0..n_atoms).map(|i| format!("A{}", i)).collect(),
211            element: vec!["C".into(); n_atoms],
212            alt_loc: vec![None; n_atoms],
213            formal_charge: vec![None; n_atoms],
214        };
215        let residues = ResiduesTable {
216            comp_id: vec!["ALA".into()],
217            label_seq_id: vec![0],
218            auth_seq_id: vec![1],
219            ins_code: vec![None],
220            group: vec![ResidueGroup::Polymer],
221        };
222        let chains = ChainsTable {
223            label_asym_id: vec!["A".into()],
224            auth_asym_id: vec!["A".into()],
225            entity_id: vec!["1".into()],
226        };
227        let atom_to_residue = Segmentation::from_offsets(vec![0, n_atoms as u32]);
228        let residue_to_chain = Segmentation::from_offsets(vec![0, 1]);
229        let bonds = Bonds::from_unsorted(vec![], vec![], vec![], n_atoms);
230
231        let hierarchy = Arc::new(AtomicHierarchy {
232            atoms,
233            residues,
234            chains,
235            atom_to_residue,
236            residue_to_chain,
237            bonds,
238        });
239
240        let conformation = AtomicConformation {
241            x: (0..n_atoms).map(|i| i as f32).collect(),
242            y: (0..n_atoms).map(|i| i as f32 * 10.0).collect(),
243            z: (0..n_atoms).map(|_| 0.0).collect(),
244            occupancy: None,
245            b_iso: None,
246            confidence: None,
247        };
248
249        Model::new(hierarchy, conformation)
250    }
251
252    fn make_backbone_model() -> Model {
253        let atom_names = vec![
254            "N".into(),
255            "CA".into(),
256            "C".into(),
257            "O".into(),
258            "CB".into(),
259        ];
260        let n_atoms = atom_names.len();
261        let atoms = AtomsTable {
262            atom_name: atom_names,
263            element: vec!["C".into(); n_atoms],
264            alt_loc: vec![None; n_atoms],
265            formal_charge: vec![None; n_atoms],
266        };
267        let residues = ResiduesTable {
268            comp_id: vec!["ALA".into()],
269            label_seq_id: vec![0],
270            auth_seq_id: vec![1],
271            ins_code: vec![None],
272            group: vec![ResidueGroup::Polymer],
273        };
274        let chains = ChainsTable {
275            label_asym_id: vec!["A".into()],
276            auth_asym_id: vec!["A".into()],
277            entity_id: vec!["1".into()],
278        };
279        let atom_to_residue = Segmentation::from_offsets(vec![0, n_atoms as u32]);
280        let residue_to_chain = Segmentation::from_offsets(vec![0, 1]);
281        let bonds = Bonds::from_unsorted(vec![], vec![], vec![], n_atoms);
282
283        let hierarchy = Arc::new(AtomicHierarchy {
284            atoms,
285            residues,
286            chains,
287            atom_to_residue,
288            residue_to_chain,
289            bonds,
290        });
291
292        let conformation = AtomicConformation {
293            x: vec![0.0, 1.0, 2.0, 3.0, 4.0],
294            y: vec![0.0, 0.0, 0.0, 0.0, 0.0],
295            z: vec![0.0, 0.0, 0.0, 0.0, 0.0],
296            occupancy: None,
297            b_iso: None,
298            confidence: None,
299        };
300
301        Model::new(hierarchy, conformation)
302    }
303
304    #[test]
305    fn test_unit_from_mask_no_copy() {
306        let model = make_test_model(10);
307        let mask: Vec<bool> = (0..10).map(|i| i % 2 == 0).collect();
308
309        let unit = Unit::from_mask(&model, &mask);
310
311        assert_eq!(unit.len(), 5);
312        assert!(std::ptr::eq(unit.model(), &model));
313
314        let indices: Vec<u32> = unit.atom_indices().collect();
315        assert_eq!(indices, vec![0, 2, 4, 6, 8]);
316    }
317
318    #[test]
319    fn test_unit_all() {
320        let model = make_test_model(5);
321        let unit = Unit::all(&model);
322
323        assert_eq!(unit.len(), 5);
324        let indices: Vec<u32> = unit.atom_indices().collect();
325        assert_eq!(indices, vec![0, 1, 2, 3, 4]);
326    }
327
328    #[test]
329    fn test_unit_iteration_lazy() {
330        let model = make_test_model(100);
331        let mask: Vec<bool> = (0..100).map(|i| i < 3).collect();
332
333        let unit = Unit::from_mask(&model, &mask);
334
335        let coords: Vec<[f32; 3]> = unit.coords().collect();
336        assert_eq!(coords.len(), 3);
337        assert_eq!(coords[0], [0.0, 0.0, 0.0]);
338        assert_eq!(coords[1], [1.0, 10.0, 0.0]);
339        assert_eq!(coords[2], [2.0, 20.0, 0.0]);
340    }
341
342    #[test]
343    fn test_unit_set_operations() {
344        let model = make_test_model(10);
345
346        let mask_a: Vec<bool> = (0..10).map(|i| i < 5).collect();
347        let mask_b: Vec<bool> = (0..10).map(|i| i >= 3 && i < 8).collect();
348
349        let unit_a = Unit::from_mask(&model, &mask_a);
350        let unit_b = Unit::from_mask(&model, &mask_b);
351
352        let union = unit_a.union(&unit_b);
353        let union_indices: Vec<u32> = union.atom_indices().collect();
354        assert_eq!(union_indices, vec![0, 1, 2, 3, 4, 5, 6, 7]);
355
356        let intersection = unit_a.intersection(&unit_b);
357        let int_indices: Vec<u32> = intersection.atom_indices().collect();
358        assert_eq!(int_indices, vec![3, 4]);
359
360        let difference = unit_a.difference(&unit_b);
361        let diff_indices: Vec<u32> = difference.atom_indices().collect();
362        assert_eq!(diff_indices, vec![0, 1, 2]);
363    }
364
365    #[test]
366    fn test_unit_empty() {
367        let model = make_test_model(5);
368        let mask = vec![false; 5];
369
370        let unit = Unit::from_mask(&model, &mask);
371        assert!(unit.is_empty());
372        assert_eq!(unit.len(), 0);
373    }
374
375    #[test]
376    fn test_model_filter() {
377        let model = make_test_model(5);
378        let mask = vec![true, false, true, false, true];
379
380        let unit = model.filter(&mask);
381        assert_eq!(unit.len(), 3);
382    }
383
384    #[test]
385    fn test_unit_chain() {
386        let model = make_test_model(5);
387
388        let chain_a = Unit::chain(&model, "A");
389        assert!(chain_a.is_some());
390        assert_eq!(chain_a.unwrap().len(), 5);
391
392        let chain_b = Unit::chain(&model, "B");
393        assert!(chain_b.is_none());
394    }
395
396    #[test]
397    fn test_unit_backbone() {
398        let model = make_backbone_model();
399
400        let backbone = Unit::backbone(&model);
401        assert_eq!(backbone.len(), 4);
402
403        let indices: Vec<u32> = backbone.atom_indices().collect();
404        assert_eq!(indices, vec![0, 1, 2, 3]);
405    }
406
407    #[test]
408    #[should_panic(expected = "mask length must equal number of atoms")]
409    fn test_unit_from_mask_wrong_length() {
410        let model = make_test_model(5);
411        let mask = vec![true, false, true];
412        Unit::from_mask(&model, &mask);
413    }
414}
ferritin_core/unit/unit.rs

ferritin_core/unit/
unit.rs
