From 230330347233489e4b17f4514a3aa9bf37256d1a Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 16 Apr 2020 13:47:08 +0200 Subject: [PATCH] XTCTrajectory and TRRTrajectory --- Cargo.toml | 1 + src/frame.rs | 103 +++++++++++++++ src/lib.rs | 338 +++++++++++++++++++++++++++++++++++++++++++++++++ src/xdrfile.rs | 0 4 files changed, 442 insertions(+) create mode 100644 src/frame.rs create mode 100644 src/xdrfile.rs diff --git a/Cargo.toml b/Cargo.toml index f1c63cd..c9ded1b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ lazy-init = "0.3" [dev-dependencies] tempfile = "3.1.0" +assert_approx_eq = "1.1.0" [build-dependencies] cc = { version = "1.0", features = ["parallel" ]} diff --git a/src/frame.rs b/src/frame.rs new file mode 100644 index 0000000..82dcf21 --- /dev/null +++ b/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 = 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); + } +} diff --git a/src/lib.rs b/src/lib.rs index 268972c..3b5f62d 100644 --- a/src/lib.rs +++ b/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 { + 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; +} + +pub struct XTCTrajectory { + handle: XDRFile, + precision: Cell, // internal mutability required for read method + num_atoms: Lazy> +} + +impl XTCTrajectory { + pub fn open(path: &Path, filemode: FileMode) -> Result { + 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 { + 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> +} + +impl TRRTrajectory { + pub fn open(path: &Path, filemode: FileMode) -> Result { + 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 { + 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); + } +} + diff --git a/src/xdrfile.rs b/src/xdrfile.rs new file mode 100644 index 0000000..e69de29