Skip to main content

ferritin_core/trajectory/
coordinates.rs

1//! Coordinate data for trajectory frames.
2//!
3//! [`Frame`] holds a single snapshot of coordinates (topology-free), and
4//! [`Coordinates`] is an ordered collection of frames for trajectory storage.
5
6use crate::model::ModelError;
7
8/// Unit cell parameters for periodic boundary condition systems.
9#[derive(Clone, Debug, PartialEq)]
10pub struct UnitCell {
11    /// Length of the a-axis in Å.
12    pub a: f32,
13    /// Length of the b-axis in Å.
14    pub b: f32,
15    /// Length of the c-axis in Å.
16    pub c: f32,
17    /// Angle between b and c axes in degrees.
18    pub alpha: f32,
19    /// Angle between a and c axes in degrees.
20    pub beta: f32,
21    /// Angle between a and b axes in degrees.
22    pub gamma: f32,
23}
24
25/// Single frame of coordinates — topology-free.
26///
27/// Stores Cartesian coordinates as parallel arrays (SoA layout). Optional
28/// fields are `None` when not relevant (e.g. non-periodic systems have no
29/// `cell`; single-structure files have no `time`).
30#[derive(Clone, Debug)]
31pub struct Frame {
32    /// X coordinates (Å), one per atom.
33    pub x: Vec<f32>,
34    /// Y coordinates (Å), one per atom.
35    pub y: Vec<f32>,
36    /// Z coordinates (Å), one per atom.
37    pub z: Vec<f32>,
38    /// Periodic box parameters, or `None` for non-periodic systems.
39    pub cell: Option<UnitCell>,
40    /// Simulation time in picoseconds, or `None` if not applicable.
41    pub time: Option<f64>,
42}
43
44/// Ordered collection of coordinate frames.
45///
46/// Frames are indexed by position. Use [`Coordinates::frame`] for random access.
47#[derive(Clone, Debug)]
48pub struct Coordinates {
49    frames: Vec<Frame>,
50}
51
52impl Coordinates {
53    /// Construct a validated coordinate collection.
54    pub fn try_new(frames: Vec<Frame>) -> Result<Self, ModelError> {
55        let expected = frames.first().map_or(0, |frame| frame.x.len());
56        for (frame_index, frame) in frames.iter().enumerate() {
57            for (axis, len) in [
58                ("x", frame.x.len()),
59                ("y", frame.y.len()),
60                ("z", frame.z.len()),
61            ] {
62                if len != expected {
63                    return Err(ModelError::new(format!(
64                        "coordinates frame {frame_index}.{axis} has length {len}, expected {expected}"
65                    )));
66                }
67            }
68        }
69        Ok(Self { frames })
70    }
71
72    /// Construct from a `Vec<Frame>`.
73    ///
74    /// # Panics
75    /// Panics if coordinate column lengths differ. Use [`Coordinates::try_new`]
76    /// to report invalid input without panicking.
77    pub fn new(frames: Vec<Frame>) -> Self {
78        Self::try_new(frames).expect("invalid coordinate frames")
79    }
80
81    /// Number of frames.
82    pub fn len(&self) -> usize {
83        self.frames.len()
84    }
85
86    /// Returns `true` if there are no frames.
87    pub fn is_empty(&self) -> bool {
88        self.frames.is_empty()
89    }
90
91    /// Returns a reference to the frame at `index`.
92    ///
93    /// # Panics
94    /// Panics if `index >= len()`.
95    pub fn frame(&self, index: usize) -> &Frame {
96        &self.frames[index]
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_coordinates_len() {
106        let frames: Vec<Frame> = (0..3)
107            .map(|i| Frame {
108                x: vec![i as f32],
109                y: vec![0.0],
110                z: vec![0.0],
111                cell: None,
112                time: None,
113            })
114            .collect();
115        let coords = Coordinates::new(frames);
116        assert_eq!(coords.len(), 3);
117        assert!(!coords.is_empty());
118    }
119
120    #[test]
121    fn test_frame_fields() {
122        let cell = UnitCell {
123            a: 10.0,
124            b: 20.0,
125            c: 30.0,
126            alpha: 90.0,
127            beta: 90.0,
128            gamma: 120.0,
129        };
130        let frame = Frame {
131            x: vec![1.0, 2.0],
132            y: vec![3.0, 4.0],
133            z: vec![5.0, 6.0],
134            cell: Some(cell.clone()),
135            time: Some(0.5),
136        };
137
138        assert_eq!(frame.x, vec![1.0, 2.0]);
139        assert_eq!(frame.y, vec![3.0, 4.0]);
140        assert_eq!(frame.z, vec![5.0, 6.0]);
141        assert_eq!(frame.time, Some(0.5));
142        let c = frame.cell.as_ref().unwrap();
143        assert_eq!(c.a, 10.0);
144        assert_eq!(c.b, 20.0);
145        assert_eq!(c.c, 30.0);
146        assert_eq!(c.alpha, 90.0);
147        assert_eq!(c.beta, 90.0);
148        assert_eq!(c.gamma, 120.0);
149    }
150
151    #[test]
152    fn test_coordinates_validation_rejects_frame_size_mismatch() {
153        let frames = vec![
154            Frame {
155                x: vec![0.0],
156                y: vec![0.0],
157                z: vec![0.0],
158                cell: None,
159                time: None,
160            },
161            Frame {
162                x: vec![0.0, 1.0],
163                y: vec![0.0, 1.0],
164                z: vec![0.0, 1.0],
165                cell: None,
166                time: None,
167            },
168        ];
169        assert!(Coordinates::try_new(frames).is_err());
170    }
171}