ferritin_core/trajectory/
array_trajectory.rs

1//! Eager trajectory: all frames held in memory as `Vec<Model>`.
2
3use std::borrow::Cow;
4use std::sync::Arc;
5use crate::model::Model;
6use super::Trajectory;
7
8/// Error type for [`ArrayTrajectory`] construction.
9#[derive(Debug, Clone)]
10pub enum TrajectoryError {
11    /// Trajectory must contain at least one frame.
12    Empty,
13    /// All frames must share the same `Arc<AtomicHierarchy>`.
14    MixedTopologies,
15}
16
17impl std::fmt::Display for TrajectoryError {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        match self {
20            TrajectoryError::Empty => write!(f, "trajectory must contain at least one frame"),
21            TrajectoryError::MixedTopologies => {
22                write!(f, "all frames must share the same Arc<AtomicHierarchy>")
23            }
24        }
25    }
26}
27
28impl std::error::Error for TrajectoryError {}
29
30/// Eager trajectory: all models held in memory.
31///
32/// All frames **must** share the same `Arc<AtomicHierarchy>` (verified on
33/// construction). This invariant ensures that topology does not vary across
34/// frames, which is required for correct trajectory analysis.
35pub struct ArrayTrajectory {
36    models: Vec<Model>,
37}
38
39impl ArrayTrajectory {
40    /// Construct from a `Vec<Model>`.
41    ///
42    /// Returns `Err(TrajectoryError::Empty)` if `models` is empty.
43    /// Returns `Err(TrajectoryError::MixedTopologies)` if any two models have
44    /// different `Arc<AtomicHierarchy>` pointers.
45    pub fn new(models: Vec<Model>) -> Result<Self, TrajectoryError> {
46        if models.is_empty() {
47            return Err(TrajectoryError::Empty);
48        }
49        let first = Arc::as_ptr(&models[0].hierarchy);
50        for m in &models[1..] {
51            if Arc::as_ptr(&m.hierarchy) != first {
52                return Err(TrajectoryError::MixedTopologies);
53            }
54        }
55        Ok(Self { models })
56    }
57
58    /// Returns a slice of all models in the trajectory.
59    pub fn models(&self) -> &[Model] {
60        &self.models
61    }
62}
63
64impl Trajectory for ArrayTrajectory {
65    fn frame_count(&self) -> usize {
66        self.models.len()
67    }
68
69    fn representative(&self) -> &Model {
70        &self.models[0]
71    }
72
73    fn frame(&self, index: usize) -> Cow<'_, Model> {
74        Cow::Borrowed(&self.models[index])
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use std::borrow::Cow;
82    use crate::data::Segmentation;
83    use crate::model::{AtomicConformation, AtomicHierarchy, Bonds};
84    use crate::model::tables::{AtomsTable, ChainsTable, ResidueGroup, ResiduesTable};
85
86    fn make_simple_hierarchy(n_residues: usize) -> Arc<AtomicHierarchy> {
87        let n_atoms = n_residues;
88        let atoms = AtomsTable {
89            atom_name: (0..n_atoms).map(|_| "CA".to_string()).collect(),
90            element: (0..n_atoms).map(|_| "C".to_string()).collect(),
91            alt_loc: vec![None; n_atoms],
92            formal_charge: vec![None; n_atoms],
93        };
94        let residues = ResiduesTable {
95            comp_id: (0..n_residues).map(|i| format!("R{}", i)).collect(),
96            label_seq_id: (0..n_residues as i32).collect(),
97            auth_seq_id: (1..=n_residues as i32).collect(),
98            ins_code: vec![None; n_residues],
99            group: vec![ResidueGroup::Polymer; n_residues],
100        };
101        let chains = ChainsTable {
102            label_asym_id: vec!["A".into()],
103            auth_asym_id: vec!["A".into()],
104            entity_id: vec!["1".into()],
105        };
106        let atom_offsets: Vec<u32> = (0..=n_residues as u32).collect();
107        let atom_to_residue = Segmentation::from_offsets(atom_offsets);
108        let residue_to_chain = Segmentation::from_offsets(vec![0, n_residues as u32]);
109        let bonds = Bonds::from_unsorted(vec![], vec![], vec![], n_atoms);
110        Arc::new(AtomicHierarchy { atoms, residues, chains, atom_to_residue, residue_to_chain, bonds })
111    }
112
113    fn make_conformation(n: usize, offset: f32) -> AtomicConformation {
114        AtomicConformation {
115            x: (0..n).map(|i| i as f32 + offset).collect(),
116            y: (0..n).map(|i| i as f32 * 10.0).collect(),
117            z: vec![0.0; n],
118            occupancy: None,
119            b_iso: None,
120            confidence: None,
121        }
122    }
123
124    #[test]
125    fn test_array_trajectory_frame_count() {
126        let hierarchy = make_simple_hierarchy(3);
127        let models: Vec<Model> = (0..3)
128            .map(|i| Model::new(Arc::clone(&hierarchy), make_conformation(3, i as f32 * 100.0)))
129            .collect();
130        let traj = ArrayTrajectory::new(models).unwrap();
131        assert_eq!(traj.frame_count(), 3);
132    }
133
134    #[test]
135    fn test_array_trajectory_borrowed() {
136        let hierarchy = make_simple_hierarchy(2);
137        let models: Vec<Model> = (0..3)
138            .map(|i| Model::new(Arc::clone(&hierarchy), make_conformation(2, i as f32 * 10.0)))
139            .collect();
140        let traj = ArrayTrajectory::new(models).unwrap();
141        for i in 0..3 {
142            let f = traj.frame(i);
143            assert!(matches!(f, Cow::Borrowed(_)), "frame({}) should be Cow::Borrowed", i);
144        }
145    }
146
147    #[test]
148    fn test_array_trajectory_arc_shared() {
149        let hierarchy = make_simple_hierarchy(4);
150        let models: Vec<Model> = (0..3)
151            .map(|i| Model::new(Arc::clone(&hierarchy), make_conformation(4, i as f32 * 5.0)))
152            .collect();
153        let traj = ArrayTrajectory::new(models).unwrap();
154        let h0 = &traj.frame(0).hierarchy;
155        let h2 = &traj.frame(2).hierarchy;
156        assert!(Arc::ptr_eq(h0, h2), "frame(0) and frame(2) must share the same Arc<AtomicHierarchy>");
157    }
158
159    #[test]
160    fn test_array_trajectory_empty_error() {
161        let result = ArrayTrajectory::new(vec![]);
162        assert!(matches!(result, Err(TrajectoryError::Empty)));
163    }
164
165    #[test]
166    fn test_array_trajectory_mixed_topologies_error() {
167        let h1 = make_simple_hierarchy(3);
168        let h2 = make_simple_hierarchy(3); // different Arc allocation
169        let models = vec![
170            Model::new(Arc::clone(&h1), make_conformation(3, 0.0)),
171            Model::new(Arc::clone(&h2), make_conformation(3, 1.0)),
172        ];
173        let result = ArrayTrajectory::new(models);
174        assert!(matches!(result, Err(TrajectoryError::MixedTopologies)));
175    }
176
177    #[test]
178    fn test_array_trajectory_object_safe() {
179        let hierarchy = make_simple_hierarchy(2);
180        let models: Vec<Model> = (0..2)
181            .map(|i| Model::new(Arc::clone(&hierarchy), make_conformation(2, i as f32)))
182            .collect();
183        let traj = ArrayTrajectory::new(models).unwrap();
184        let boxed: Box<dyn Trajectory> = Box::new(traj);
185        assert_eq!(boxed.frame_count(), 2);
186    }
187}
ferritin_core/trajectory/array_trajectory.rs

ferritin_core/trajectory/
array_trajectory.rs
