ferritin_core/unit/
unit.rs

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

ferritin_core/unit/
unit.rs
