use molly::{selection::AtomSelection, Frame};
use crate::{
errors::ReadTrajError,
io::traj_read::TrajGroupReadOpen,
prelude::{
FrameData, FrameDataTime, TrajFile, TrajFullReadOpen, TrajRangeRead, TrajRead,
TrajReadOpen, TrajReader, TrajStepRead, TrajStepTimeRead, Vector3D,
},
structures::{container::AtomContainer, group::Group},
system::System,
};
use std::{
fs::File,
io::{ErrorKind, Read, Seek, SeekFrom},
marker::PhantomData,
path::Path,
};
const TIME_PRECISION: f32 = 0.001;
pub struct MollyXtc {
reader: molly::XTCReader<File>,
group: AtomContainer,
atom_selection: AtomSelection,
}
impl MollyXtc {
fn new(
filename: impl AsRef<Path>,
n_atoms: usize,
group: Option<&Group>,
) -> Result<MollyXtc, ReadTrajError> {
let file = std::fs::File::open(&filename)
.map_err(|_| ReadTrajError::FileNotFound(Box::from(filename.as_ref())))?;
let mut xtc = match group {
None => MollyXtc {
reader: molly::XTCReader { file, step: 0 },
group: AtomContainer::from_ranges(vec![(0, n_atoms)], n_atoms),
atom_selection: AtomSelection::All,
},
Some(group) => {
let number = group.get_atoms().last().map(|x| x + 1).unwrap_or(0);
MollyXtc {
reader: molly::XTCReader { file, step: 0 },
group: group.get_atoms().clone(),
atom_selection: AtomSelection::Until(number as u32),
}
}
};
match xtc.read_i32() {
Ok(x) if x == 1995 || x == 2023 => x,
Ok(_) => return Err(ReadTrajError::NotXtc(Box::from(filename.as_ref()))),
Err(_) => return Err(ReadTrajError::CouldNotReadMagic),
};
match xtc.read_i32() {
Ok(x) if x == n_atoms as i32 => (),
Ok(_) => {
return Err(ReadTrajError::AtomsNumberMismatch(Box::from(
filename.as_ref(),
)))
}
Err(_) => return Err(ReadTrajError::FrameNotFound),
};
xtc.reader
.home()
.expect("FATAL GROAN ERROR | MollyXtc::new | Could not jump to the start of the file.");
Ok(xtc)
}
#[inline]
fn read_n_bytes(&mut self, magic: i32) -> Result<i64, ReadTrajError> {
let bytes = match magic {
1995 => self.read_i32().map_err(|_| ReadTrajError::SkipFailed)? as i64,
2023 => self.read_i64().map_err(|_| ReadTrajError::SkipFailed)?,
_ => panic!("FATAL GROAN ERROR | MollyXtc::read_n_bytes | Unexpected magic number `{}` provided.", magic),
};
Ok(bytes + (4 - (bytes % 4)) % 4)
}
fn skip_frame(&mut self) -> Result<bool, ReadTrajError> {
let magic_number = match self.read_i32() {
Ok(x) => x,
Err(_) => return Ok(false),
};
if magic_number != 1995 && magic_number != 2023 {
return Err(ReadTrajError::FrameNotFound);
}
self.xdr_jump(84).map_err(|_| ReadTrajError::SkipFailed)?;
let size = self.read_n_bytes(magic_number)?;
if self.xdr_jump(size).is_err() {
return Ok(false);
}
Ok(true)
}
fn skip_frame_with_time(&mut self) -> Result<Option<f32>, ReadTrajError> {
let magic_number = match self.read_i32() {
Ok(x) => x,
Err(_) => return Ok(None),
};
if magic_number != 1995 && magic_number != 2023 {
return Err(ReadTrajError::FrameNotFound);
}
self.xdr_jump(8).map_err(|_| ReadTrajError::SkipFailed)?;
let time = self.read_f32().map_err(|_| ReadTrajError::SkipFailed)?;
self.xdr_jump(72).map_err(|_| ReadTrajError::SkipFailed)?;
let size = self.read_n_bytes(magic_number)?;
if self.xdr_jump(size).is_err() {
return Ok(None);
}
Ok(Some(time))
}
fn jump_to_start(&mut self, target_time: f32) -> Result<(), ReadTrajError> {
loop {
match self.get_jump_info(target_time, TIME_PRECISION)? {
0 => break, jump => self.xdr_jump(jump).map_err(|_| ReadTrajError::SkipFailed)?,
}
}
self.xdr_jump(-16).map_err(|_| ReadTrajError::SkipFailed)?;
Ok(())
}
fn get_jump_info(
&mut self,
target_time: f32,
time_precision: f32,
) -> Result<i64, ReadTrajError> {
let magic_number = self
.read_i32()
.map_err(|_| ReadTrajError::StartNotFound(target_time.to_string()))?;
if magic_number != 1995 && magic_number != 2023 {
return Err(ReadTrajError::FrameNotFound);
}
self.xdr_jump(8)
.map_err(|_| ReadTrajError::StartNotFound(target_time.to_string()))?;
let time = self
.read_f32()
.map_err(|_| ReadTrajError::StartNotFound(target_time.to_string()))?;
if time >= target_time - time_precision {
return Ok(0);
}
self.xdr_jump(72)
.map_err(|_| ReadTrajError::StartNotFound(target_time.to_string()))?;
let size = self.read_n_bytes(magic_number)?;
Ok(size)
}
#[inline(always)]
fn read_i32(&mut self) -> std::io::Result<i32> {
let mut buf = [0u8; 4];
self.reader.file.read_exact(&mut buf)?;
Ok(i32::from_be_bytes(buf))
}
#[inline(always)]
fn read_f32(&mut self) -> std::io::Result<f32> {
let mut buf = [0u8; 4];
self.reader.file.read_exact(&mut buf)?;
Ok(f32::from_be_bytes(buf))
}
pub(crate) fn read_i64(&mut self) -> std::io::Result<i64> {
let mut buf: [u8; 8] = Default::default();
self.reader.file.read_exact(&mut buf)?;
Ok(i64::from_be_bytes(buf))
}
#[inline(always)]
fn xdr_jump(&mut self, offset: i64) -> std::io::Result<()> {
self.reader.file.seek(SeekFrom::Current(offset))?;
Ok(())
}
}
pub struct XtcReader<'a> {
system: *mut System,
xtc: MollyXtc,
phantom: PhantomData<&'a mut System>,
}
impl TrajFile for MollyXtc {}
pub struct XtcFrameData {
frame: molly::Frame,
}
impl FrameData for XtcFrameData {
type TrajFile = MollyXtc;
#[inline(always)]
fn from_frame(traj_file: &mut MollyXtc, _system: &System) -> Option<Result<Self, ReadTrajError>>
where
Self: Sized,
{
let mut frame = Frame::default();
match traj_file
.reader
.read_frame_with_selection(&mut frame, &traj_file.atom_selection)
{
Ok(_) => {
let frame_data = XtcFrameData { frame };
Some(Ok(frame_data))
}
Err(e) if e.kind() == ErrorKind::UnexpectedEof => None,
Err(e) => Some(Err(ReadTrajError::MollyXtcError(e.to_string()))),
}
}
#[inline]
fn update_system(self, system: &mut System) {
for (atom, pos) in system.get_atoms_mut().iter_mut().zip(self.frame.coords()) {
atom.set_position(Vector3D::new(pos[0], pos[1], pos[2]));
atom.reset_velocity();
atom.reset_force();
}
system.set_simulation_step(self.frame.step as u64);
system.set_simulation_time(self.frame.time);
let b = self.frame.boxvec;
system.set_box([b[0], b[4], b[8], b[1], b[2], b[3], b[5], b[6], b[7]].into());
system.set_precision(self.frame.precision as u64);
}
}
impl FrameDataTime for XtcFrameData {
#[inline(always)]
fn get_time(&self) -> f32 {
self.frame.time
}
}
impl<'a> TrajRead<'a> for XtcReader<'a> {
type FrameData = XtcFrameData;
#[inline(always)]
fn get_system(&mut self) -> *mut System {
self.system
}
#[inline(always)]
fn get_file_handle(&mut self) -> &mut MollyXtc {
&mut self.xtc
}
}
impl<'a> TrajFullReadOpen<'a> for XtcReader<'a> {
#[inline(always)]
fn new(
system: &'a mut System,
filename: impl AsRef<Path>,
) -> Result<XtcReader<'a>, ReadTrajError> {
let xtc = MollyXtc::new(&filename, system.get_n_atoms(), None)?;
Ok(XtcReader {
system,
xtc,
phantom: PhantomData,
})
}
}
impl<'a> TrajReadOpen<'a> for XtcReader<'a> {
#[inline(always)]
fn initialize(
system: &'a mut System,
filename: impl AsRef<Path>,
group: Option<&str>,
) -> Result<Self, ReadTrajError>
where
Self: Sized,
{
match group {
None => TrajFullReadOpen::new(system, filename),
Some(_) => panic!("FATAL GROAN ERROR | XtcReader::initialize | XtcReader does not support partial-frame reading. Use `GroupXtcReader` instead."),
}
}
}
impl<'a> TrajRangeRead<'a> for XtcReader<'a> {
#[inline(always)]
fn jump_to_start(&mut self, start_time: f32) -> Result<(), ReadTrajError> {
self.xtc.jump_to_start(start_time)
}
}
impl<'a> TrajStepRead<'a> for XtcReader<'a> {
#[inline(always)]
fn skip_frame(&mut self) -> Result<bool, ReadTrajError> {
self.xtc.skip_frame()
}
}
impl<'a> TrajStepTimeRead<'a> for XtcReader<'a> {
#[inline(always)]
fn skip_frame_time(&mut self) -> Result<Option<f32>, ReadTrajError> {
self.xtc.skip_frame_with_time()
}
}
pub struct GroupXtcFrameData {
frame: molly::Frame,
group: *const AtomContainer,
}
impl FrameDataTime for GroupXtcFrameData {
#[inline(always)]
fn get_time(&self) -> f32 {
self.frame.time
}
}
impl FrameData for GroupXtcFrameData {
type TrajFile = MollyXtc;
#[inline(always)]
fn from_frame(traj_file: &mut MollyXtc, _system: &System) -> Option<Result<Self, ReadTrajError>>
where
Self: Sized,
{
let mut frame = Frame::default();
match traj_file
.reader
.read_frame_with_selection_buffered(&mut frame, &traj_file.atom_selection)
{
Ok(_) => {
let frame_data = GroupXtcFrameData {
frame,
group: &traj_file.group as *const AtomContainer,
};
Some(Ok(frame_data))
}
Err(e) if e.kind() == ErrorKind::UnexpectedEof => None,
Err(e) => Some(Err(ReadTrajError::MollyXtcError(e.to_string()))),
}
}
#[inline]
fn update_system(self, system: &mut System) {
let atoms = unsafe { &*self.group }.iter();
for index in atoms {
let atom = system
.get_atom_mut(index)
.expect("FATAL GROAN ERROR | XtcFrameData::update_system | Atom should exist.");
unsafe {
let x = *self.frame.positions.get_unchecked(3 * index);
let y = *self.frame.positions.get_unchecked(3 * index + 1);
let z = *self.frame.positions.get_unchecked(3 * index + 2);
atom.set_position(Vector3D::new(x, y, z));
}
atom.reset_velocity();
atom.reset_force();
}
system.set_simulation_step(self.frame.step as u64);
system.set_simulation_time(self.frame.time);
let b = self.frame.boxvec;
system.set_box([b[0], b[4], b[8], b[1], b[2], b[3], b[5], b[6], b[7]].into());
system.set_precision(self.frame.precision as u64);
}
}
pub struct GroupXtcReader<'a> {
system: *mut System,
xtc: MollyXtc,
phantom: PhantomData<&'a mut System>,
}
impl<'a> TrajRead<'a> for GroupXtcReader<'a> {
type FrameData = GroupXtcFrameData;
#[inline(always)]
fn get_system(&mut self) -> *mut System {
self.system
}
#[inline(always)]
fn get_file_handle(&mut self) -> &mut MollyXtc {
&mut self.xtc
}
}
impl<'a> TrajReadOpen<'a> for GroupXtcReader<'a> {
#[inline(always)]
fn initialize(
system: &'a mut System,
filename: impl AsRef<Path>,
group: Option<&str>,
) -> Result<Self, ReadTrajError>
where
Self: Sized,
{
match group {
None => panic!("FATAL GROAN ERROR | GroupXtcReader::initialize | GroupXtcReader requires a group to be initialized. Use `XtcReader` instead."),
Some(x) => TrajGroupReadOpen::new(system, filename, x),
}
}
}
impl<'a> TrajGroupReadOpen<'a> for GroupXtcReader<'a> {
#[inline(always)]
fn new(
system: &'a mut System,
filename: impl AsRef<Path>,
group: &str,
) -> Result<Self, ReadTrajError>
where
Self: Sized,
{
let group = system
.get_groups()
.get(group)
.map_err(|_| ReadTrajError::GroupNotFound(group.to_owned()))?;
let xtc = MollyXtc::new(&filename, system.get_n_atoms(), Some(group))?;
Ok(GroupXtcReader {
system,
xtc,
phantom: PhantomData,
})
}
}
impl<'a> TrajRangeRead<'a> for GroupXtcReader<'a> {
#[inline(always)]
fn jump_to_start(&mut self, start_time: f32) -> Result<(), ReadTrajError> {
self.xtc.jump_to_start(start_time)
}
}
impl<'a> TrajStepRead<'a> for GroupXtcReader<'a> {
#[inline(always)]
fn skip_frame(&mut self) -> Result<bool, ReadTrajError> {
self.xtc.skip_frame()
}
}
impl<'a> TrajStepTimeRead<'a> for GroupXtcReader<'a> {
#[inline(always)]
fn skip_frame_time(&mut self) -> Result<Option<f32>, ReadTrajError> {
self.xtc.skip_frame_with_time()
}
}
impl System {
pub fn group_xtc_iter(
&mut self,
filename: impl AsRef<Path>,
group: &str,
) -> Result<TrajReader<'_, GroupXtcReader<'_>>, ReadTrajError> {
Ok(TrajReader::wrap_traj(GroupXtcReader::new(
self, filename, group,
)?))
}
}
#[cfg(test)]
mod tests_group {
use crate::{
errors::ReadTrajError,
system::System,
test_utilities::utilities::{compare_atoms, compare_box},
};
#[test]
fn group_xtc_iter_all() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "all")
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for (a1, a2) in g.atoms_iter().zip(f.atoms_iter()) {
compare_atoms(a1, a2);
}
}
}
#[test]
fn group_xtc_iter_popc() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
group_system
.group_create("Membrane", "resname POPC")
.unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
let original_system = full_system.clone();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "Membrane")
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for ((a1, a2), a3) in g
.atoms_iter()
.zip(f.atoms_iter())
.zip(original_system.atoms_iter())
{
if g.group_isin("Membrane", a1.get_index()).unwrap() {
compare_atoms(a1, a2);
} else {
compare_atoms(a1, a3);
}
}
}
}
#[test]
fn group_xtc_iter_protein() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
group_system.group_create("Protein", "@protein").unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
let original_system = full_system.clone();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "Protein")
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for ((a1, a2), a3) in g
.atoms_iter()
.zip(f.atoms_iter())
.zip(original_system.atoms_iter())
{
if g.group_isin("Protein", a1.get_index()).unwrap() {
compare_atoms(a1, a2);
} else {
compare_atoms(a1, a3);
}
}
}
}
#[test]
fn group_xtc_iter_single_particle() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
group_system.group_create("Particle", "serial 1").unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
let original_system = full_system.clone();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "Particle")
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for ((a1, a2), a3) in g
.atoms_iter()
.zip(f.atoms_iter())
.zip(original_system.atoms_iter())
{
if g.group_isin("Particle", a1.get_index()).unwrap() {
compare_atoms(a1, a2);
} else {
compare_atoms(a1, a3);
}
}
}
}
#[test]
fn group_xtc_iter_no_particle() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
group_system.group_create("Nothing", "not all").unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
let original_system = full_system.clone();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "Nothing")
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for (a1, a3) in g.atoms_iter().zip(original_system.atoms_iter()) {
compare_atoms(a1, a3);
}
}
}
#[test]
fn group_xtc_iter_range() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
group_system
.group_create("Membrane", "resname POPC")
.unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
let original_system = full_system.clone();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "Membrane")
.unwrap()
.with_range(300.0, 800.0)
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap()
.with_range(300.0, 800.0)
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for ((a1, a2), a3) in g
.atoms_iter()
.zip(f.atoms_iter())
.zip(original_system.atoms_iter())
{
if g.group_isin("Membrane", a1.get_index()).unwrap() {
compare_atoms(a1, a2);
} else {
compare_atoms(a1, a3);
}
}
}
}
#[test]
fn group_xtc_range_start_not_found() {
let mut system = System::from_file("test_files/example.gro").unwrap();
system.group_create("Membrane", "resname POPC").unwrap();
match system
.group_xtc_iter("test_files/short_trajectory.xtc", "Membrane")
.unwrap()
.with_range(12000.0, 20000.0)
{
Ok(_) => panic!("Iterator should not have been constructed."),
Err(ReadTrajError::StartNotFound(_)) => (),
Err(e) => panic!("Incorrect error type {} returned", e),
}
}
#[test]
fn group_xtc_step_3() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
group_system
.group_create("Membrane", "resname POPC")
.unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
let original_system = full_system.clone();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "Membrane")
.unwrap()
.with_step(3)
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap()
.with_step(3)
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for ((a1, a2), a3) in g
.atoms_iter()
.zip(f.atoms_iter())
.zip(original_system.atoms_iter())
{
if g.group_isin("Membrane", a1.get_index()).unwrap() {
compare_atoms(a1, a2);
} else {
compare_atoms(a1, a3);
}
}
}
}
#[test]
fn group_xtc_iter_range_and_step() {
let mut group_system = System::from_file("test_files/example.gro").unwrap();
group_system
.group_create("Membrane", "resname POPC")
.unwrap();
let mut full_system = System::from_file("test_files/example.gro").unwrap();
let original_system = full_system.clone();
for (group_frame, full_frame) in group_system
.group_xtc_iter("test_files/short_trajectory.xtc", "Membrane")
.unwrap()
.with_range(300.0, 800.0)
.unwrap()
.with_step(2)
.unwrap()
.zip(
full_system
.xtc_iter("test_files/short_trajectory.xtc")
.unwrap()
.with_range(300.0, 800.0)
.unwrap()
.with_step(2)
.unwrap(),
)
{
let g = group_frame.unwrap();
let f = full_frame.unwrap();
compare_box(g.get_box().unwrap(), f.get_box().unwrap());
for ((a1, a2), a3) in g
.atoms_iter()
.zip(f.atoms_iter())
.zip(original_system.atoms_iter())
{
if g.group_isin("Membrane", a1.get_index()).unwrap() {
compare_atoms(a1, a2);
} else {
compare_atoms(a1, a3);
}
}
}
}
#[test]
fn group_xtc_nonexistent_file() {
let mut system = System::from_file("test_files/example.gro").unwrap();
system.group_create("Membrane", "resname POPC").unwrap();
match system.group_xtc_iter("test_files/nonexistent.xtc", "Membrane") {
Err(ReadTrajError::FileNotFound(_)) => (),
_ => panic!("XTC file should not exist."),
}
}
#[test]
fn group_xtc_nonexistent_group() {
let mut system = System::from_file("test_files/example.gro").unwrap();
match system.group_xtc_iter("test_files/short_trajectory.xtc", "Membrane") {
Err(ReadTrajError::GroupNotFound(x)) => assert_eq!(x, "Membrane"),
_ => panic!("XTC file should not exist."),
}
}
}