From ceb09f637110216863a99cc8f322ec723725da6e Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 16 Apr 2020 15:00:20 +0200 Subject: [PATCH] iterators --- src/iterator.rs | 139 +++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 27 ++++++++- tests/integration.rs | 32 ++++++++++ 3 files changed, 196 insertions(+), 2 deletions(-) create mode 100644 src/iterator.rs create mode 100644 tests/integration.rs diff --git a/src/iterator.rs b/src/iterator.rs new file mode 100644 index 0000000..586e2e2 --- /dev/null +++ b/src/iterator.rs @@ -0,0 +1,139 @@ +use crate::*; +use crate::c_abi::xdrfile::exdrENDOFFILE; +use failure::Error; +use std::rc::Rc; + +impl IntoIterator for XTCTrajectory { + type Item = Result, Error>; + type IntoIter = XTCTrajectoryIterator; + + fn into_iter(mut self) -> Self::IntoIter { + // TODO this should be handled without the requirement of unwrap + let num_atoms = self.get_num_atoms().unwrap(); + XTCTrajectoryIterator { + trajectory: self, + item: Rc::new(Frame::with_capacity(num_atoms)), + has_error: false + } + } +} + +/* +Iterator for trajectories. This iterator yields a Result +for each frame in the trajectory file and stops with yielding None once the +trajectory is finished. Also yields None after the first occurence of an error +*/ +pub struct XTCTrajectoryIterator { + trajectory: XTCTrajectory, + item: Rc, + has_error: bool, +} + +impl Iterator for XTCTrajectoryIterator { + type Item = Result, Error>; + + fn next(&mut self) -> Option { // Reuse old frame + if self.has_error { + return None; + } + + let item: &mut Frame = match Rc::get_mut(&mut self.item) { + Some(item) => item, + None => { // caller kept frame. Create new one + self.item = Rc::new(Frame::with_capacity(self.item.num_atoms)); + Rc::get_mut(&mut self.item).unwrap() + } + + }; + match self.trajectory.read(item) { + Ok(()) => Some(Ok(Rc::clone(&self.item))), + Err(msg) => { + if msg.to_string().contains(&exdrENDOFFILE.to_string()) { + None + } else { + self.has_error = true; + Some(Err(msg)) + } + } + } + } +} + +impl IntoIterator for TRRTrajectory { + type Item = Result, Error>; + type IntoIter = TRRTrajectoryIterator; + + fn into_iter(mut self) -> Self::IntoIter { + // TODO this should be handled without the requirement of unwrap + let num_atoms = self.get_num_atoms().unwrap(); + TRRTrajectoryIterator { + trajectory: self, + item: Rc::new(Frame::with_capacity(num_atoms)), + has_error: false + } + } +} + +/* +Iterator for trajectories. This iterator yields a Result +for each frame in the trajectory file and stops with yielding None once the +trajectory is finished. Also yields None after the first occurence of an error +*/ +pub struct TRRTrajectoryIterator { + trajectory: TRRTrajectory, + item: Rc, + has_error: bool, +} + +impl Iterator for TRRTrajectoryIterator { + type Item = Result, Error>; + + fn next(&mut self) -> Option { // Reuse old frame + if self.has_error { + return None; + } + + let item: &mut Frame = match Rc::get_mut(&mut self.item) { + Some(item) => item, + None => { // caller kept frame. Create new one + self.item = Rc::new(Frame::with_capacity(self.item.num_atoms)); + Rc::get_mut(&mut self.item).unwrap() + } + + }; + match self.trajectory.read(item) { + Ok(()) => Some(Ok(Rc::clone(&self.item))), + Err(msg) => { + if msg.to_string().contains(&exdrENDOFFILE.to_string()) { + None + } else { + self.has_error = true; + Some(Err(msg)) + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + pub fn test_xtc_trajectory_iterator() { + let traj = XTCTrajectory::open(Path::new("tests/1l2y.xtc"), FileMode::Read).unwrap(); + let frames: Vec> = traj.into_iter().filter_map(Result::ok).collect(); + assert!(frames.len() == 38); + assert!(frames[0].step == 1, frames[0].step); + assert!(frames[37].step == 38); + } + + #[test] + pub fn test_trr_trajectory_iterator() { + let traj = TRRTrajectory::open(Path::new("tests/1l2y.trr"), FileMode::Read).unwrap(); + let frames: Vec> = traj.into_iter().filter_map(Result::ok).collect(); + assert!(frames.len() == 38); + assert!(frames[0].step == 1, frames[0].step); + assert!(frames[37].step == 38); + } +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index c8ca9bd..0116968 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,7 +21,7 @@ //! // without instantiating data arrays for every step //! let mut frame = Frame::with_capacity(num_atoms); //! -//! // read the first trame of the trajectory +//! // read the first frame of the trajectory //! let result = trj.read(&mut frame); //! match result { //! Ok(_) => { @@ -35,7 +35,28 @@ //! panic!("Something went wrong: {}", msg); //! } //! } -//! ``` +//! ``` +//! +//! # Frame iteration +//! For convenience, the trajectory implementations provide "into_iter" to +//! be turned into an iterator that yields Rc. If a frame is not kept +//! during iteration, the Iterator reuses it for better performance (and hence, +//! Rc is required) +//! +//! ```rust +//! use xdrfile::*; +//! use std::path::Path; +//! +//! let mut path = Path::new("tests/1l2y.xtc"); +//! // get a handle to the file +//! let trj = XTCTrajectory::open(path, FileMode::Read).unwrap(); +//! +//! // iterate over all frames +//! for (idx, frame) in trj.into_iter().filter_map(Result::ok).enumerate() { +//! println!("{}", frame.time); +//! assert_eq!(idx+1, frame.step as usize); +//! } +//! ``` #[cfg(test)] @@ -43,8 +64,10 @@ extern crate lazy_init; mod frame; +mod iterator; pub mod c_abi; pub use frame::Frame; +pub use iterator::*; use c_abi::xdrfile; use c_abi::xdrfile::XDRFILE; diff --git a/tests/integration.rs b/tests/integration.rs new file mode 100644 index 0000000..70e0b2b --- /dev/null +++ b/tests/integration.rs @@ -0,0 +1,32 @@ +#[cfg(test)] +mod integration { + + use std::rc::Rc; + use xdrfile::*; + use std::path::Path; + + #[test] + fn test_use_library() { + let path = Path::new("tests/1l2y.xtc"); + + let mut trj = XTCTrajectory::open(path, FileMode::Read).unwrap(); + let num_atoms = trj.get_num_atoms().unwrap(); + let mut frame = Frame::with_capacity(num_atoms); + + trj.read(&mut frame).unwrap(); + trj.read(&mut frame).unwrap(); + assert_eq!(frame.step, 2); + } + + #[test] + fn test_use_library_iterator() { + let path = Path::new("tests/1l2y.xtc"); + + let trj = XTCTrajectory::open(path, FileMode::Read).unwrap(); + let frames: Vec> = trj.into_iter().filter_map(Result::ok).collect(); + for (idx, frame) in frames.iter().enumerate() { + assert_eq!(frame.step as usize, idx+1); + } + } + +} \ No newline at end of file