ferritin_core/trajectory/
array_trajectory.rs

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

ferritin_core/trajectory/
array_trajectory.rs
