1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
use crate::*;
use std::rc::Rc;
fn into_iter_inner<T: Trajectory>(mut traj: T) -> TrajectoryIterator<T> {
let num_atoms = traj.get_num_atoms();
let frame = match &num_atoms {
Ok(num_atoms) => Frame::with_len(*num_atoms),
Err(_) => Frame::new(),
};
TrajectoryIterator {
trajectory: traj,
item: Rc::new(frame),
has_error: false,
}
}
impl IntoIterator for XTCTrajectory {
type Item = Result<Rc<Frame>>;
type IntoIter = TrajectoryIterator<XTCTrajectory>;
fn into_iter(self) -> Self::IntoIter {
into_iter_inner(self)
}
}
impl IntoIterator for TRRTrajectory {
type Item = Result<Rc<Frame>>;
type IntoIter = TrajectoryIterator<TRRTrajectory>;
fn into_iter(self) -> Self::IntoIter {
into_iter_inner(self)
}
}
pub struct TrajectoryIterator<T> {
trajectory: T,
item: Rc<Frame>,
has_error: bool,
}
impl<T: Trajectory> TrajectoryIterator<T> {
fn next_inner(&mut self) -> <Self as Iterator>::Item {
let num_atoms = match &self.trajectory.get_num_atoms() {
&Ok(n) => n,
Err(e) => Err(Error::CouldNotCheckNAtoms(Box::new(e.clone())))?,
};
let item: &mut Frame = match Rc::get_mut(&mut self.item) {
Some(item) => item,
None => {
self.item = Rc::new(Frame::with_len(num_atoms as usize));
Rc::get_mut(&mut self.item).expect("Could not get mutable access to new Rc")
}
};
self.trajectory.read(item)?;
Ok(Rc::clone(&self.item))
}
}
impl<T> Iterator for TrajectoryIterator<T>
where
T: Trajectory,
{
type Item = Result<Rc<Frame>>;
fn next(&mut self) -> Option<Self::Item> {
if self.has_error {
return None;
}
match self.next_inner() {
Ok(item) => Some(Ok(item)),
Err(e) if e.is_eof() => None,
Err(e) => {
self.has_error = true;
Some(Err(e))
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
pub fn test_xtc_trajectory_iterator() -> Result<()> {
let traj = XTCTrajectory::open_read("tests/1l2y.xtc")?;
let frames: Result<Vec<Rc<Frame>>> = traj.into_iter().collect();
let frames = frames?;
assert!(frames.len() == 38);
assert!(frames[0].step == 1, frames[0].step);
assert!(frames[37].step == 38);
Ok(())
}
#[test]
pub fn test_trr_trajectory_iterator() -> Result<()> {
let traj = TRRTrajectory::open_read("tests/1l2y.trr")?;
let frames: Result<Vec<Rc<Frame>>> = traj.into_iter().collect();
let frames = frames?;
assert!(frames.len() == 38);
assert!(frames[0].step == 1, frames[0].step);
assert!(frames[37].step == 38);
Ok(())
}
}