XTCTrajectory and TRRTrajectory

src/frame.rs

@@ -0,0 +1,103 @@
+use std::fmt;
+
+/// Representation of a single frame of an MD trajectory
+#[derive(Clone)]
+pub struct Frame {
+    /// Number of atoms in the frame
+    pub num_atoms: u32,
+
+    /// Trajectory step
+    pub step: u32,
+
+    /// Time step
+    pub time: f32,
+
+    /// 3x3 box vector 
+    pub box_vector: [[f32; 3usize]; 3usize],
+    
+    /// 3D coordinates for N atoms where N is num_atoms 
+    pub coords: Vec<[f32; 3usize]>,
+}
+
+impl Default for Frame {
+    fn default() -> Frame {
+        return Frame { 
+            num_atoms: 0,
+            step: 0,
+            time: 0.0,
+            box_vector: [[0.0; 3]; 3],
+            coords: Vec::with_capacity(0)
+        };
+    }
+}
+
+impl fmt::Debug for Frame {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "Frame {{ atoms: {}, step: {}, time: {}, \
+            box: {:?}, coords: {:?} }}", self.num_atoms, self.step, self.time,
+            self.box_vector, self.coords)
+    }
+}
+
+impl Frame {
+    pub fn new() -> Frame {
+        Frame{ ..Default::default() }
+    }
+
+    pub fn with_capacity(num_atoms: u32) -> Frame {
+        Frame {
+            num_atoms: num_atoms,
+            coords: vec![[0.0, 0.0, 0.0]; num_atoms as usize],
+            ..Default::default()
+        }
+
+    }
+
+    pub fn filter_coords(self: &mut Frame, indeces: &[usize]) {
+        self.coords = self.coords.iter() 
+            .map(|elem| elem.clone())
+            .enumerate()
+            .filter(|&(i, _)| indeces.contains(&i))
+            .map(|(_, elem)| elem)
+            .collect();
+        self.num_atoms = self.coords.len() as u32;
+    }   
+
+    pub fn len(self: &Frame) -> usize {
+        self.num_atoms as usize
+    } 
+}
+
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_frame_with_capacity() {
+        let frame = Frame::with_capacity(10);
+        println!("{:?}", frame.coords);
+        assert_eq!(frame.coords.len(), 10);
+    }
+
+    #[test]
+    fn test_frame_filter_atoms() {
+        let mut frame = Frame::with_capacity(3);
+        frame.coords[0] = [1.0, 2.0, 3.0];
+        frame.coords[1] = [4.0, 5.0, 6.0];
+        frame.coords[2] = [7.0, 8.0, 9.0];
+        let filter: Vec<usize> = vec![1,2];
+        let mut frame_new = frame.clone();
+        frame_new.filter_coords(&filter);
+        assert!(frame_new.num_atoms as usize == filter.len());
+        assert!(frame_new.coords[0] == frame.coords[1]);
+        assert!(frame_new.coords[1] == frame.coords[2]);           
+    }
+
+    #[test]
+    fn test_frame_len() {
+        let frame = Frame::with_capacity(10);
+        assert_eq!(frame.len(), 10);
+    }
+}

src/lib.rs

@@ -1 +1,339 @@
+#[cfg(test)]
+#[macro_use] extern crate assert_approx_eq;
+
+mod frame;
 pub mod c_abi; 
+pub use frame::Frame;
+
+use c_abi::xdrfile;
+use c_abi::xdrfile::XDRFILE;
+use c_abi::xdr_seek;
+use c_abi::xdrfile_xtc;
+use c_abi::xdrfile_trr;
+
+use std::ffi::CString;
+use std::path::Path;
+use failure::{Error,err_msg};
+use std::cell::Cell;
+use lazy_init::Lazy;
+
+pub enum FileMode {
+    Write,
+    Append,
+    Read    
+}
+
+impl FileMode {
+    pub fn value(&self) -> &str {
+        return match *self {
+            FileMode::Write => "w",
+            FileMode::Append => "a",
+            FileMode::Read => "r",
+        }
+    }
+}
+
+fn path_to_cstring(path: &Path) -> CString {
+        CString::new(path.to_str().unwrap()).unwrap()
+    }
+
+struct XDRFile {
+    xdrfile: *mut XDRFILE,
+    filemode: FileMode,
+    path: String,
+}
+
+impl XDRFile {
+
+    pub fn open(path: &Path, filemode: FileMode) -> Result<XDRFile, Error> {
+        let path_p = path_to_cstring(path).into_raw();
+        let mode_p = CString::new(filemode.value()).unwrap().into_raw();
+
+        unsafe {
+            let xdrfile = xdrfile::xdrfile_open(path_p, mode_p);
+        
+            if ! xdrfile.is_null() {
+                let path = String::from(path.to_str().unwrap());
+                return Ok(XDRFile { xdrfile, filemode, path });
+            } else {  // Something went wrong. But the C api does not tell us what
+                return Err(err_msg("Failed to open trajectory file"));
+            }
+        }
+    }
+}
+
+impl Drop for XDRFile {
+    // Close the underlying xdr file on drop
+    fn drop(&mut self) {
+        unsafe {
+            xdrfile::xdrfile_close(self.xdrfile);
+        }
+    } 
+}
+
+pub trait Trajectory {
+    fn read(&mut self, frame: &mut Frame) -> Result<(), Error>;
+    fn write(&mut self, frame: &Frame) -> Result<(), Error>;
+    fn flush(&mut self) -> Result<(), Error>;
+    fn get_num_atoms(&mut self) -> Result<u32, Error>; 
+}
+
+pub struct XTCTrajectory {
+    handle: XDRFile,
+    precision: Cell<f32>,  // internal mutability required for read method
+    num_atoms: Lazy<Result<u32, Error>>
+}
+
+impl XTCTrajectory {
+    pub fn open(path: &Path, filemode: FileMode) -> Result<XTCTrajectory, Error> {
+        let xdr = XDRFile::open(path, filemode)?;
+        Ok(XTCTrajectory { handle: xdr, precision: Cell::new(1000.0), num_atoms: Lazy::new() })
+    }
+}
+
+impl Trajectory for XTCTrajectory {
+
+    fn read(&mut self, frame: &mut Frame) -> Result<(), Error> {
+        unsafe {
+            // C lib requires an i32 to be passed, but step is exposed it as u32
+            // (A step cannot be negative, can it?). So we need to create a step
+            // variable to pass to read_xtc and cast it afterwards to u32
+            let mut step: i32 = 0;
+            let code = xdrfile_xtc::read_xtc(self.handle.xdrfile,
+                frame.num_atoms as i32,
+                &mut step,
+                &mut frame.time,
+                &mut frame.box_vector,
+                frame.coords.as_ptr() as *mut [f32; 3],
+                &mut self.precision.get(),
+                ) as u32;
+            frame.step = step as u32;
+            match code {
+                xdrfile::exdrOK => Ok(()),
+                _ => Err(err_msg(format!("Failed to read trajectory. Error code: {}", code))), 
+            }
+        }
+    }
+
+    fn write(&mut self, frame: &Frame) -> Result<(), Error> {
+        unsafe {
+            let code = xdrfile_xtc::write_xtc(self.handle.xdrfile,
+                frame.num_atoms as i32,
+                frame.step as i32,
+                frame.time,
+                frame.box_vector.as_ptr() as *mut [[f32; 3]; 3],
+                frame.coords[..].as_ptr() as *mut [f32; 3],
+                1000.0) as u32;
+            match code {
+                xdrfile::exdrOK => Ok(()),
+                _ => Err(err_msg(format!("Failed to write trajectory. Error code: {}", code)))
+            }
+        }
+    }
+
+    fn flush(&mut self) -> Result<(), Error> {
+        unsafe {
+            let code = xdr_seek::xdr_flush(self.handle.xdrfile) as u32;
+            match code {
+                xdrfile::exdrOK => Ok(()),
+                _ => Err(err_msg(format!("Failed to flush trajectory. Error code: {}", code)))
+            }
+        } 
+    }
+
+    fn get_num_atoms(&mut self) -> Result<u32, Error> {
+        let result = self.num_atoms.get_or_create(|| {
+            let mut num_atoms: i32 = 0;
+            unsafe {
+                let path = CString::new(self.handle.path.as_str()).unwrap();
+                let path_p = path.into_raw();
+                let code = xdrfile_xtc::read_xtc_natoms(path_p, &mut num_atoms as *const i32) as u32;
+                match code {
+                    xdrfile::exdrOK => Ok(num_atoms as u32),
+                    _ => Err(err_msg(format!("Failed to read atom number from trajectory. Error code: {}", code)))
+                }
+            }
+        });
+        match result {
+            Ok(val) => Ok(*val),
+            // ugly hack because failure::Error is not "Clone"
+            Err(err) => Err(err_msg(format!("{}", err)))
+        }
+    }
+}
+
+pub struct TRRTrajectory {
+    handle: XDRFile,
+    num_atoms: Lazy<Result<u32, Error>>
+}
+
+impl TRRTrajectory {
+    pub fn open(path: &Path, filemode: FileMode) -> Result<TRRTrajectory, Error> {
+        let xdr = XDRFile::open(path, filemode)?;
+        Ok(TRRTrajectory { handle: xdr, num_atoms: Lazy::new() })
+    }
+}
+
+impl Trajectory for TRRTrajectory {
+
+    fn read(&mut self, frame: &mut Frame) -> Result<(), Error> {
+        unsafe {
+            // C lib requires an i32 to be passed, but step is exposed it as u32
+            // (A step cannot be negative, can it?). So we need to create a step
+            // variable to pass to read_trr and cast it afterwards to u32.
+            // Similar for lambda.
+            let mut step: i32 = 0;
+            let mut lambda: f32 = 0.0;
+            let code = xdrfile_trr::read_trr(self.handle.xdrfile,
+                frame.num_atoms as i32,
+                &mut step,
+                &mut frame.time,
+                &mut lambda,
+                &mut frame.box_vector,
+                frame.coords.as_ptr() as *mut [f32; 3],
+                std::ptr::null_mut(),
+                std::ptr::null_mut()
+                ) as u32;
+            frame.step = step as u32;
+            match code {
+                xdrfile::exdrOK => Ok(()),
+                _ => Err(err_msg(format!("Failed to read trajectory. Error code: {}", code)))
+            }
+        }
+    }
+
+    fn write(&mut self, frame: &Frame) -> Result<(), Error> {
+        unsafe {
+            let code = xdrfile_trr::write_trr(self.handle.xdrfile,
+                frame.num_atoms as i32,
+                frame.step as i32,
+                frame.time,
+                0.0,
+                frame.box_vector.as_ptr() as *mut [[f32; 3]; 3],
+                frame.coords[..].as_ptr() as *mut [f32; 3],
+                std::ptr::null_mut(),
+                std::ptr::null_mut()) as u32;
+            match code {
+                xdrfile::exdrOK => Ok(()),
+                _ => Err(err_msg(format!("Failed to write trajectory. Error code: {}", code)))
+            }
+        }
+    }
+
+    fn flush(&mut self) -> Result<(), Error> {
+        unsafe {
+            let code = xdr_seek::xdr_flush(self.handle.xdrfile) as u32;
+            match code {
+                xdrfile::exdrOK => Ok(()),
+                _ => Err(err_msg(format!("Failed to flush trajectory. Error code: {}", code))) 
+            }
+        } 
+    }
+
+    fn get_num_atoms(&mut self) -> Result<u32, Error> {
+        let result = self.num_atoms.get_or_create(|| {
+            let mut num_atoms: i32 = 0;
+            unsafe {
+                let path = CString::new(self.handle.path.as_str()).unwrap();
+                let path_p = path.into_raw();
+                let code = xdrfile_trr::read_trr_natoms(path_p, &mut num_atoms as *const i32) as u32;
+                match code {
+                    xdrfile::exdrOK => Ok(num_atoms as u32),
+                    _ => Err(err_msg(format!("Failed to read atom number from trajectory. Error code: {}", code)))
+                }
+            }
+        });
+        match result {
+            Ok(val) => Ok(*val),
+            // ugly hack because failure::Error is not "Clone"
+            Err(err) => Err(err_msg(format!("{}", err)))
+        }
+    }
+}
+
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+    use tempfile::NamedTempFile; 
+
+    #[test]
+    fn test_read_write_xtc() {
+        let tempfile = NamedTempFile::new().unwrap();
+        let tmp_path = tempfile.path();
+
+        let natoms: u32 = 2;
+        let frame = Frame {
+            num_atoms: natoms,
+            step: 5,
+            time: 2.0,
+            box_vector: [[1.0, 2.0, 3.0], [2.0, 1.0, 3.0], [3.0, 2.0, 1.0]],
+            coords: vec![[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
+        };
+        let mut f = XTCTrajectory::open(tmp_path, FileMode::Write).unwrap();
+        let write_status = f.write(&frame);
+        match write_status {
+            Err(_) => panic!("Failed"),
+            Ok(()) => {}
+        }
+        f.flush().unwrap();
+
+        let mut new_frame = Frame::with_capacity(natoms);
+        let mut f = XTCTrajectory::open(tmp_path, FileMode::Read).unwrap();
+        let num_atoms = f.get_num_atoms().unwrap();
+        assert_eq!(num_atoms, natoms);
+
+        let read_status = f.read(&mut new_frame);
+        match read_status {
+            Err(e) => assert!(false, "{:?}", e),
+            Ok(()) => {}
+        }
+
+        assert_eq!(new_frame.num_atoms, frame.num_atoms);
+        assert_eq!(new_frame.step, frame.step);
+        assert_approx_eq!(new_frame.time, frame.time);
+        assert_eq!(new_frame.box_vector, frame.box_vector);
+        assert_eq!(new_frame.coords, frame.coords);
+    }
+
+    #[test]
+    fn test_read_write_trr() {
+        let tempfile = NamedTempFile::new().unwrap();
+        let tmp_path = tempfile.path();
+
+        let natoms: u32 = 2;
+        let frame = Frame {
+            num_atoms: natoms,
+            step: 5,
+            time: 2.0,
+            box_vector: [[1.0, 2.0, 3.0], [2.0, 1.0, 3.0], [3.0, 2.0, 1.0]],
+            coords: vec![[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
+        };
+        let mut f = TRRTrajectory::open(tmp_path, FileMode::Write).unwrap();
+        let write_status = f.write(&frame);
+        match write_status {
+            Err(_) => panic!("Failed"),
+            Ok(()) => {}
+        }
+        f.flush().unwrap();
+
+        let mut new_frame = Frame::with_capacity(natoms);
+        let mut f = TRRTrajectory::open(tmp_path, FileMode::Read).unwrap();
+        // let num_atoms = f.get_num_atoms().unwrap();
+        // assert_eq!(num_atoms, natoms);
+
+        let read_status = f.read(&mut new_frame);
+        match read_status {
+            Err(e) => assert!(false, "{:?}", e),
+            Ok(()) => {}
+        }
+
+        assert_eq!(new_frame.num_atoms, frame.num_atoms);
+        assert_eq!(new_frame.step, frame.step);
+        assert_eq!(new_frame.time, frame.time);
+        assert_eq!(new_frame.box_vector, frame.box_vector);
+        assert_eq!(new_frame.coords, frame.coords);
+    }
+}
+