use crate::{error::CError, frame::Frame, property::Property, unit_cell::UnitCell};
use core::f64;
use log::{debug, warn};
use netcdf3::{DataSet, DataType, FileReader, FileWriter, Variable, Version};
use std::{
fs::OpenOptions,
path::{Path, PathBuf},
};
#[derive(Debug)]
struct VariableWithScale {
var: Variable,
scale: f64,
}
#[derive(Debug, Default)]
struct Variables {
coordinates: Option<VariableWithScale>,
velocities: Option<VariableWithScale>,
cell_lengths: Option<VariableWithScale>,
cell_angles: Option<VariableWithScale>,
time: Option<VariableWithScale>,
}
#[derive(Debug)]
struct BufferedFrame {
positions: Vec<f32>,
velocities: Option<Vec<f32>>,
cell_lengths: [f32; 3],
cell_angles: [f32; 3],
time: Option<f32>,
}
#[derive(Debug, Default)]
pub struct AMBERTrajFormat {
reader: Option<FileReader>,
variables: Variables,
index: usize,
file_title: String,
n_atoms: usize,
convention: Convention,
write_path: Option<PathBuf>,
buffered_frames: Vec<BufferedFrame>,
has_velocities: bool,
initialized: bool,
mode: FileMode,
version: Option<Version>,
}
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
pub(crate) enum FileMode {
#[default]
Read,
Append,
}
#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
pub(crate) enum Convention {
#[default]
Amber,
Restart,
}
impl Convention {
fn as_str(self) -> &'static str {
match self {
Convention::Amber => "AMBER",
Convention::Restart => "AMBERRESTART",
}
}
}
enum NumericData {
F32(Vec<f32>),
F64(Vec<f64>),
}
fn read_values(
file_reader: &mut FileReader,
index: usize,
convention: Convention,
variable: &Variable,
invalid_type_message: &str,
) -> Result<NumericData, CError> {
let name = variable.name();
match variable.data_type() {
DataType::I8 | DataType::U8 | DataType::I16 | DataType::I32 => {
Err(CError::GenericError(invalid_type_message.to_string()))
}
DataType::F32 => {
let values = match convention {
Convention::Amber => file_reader.read_record_f32(name, index)?,
Convention::Restart => file_reader.read_var_f32(name)?,
};
Ok(NumericData::F32(values))
}
DataType::F64 => {
let values = match convention {
Convention::Amber => file_reader.read_record_f64(name, index)?,
Convention::Restart => file_reader.read_var_f64(name)?,
};
Ok(NumericData::F64(values))
}
}
}
fn read_triplet(
file_reader: &mut FileReader,
index: usize,
convention: Convention,
variable: &VariableWithScale,
) -> Result<[f64; 3], CError> {
match read_values(
file_reader,
index,
convention,
&variable.var,
"invalid type for variable, expected f32/f64 data",
)? {
NumericData::F32(values) => Ok([
variable.scale * f64::from(values[0]),
variable.scale * f64::from(values[1]),
variable.scale * f64::from(values[2]),
]),
NumericData::F64(values) => Ok([
variable.scale * values[0],
variable.scale * values[1],
variable.scale * values[2],
]),
}
}
fn read_scalar(
file_reader: &mut FileReader,
index: usize,
convention: Convention,
variable: &VariableWithScale,
invalid_type_message: &str,
) -> Result<f64, CError> {
match read_values(
file_reader,
index,
convention,
&variable.var,
invalid_type_message,
)? {
NumericData::F32(values) => Ok(variable.scale * f64::from(values[0])),
NumericData::F64(values) => Ok(variable.scale * values[0]),
}
}
fn read_array(
file_reader: &mut FileReader,
index: usize,
convention: Convention,
variable: &VariableWithScale,
array: &mut [[f64; 3]],
) -> Result<(), CError> {
match read_values(
file_reader,
index,
convention,
&variable.var,
"invalid type for variable, expected f32/f64 data",
)? {
NumericData::F32(values) => {
for (chunk, item) in values.chunks_exact(3).zip(array.iter_mut()) {
item[0] = variable.scale * f64::from(chunk[0]);
item[1] = variable.scale * f64::from(chunk[1]);
item[2] = variable.scale * f64::from(chunk[2]);
}
}
NumericData::F64(values) => {
for (chunk, item) in values.chunks_exact(3).zip(array.iter_mut()) {
item[0] = variable.scale * chunk[0];
item[1] = variable.scale * chunk[1];
item[2] = variable.scale * chunk[2];
}
}
}
Ok(())
}
fn validate_common(file_reader: &FileReader, convention: &str) -> Result<(), CError> {
let data_set = file_reader.data_set();
fn check_attr(data_set: &DataSet, attr: &str, expected_attr: &str) -> Result<(), CError> {
let read_attr = data_set.get_global_attr(attr).ok_or(CError::GenericError(
"expected a 'Conventions' attribute to be defined".to_string(),
))?;
let read_attr = read_attr.get_as_string();
if let Some(read_attr) = read_attr {
if read_attr != expected_attr {
return Err(CError::GenericError(format!(
"expected '{expected_attr}', got {read_attr}"
)));
}
} else {
return Err(CError::GenericError(
"could not read attr as string".to_string(),
));
}
Ok(())
}
check_attr(data_set, "Conventions", convention)?;
check_attr(data_set, "ConventionVersion", "1.0")?;
if let Some(spatial) = file_reader.data_set().get_dim("spatial") {
if spatial.size() != 3 {
return Err(CError::GenericError(format!(
"'spatial' dimension must have a size of 3, got {}",
spatial.size()
)));
}
} else {
return Err(CError::GenericError(
"missing 'spatial' dimension".to_string(),
));
}
file_reader
.data_set()
.get_dim("atom")
.ok_or(CError::GenericError("missing 'atom' dimension".to_string()))?;
if let Some(cell_spatial) = file_reader.data_set().get_dim("cell_spatial")
&& cell_spatial.size() != 3
{
return Err(CError::GenericError(format!(
"'cell_spatial' dimension must have a size of 3, got {}",
cell_spatial.size()
)));
}
if let Some(cell_angular) = file_reader.data_set().get_dim("cell_angular")
&& cell_angular.size() != 3
{
return Err(CError::GenericError(format!(
"'cell_angular' dimension must have a size of 3, got {}",
cell_angular.size()
)));
}
Ok(())
}
fn scale_for_distance(units: &str) -> f64 {
match units.to_lowercase().as_str() {
"angstroms" | "angstrom" | "a" => 1.0,
"meters" | "meter" | "m" => 1e10,
"centimeters" | "centimeter" | "cm" => 1e8,
"micrometers" | "micrometer" | "µm" | "um" => 1e4,
"nanometers" | "nanometer" | "nm" => 1e4,
"bohrs" | "bohr" => 0.52918,
unknown => {
warn!("unknown unit ({unknown}) for distances");
1.0
}
}
}
fn scale_for_velocity(units: &str) -> f64 {
let lower_case = units.to_lowercase();
let mut splitted = lower_case.split('/');
if let Some(scale) = splitted.next()
&& let Some(time_unit) = splitted.next().as_mut()
{
let scale = scale_for_distance(scale);
let time_unit = scale_for_time(time_unit);
return scale / time_unit;
}
warn!("unknown unit ({units}) for velocities");
1.0
}
fn scale_for_time(units: &str) -> f64 {
match units.to_lowercase().as_str() {
"picoseconds" | "picosecond" | "ps" => 1.0,
"femtoseconds" | "femtosecond" | "fs" => 1e-3,
"nanoseconds" | "nanosecond" | "ns" => 1e3,
"microseconds" | "microsecond" | "µs" | "us" => 1e6,
"seconds" | "second" | "s" => 1e12,
unknown => {
warn!("unknown unit ({unknown}) for time");
1.0
}
}
}
fn scale_factor(variable: &Variable) -> Result<f64, CError> {
if let Some(scale) = variable
.get_attr_f32("scale_factor")
.and_then(|scale| scale.first().copied())
{
return Ok(f64::from(scale));
}
if let Some(scale) = variable
.get_attr_f64("scale_factor")
.and_then(|scale| scale.first().copied())
{
return Ok(scale);
}
if variable.get_attr("scale_factor").is_some() {
return Err(CError::GenericError(format!(
"scale_factor attribute for '{}' must be a floating point value",
variable.name()
)));
}
Ok(1.0)
}
impl AMBERTrajFormat {
fn read_cell(&mut self) -> Result<Option<UnitCell>, CError> {
if self.variables.cell_lengths.is_none() || self.variables.cell_angles.is_none() {
return Ok(None);
}
let convention = self.convention;
let index = self.index;
let reader = self
.reader
.as_mut()
.expect("reader should've been initialized");
let cell_lengths = self
.variables
.cell_lengths
.as_ref()
.expect("we just checked for None");
let lengths = read_triplet(reader, index, convention, cell_lengths)?;
let cell_angles = self
.variables
.cell_angles
.as_ref()
.expect("we just checked for None");
let mut angles = read_triplet(reader, index, convention, cell_angles)?;
Ok(Some(UnitCell::new_from_lengths_angles(
lengths,
&mut angles,
)?))
}
pub(crate) fn open(
path: &Path,
mode: FileMode,
convention: Convention,
) -> Result<Self, CError> {
let file_reader = match FileReader::open(path) {
Ok(r) => r,
Err(_) if mode == FileMode::Append => {
return Ok(Self {
write_path: Some(path.to_path_buf()),
mode,
convention,
..Default::default()
});
}
Err(e) => return Err(CError::GenericError(e.to_string())),
};
validate_common(&file_reader, convention.as_str())?;
let file_title = file_reader
.data_set()
.get_global_attr_as_string("title")
.unwrap_or_default();
let n_atoms = file_reader
.data_set()
.get_dim("atom")
.expect("we just just validated that 'atom' exists")
.size();
let coordinates = if let Some(coordinates) = file_reader.data_set().get_var("coordinates") {
let mut coords = Some(VariableWithScale {
var: coordinates.clone(),
scale: scale_factor(coordinates)?,
});
if let Some(units) = coordinates.get_attr_as_string("units") {
coords.as_mut().expect("we just init'ed").scale *=
scale_for_distance(units.as_str());
} else {
debug!("not scaling coordinates");
}
coords
} else {
warn!("the coordinates variable is not defined in this file.");
None
};
let velocities = if let Some(velocities) = file_reader.data_set().get_var("velocities") {
let mut velo = Some(VariableWithScale {
var: velocities.clone(),
scale: scale_factor(velocities)?,
});
if let Some(units) = velocities.get_attr_as_string("units") {
velo.as_mut().expect("we just init'ed").scale *= scale_for_velocity(units.as_str());
} else {
debug!("not scaling velocities");
}
velo
} else {
warn!("the velocities variable is not defined in this file.");
None
};
let read_cell_lengths = file_reader.data_set().get_var("cell_lengths");
let read_cell_angles = file_reader.data_set().get_var("cell_angles");
let mut cell_lengths = None;
let mut cell_angles = None;
if let Some(read_cell_angles) = read_cell_angles
&& let Some(read_cell_lengths) = read_cell_lengths
{
cell_lengths = Some(VariableWithScale {
var: read_cell_lengths.clone(),
scale: scale_factor(read_cell_lengths)?,
});
if let Some(units) = read_cell_lengths.get_attr_as_string("units") {
cell_lengths.as_mut().expect("we just init'ed it").scale *=
scale_for_distance(units.as_str());
} else {
debug!("not scaling cell lengths");
}
cell_angles = Some(VariableWithScale {
var: read_cell_angles.clone(),
scale: scale_factor(read_cell_angles)?,
});
if let Some(units) = read_cell_angles.get_attr_as_string("units") {
let scaling_factor = match units.to_lowercase().as_str() {
"" | "degrees" | "degree" => 1.0,
"radians" | "radian" => 180.0 / f64::consts::PI,
unknown => {
warn!("unknown unit ({unknown}) for angles");
1.0
}
};
cell_angles.as_mut().expect("we just init'ed it").scale *= scaling_factor;
} else {
debug!("not scaling cell angles");
}
} else if read_cell_lengths.is_some() {
if read_cell_angles.is_none() {
return Err(CError::GenericError(
"cell_lengths requires cell_angles to be defined".to_string(),
));
}
if coordinates.is_none() {
return Err(CError::GenericError(
"cell_lengths requires coordinates to be defined.".to_string(),
));
}
}
let time = if let Some(time) = file_reader.data_set().get_var("time") {
let mut t = Some(VariableWithScale {
var: time.clone(),
scale: scale_factor(time)?,
});
if let Some(units) = time.get_attr_as_string("units") {
t.as_mut().expect("we just init'ed").scale *= scale_for_time(units.as_str());
} else {
debug!("not scaling time");
}
t
} else {
warn!("the timevariable is not defined in this file.");
None
};
let variables = Variables {
coordinates,
velocities,
cell_lengths,
cell_angles,
time,
};
let index = match convention {
Convention::Amber if mode == FileMode::Append => file_reader
.data_set()
.num_records()
.expect("data_set to have records"),
Convention::Restart | Convention::Amber => 0,
};
let version = Some(file_reader.version());
let has_velocities = variables.velocities.is_some();
let write_path = if mode == FileMode::Append {
Some(path.to_path_buf())
} else {
None
};
let format = AMBERTrajFormat {
reader: Some(file_reader),
index,
file_title,
variables,
n_atoms,
mode,
convention,
version,
has_velocities,
write_path,
..Default::default()
};
Ok(format)
}
pub fn create(path: &Path) -> Result<Self, CError> {
Ok(Self {
write_path: Some(path.to_path_buf()),
..Default::default()
})
}
pub fn read(&mut self) -> Result<Frame, CError> {
let frame = self.read_at(self.index);
self.index += 1;
frame
}
pub fn read_at(&mut self, index: usize) -> Result<Frame, CError> {
self.index = index;
let mut frame = Frame::new();
if let Some(unitcell) = self.read_cell()? {
frame.set_unitcell(unitcell);
}
if !self.file_title.is_empty() {
frame.properties.insert(
"name".to_string(),
crate::property::Property::String(self.file_title.clone()),
);
}
frame.resize(self.n_atoms)?;
let convention = self.convention;
let index = self.index;
let reader = self
.reader
.as_mut()
.expect("reader should've been initialized");
if let Some(coordinates) = self.variables.coordinates.as_ref() {
read_array(
reader,
index,
convention,
coordinates,
frame.positions_mut(),
)?;
}
if let Some(velocities) = self.variables.velocities.as_ref() {
frame.add_velocities();
read_array(
reader,
index,
convention,
velocities,
frame.velocities_mut().expect("just resized velocities"),
)?;
}
if let Some(time) = self.variables.time.as_ref() {
let time_value = read_scalar(
reader,
index,
convention,
time,
"invalid type for time variable",
)?;
frame
.properties
.insert("time".to_string(), Property::Double(time_value));
}
Ok(frame)
}
pub fn len(&self) -> Result<usize, CError> {
match self.convention {
Convention::Restart => Ok(1),
Convention::Amber => self
.reader
.as_ref()
.expect("we should have init reader")
.data_set()
.num_records()
.ok_or(CError::GenericError(
"expected unlimited 'frame' dimension to be defined".to_string(),
)),
}
}
pub fn is_empty(&self) -> Result<bool, CError> {
self.len().map(|n| n == 0)
}
pub fn write(&mut self, frame: &Frame) -> Result<(), CError> {
if !self.initialized {
self.n_atoms = frame.size();
self.has_velocities = frame.velocities().is_some();
if let Some(Property::String(name)) = frame.properties.get("name") {
self.file_title = name.clone();
}
self.initialized = true;
}
if frame.size() != self.n_atoms {
return Err(CError::GenericError(format!(
"this file can only write frames with {} atoms, but the frame contains {} atoms",
self.n_atoms,
frame.size()
)));
}
let mut positions = Vec::with_capacity(3 * self.n_atoms);
for pos in frame.positions() {
positions.push(pos[0] as f32);
positions.push(pos[1] as f32);
positions.push(pos[2] as f32);
}
let velocities = if self.has_velocities {
frame.velocities().map(|vels| {
let mut v = Vec::with_capacity(3 * self.n_atoms);
for vel in vels {
v.push(vel[0] as f32);
v.push(vel[1] as f32);
v.push(vel[2] as f32);
}
v
})
} else {
None
};
let cell = frame.cell();
let lengths = cell.lengths();
let angles = cell.angles();
let time = frame
.properties
.get("time")
.and_then(super::super::property::Property::as_double)
.map(|t| t as f32);
self.buffered_frames.push(BufferedFrame {
positions,
velocities,
cell_lengths: [lengths[0] as f32, lengths[1] as f32, lengths[2] as f32],
cell_angles: [angles[0] as f32, angles[1] as f32, angles[2] as f32],
time,
});
self.index += 1;
Ok(())
}
pub fn finish(&mut self) -> Result<(), CError> {
let path = match self.write_path.as_ref() {
Some(p) => p.clone(),
None => return Ok(()),
};
if self.buffered_frames.is_empty() {
return Ok(());
}
let existing_count = self.index - self.buffered_frames.len();
if self.mode == FileMode::Append && existing_count > 0 {
self.finish_append(&path)
} else {
self.finish_create(&path)
}
}
fn finish_create(&mut self, path: &Path) -> Result<(), CError> {
let has_time = self.buffered_frames.iter().any(|f| f.time.is_some());
let total_frames = self.buffered_frames.len();
let data_set = self.build_data_set(total_frames, has_time, self.has_velocities)?;
let mut writer = FileWriter::open(path)?;
writer.set_def(&data_set, Version::Offset64Bit, 0)?;
writer.write_var_u8("spatial", b"xyz")?;
writer.write_var_u8("cell_spatial", b"abc")?;
writer.write_var_u8("cell_angular", b"alphabeta gamma")?;
for (i, frame) in self.buffered_frames.iter().enumerate() {
writer.write_record_f32("coordinates", i, &frame.positions)?;
writer.write_record_f32("cell_lengths", i, &frame.cell_lengths)?;
writer.write_record_f32("cell_angles", i, &frame.cell_angles)?;
if let Some(ref velocities) = frame.velocities {
writer.write_record_f32("velocities", i, velocities)?;
}
if let Some(time) = frame.time {
writer.write_record_f32("time", i, &[time])?;
}
}
writer.close()?;
self.buffered_frames.clear();
Ok(())
}
fn finish_append(&mut self, path: &Path) -> Result<(), CError> {
let new_count = self.buffered_frames.len();
let existing_count = self.index - new_count;
let total_count = self.index;
let has_velocities = self.variables.velocities.is_some();
let has_time = self.variables.time.is_some();
let version = self.version.clone().unwrap_or(Version::Offset64Bit);
let data_set = self.build_data_set(total_count, has_time, has_velocities)?;
let output_file = OpenOptions::new().write(true).open(path)?;
let mut writer = FileWriter::open_seek_write(
path.as_os_str().to_str().expect("valid unicode path"),
Box::new(output_file),
)?;
writer.set_def(&data_set, version, 0)?;
for (i, frame) in self.buffered_frames.iter().enumerate() {
let record_index = existing_count + i;
writer.write_record_f32("coordinates", record_index, &frame.positions)?;
writer.write_record_f32("cell_lengths", record_index, &frame.cell_lengths)?;
writer.write_record_f32("cell_angles", record_index, &frame.cell_angles)?;
if let Some(ref velocities) = frame.velocities {
writer.write_record_f32("velocities", record_index, velocities)?;
}
if let Some(time) = frame.time {
writer.write_record_f32("time", record_index, &[time])?;
}
}
self.buffered_frames.clear();
Ok(())
}
fn build_data_set(
&self,
total_frames: usize,
has_time: bool,
has_velocities: bool,
) -> Result<DataSet, CError> {
let mut data_set = DataSet::new();
data_set.add_global_attr_string("Conventions", "AMBER")?;
data_set.add_global_attr_string("ConventionVersion", "1.0")?;
data_set.add_global_attr_string("program", "molio")?;
data_set.add_global_attr_string("programVersion", env!("CARGO_PKG_VERSION"))?;
if !self.file_title.is_empty() {
data_set.add_global_attr_string("title", &self.file_title)?;
}
data_set.set_unlimited_dim("frame", total_frames)?;
data_set.add_fixed_dim("spatial", 3)?;
data_set.add_fixed_dim("atom", self.n_atoms)?;
data_set.add_fixed_dim("cell_spatial", 3)?;
data_set.add_fixed_dim("cell_angular", 3)?;
data_set.add_fixed_dim("label", 5)?;
data_set.add_var_u8("spatial", &["spatial"])?;
data_set.add_var_u8("cell_spatial", &["cell_spatial"])?;
data_set.add_var_u8("cell_angular", &["cell_angular", "label"])?;
if has_time {
data_set.add_var_f32("time", &["frame"])?;
data_set.add_var_attr_string("time", "units", "picosecond")?;
}
data_set.add_var_f32("coordinates", &["frame", "atom", "spatial"])?;
data_set.add_var_attr_string("coordinates", "units", "angstrom")?;
data_set.add_var_f32("cell_lengths", &["frame", "cell_spatial"])?;
data_set.add_var_attr_string("cell_lengths", "units", "angstrom")?;
data_set.add_var_f32("cell_angles", &["frame", "cell_angular"])?;
data_set.add_var_attr_string("cell_angles", "units", "degree")?;
if has_velocities {
data_set.add_var_f32("velocities", &["frame", "atom", "spatial"])?;
data_set.add_var_attr_string("velocities", "units", "angstrom/picosecond")?;
}
Ok(data_set)
}
}
#[cfg(test)]
mod tests {
use std::path::Path;
use assert_approx_eq::assert_approx_eq;
use tempfile::Builder;
use crate::{
frame::Frame,
trajectory::Trajectory,
unit_cell::{CellShape, UnitCell},
};
#[test]
fn read_files_in_netcdf_format() {
let path = Path::new("./src/tests-data/netcdf/water.nc");
let mut trajectory = Trajectory::open(path).unwrap();
let frame = trajectory.read().unwrap().unwrap();
assert_eq!(frame.size(), 297);
assert_eq!(frame.properties.get("name"), None);
let positions = frame.positions();
assert_approx_eq!(positions[0][0], 0.4172_191, 1e-4);
assert_approx_eq!(positions[0][1], 8.303_366, 1e-4);
assert_approx_eq!(positions[0][2], 11.73_717, 1e-4);
assert_approx_eq!(positions[296][0], 6.664_049, 1e-4);
assert_approx_eq!(positions[296][1], 11.61_418, 1e-4);
assert_approx_eq!(positions[296][2], 12.96_149, 1e-4);
assert_approx_eq!(
frame.properties.get("time").unwrap().as_double().unwrap(),
2.02
);
}
#[test]
fn read_more_than_one_frame() {
let path = Path::new("./src/tests-data/netcdf/water.nc");
let mut trajectory = Trajectory::open(path).unwrap();
trajectory.read().unwrap().unwrap();
trajectory.read().unwrap().unwrap();
let mut frame = trajectory.read().unwrap().unwrap();
assert_eq!(frame.size(), 297);
assert_eq!(frame.properties.get("name"), None);
let positions = frame.positions();
assert_approx_eq!(positions[0][0], 0.2990952, 1e-4);
assert_approx_eq!(positions[0][1], 8.31003, 1e-4);
assert_approx_eq!(positions[0][2], 11.72146, 1e-4);
assert_approx_eq!(positions[296][0], 6.797599, 1e-4);
assert_approx_eq!(positions[296][1], 11.50882, 1e-4);
assert_approx_eq!(positions[296][2], 12.70423, 1e-4);
assert_approx_eq!(
frame.properties.get("time").unwrap().as_double().unwrap(),
2.04
);
while let Some(next_frame) = trajectory.read().unwrap() {
frame = next_frame;
}
let positions = frame.positions();
assert_approx_eq!(positions[0][0], 0.3185586, 1e-4);
assert_approx_eq!(positions[0][1], 8.776042, 1e-4);
assert_approx_eq!(positions[0][2], 11.8927, 1e-4);
assert_approx_eq!(positions[296][0], 7.089802, 1e-4);
assert_approx_eq!(positions[296][1], 10.35007, 1e-4);
assert_approx_eq!(positions[296][2], 12.8159, 1e-4);
assert_approx_eq!(
frame.properties.get("time").unwrap().as_double().unwrap(),
3.01
);
}
#[test]
fn no_cell() {
let path = Path::new("./src/tests-data/netcdf/no-cell.nc");
let mut trajectory = Trajectory::open(path).unwrap();
let frame = trajectory.read().unwrap().unwrap();
assert_eq!(frame.size(), 1989);
assert_eq!(
frame.properties.get("name").unwrap().as_string().unwrap(),
"Cpptraj Generated trajectory"
);
assert_eq!(*frame.cell(), UnitCell::new());
}
#[test]
fn scale_factor() {
let path = Path::new("./src/tests-data/netcdf/scaled_traj.nc");
let mut trajectory = Trajectory::open(path).unwrap();
assert_eq!(trajectory.size, 26);
let frame = trajectory.read_at(12).unwrap().unwrap();
assert_eq!(frame.size(), 1938);
assert_eq!(frame.properties.get("name"), None);
let cell = frame.cell();
assert_eq!(cell.shape, CellShape::Orthorhombic);
assert_approx_eq!(cell.lengths()[0], 1.765 * 60.9682, 1e-4);
assert_approx_eq!(cell.lengths()[1], 1.765 * 60.9682, 1e-4);
assert_approx_eq!(cell.lengths()[2], 1.765 * 0.0, 1e-4);
let positions = frame.positions();
assert_approx_eq!(positions[0][0], 1.39 * 0.455, 1e-4);
assert_approx_eq!(positions[0][1], 1.39 * 0.455, 1e-4);
assert_approx_eq!(positions[0][2], 0.0 * 0.455, 1e-4);
assert_approx_eq!(positions[296][0], 29.1 * 0.455, 1e-4);
assert_approx_eq!(positions[296][1], 37.41 * 0.455, 1e-4);
assert_approx_eq!(positions[296][2], 0.0 * 0.455, 1e-4);
let velocities = frame.velocities().unwrap();
assert_approx_eq!(velocities[1400][0], 0.6854_072 * -0.856, 1e-4);
assert_approx_eq!(velocities[1400][1], 0.09196_011 * -0.856, 1e-4);
assert_approx_eq!(velocities[1400][2], 2.260_214 * -0.856, 1e-4);
assert_approx_eq!(velocities[1600][0], -0.3342_645 * -0.856, 1e-4);
assert_approx_eq!(velocities[1600][1], 0.322_594 * -0.856, 1e-4);
assert_approx_eq!(velocities[1600][2], -2.446_901 * -0.856, 1e-4);
}
fn check_frame(frame: &Frame) {
let name = frame.properties.get("name").unwrap().as_string().unwrap();
assert_eq!(name, "Test Title 123");
let positions = frame.positions();
assert_approx_eq!(positions[0][0], 0.0, 1e-6);
assert_approx_eq!(positions[0][1], 0.0, 1e-6);
assert_approx_eq!(positions[0][2], 0.0, 1e-6);
assert_approx_eq!(positions[1][0], 1.0, 1e-6);
assert_approx_eq!(positions[1][1], 2.0, 1e-6);
assert_approx_eq!(positions[1][2], 3.0, 1e-6);
assert_approx_eq!(positions[2][0], 2.0, 1e-6);
assert_approx_eq!(positions[2][1], 4.0, 1e-6);
assert_approx_eq!(positions[2][2], 6.0, 1e-6);
assert_approx_eq!(positions[3][0], 3.0, 1e-6);
assert_approx_eq!(positions[3][1], 6.0, 1e-6);
assert_approx_eq!(positions[3][2], 9.0, 1e-6);
let velocities = frame.velocities().unwrap();
assert_approx_eq!(velocities[0][0], -3.0, 1e-6);
assert_approx_eq!(velocities[0][1], -2.0, 1e-6);
assert_approx_eq!(velocities[0][2], -1.0, 1e-6);
assert_approx_eq!(velocities[1][0], -3.0, 1e-6);
assert_approx_eq!(velocities[1][1], -2.0, 1e-6);
assert_approx_eq!(velocities[1][2], -1.0, 1e-6);
assert_approx_eq!(velocities[2][0], -3.0, 1e-6);
assert_approx_eq!(velocities[2][1], -2.0, 1e-6);
assert_approx_eq!(velocities[2][2], -1.0, 1e-6);
assert_approx_eq!(velocities[3][0], -3.0, 1e-6);
assert_approx_eq!(velocities[3][1], -2.0, 1e-6);
assert_approx_eq!(velocities[3][2], -1.0, 1e-6);
let cell = frame.cell();
let lengths = cell.lengths();
let angles = cell.angles();
assert_approx_eq!(lengths[0], 2.0, 1e-6);
assert_approx_eq!(lengths[1], 3.0, 1e-6);
assert_approx_eq!(lengths[2], 4.0, 1e-6);
assert_approx_eq!(angles[0], 80.0, 1e-6);
assert_approx_eq!(angles[1], 90.0, 1e-6);
assert_approx_eq!(angles[2], 120.0, 1e-6);
}
fn generate_frame() -> Frame {
let mut frame = Frame::from_unitcell(
UnitCell::new_from_lengths_angles([2.0, 3.0, 4.0], &mut [80.0, 90.0, 120.0]).unwrap(),
);
frame.properties.insert(
"name".to_string(),
crate::property::Property::String("Test Title 123".to_string()),
);
frame
.properties
.insert("time".to_string(), crate::property::Property::Double(2.0));
frame.add_velocities();
for i in 0..4 {
frame.add_atom_with_velocity(
crate::atom::Atom::new("X".to_string()),
[1.0 * f64::from(i), 2.0 * f64::from(i), 3.0 * f64::from(i)],
[-3.0, -2.0, -1.0],
);
}
frame
}
#[test]
fn write_files_in_netcdf_format() {
let frame = generate_frame();
let named_tmpfile = Builder::new()
.prefix("netcdf-test-write")
.suffix(".nc")
.tempfile()
.unwrap();
let mut trajectory = Trajectory::create(named_tmpfile.path()).unwrap();
trajectory.write(&frame).unwrap();
trajectory.write(&frame).unwrap();
drop(trajectory);
let mut trajectory = Trajectory::open(named_tmpfile.path()).unwrap();
check_frame(&trajectory.read().unwrap().unwrap());
check_frame(&trajectory.read().unwrap().unwrap());
}
#[test]
fn append_to_existing_file() {
let frame = generate_frame();
let named_tmpfile = Builder::new()
.prefix("netcdf-test-append")
.suffix(".nc")
.tempfile()
.unwrap();
{
let mut trajectory = Trajectory::create(named_tmpfile.path()).unwrap();
trajectory.write(&frame).unwrap();
}
{
let mut trajectory = Trajectory::append(named_tmpfile.path()).unwrap();
trajectory.write(&frame).unwrap();
}
let mut trajectory = Trajectory::open(named_tmpfile.path()).unwrap();
assert_eq!(trajectory.size, 2);
check_frame(&trajectory.read().unwrap().unwrap());
check_frame(&trajectory.read().unwrap().unwrap());
}
#[test]
fn append_to_new_file() {
let frame = generate_frame();
let named_tmpfile = Builder::new()
.prefix("netcdf-test-append")
.suffix(".nc")
.tempfile()
.unwrap();
{
let mut trajectory = Trajectory::append(named_tmpfile.path()).unwrap();
trajectory.write(&frame).unwrap();
trajectory.write(&frame).unwrap();
}
let mut trajectory = Trajectory::open(named_tmpfile.path()).unwrap();
assert_eq!(trajectory.size, 2);
check_frame(&trajectory.read().unwrap().unwrap());
check_frame(&trajectory.read().unwrap().unwrap());
}
#[test]
fn rst_water() {
let path = Path::new("./src/tests-data/netcdf/water.ncrst");
let mut trajectory = Trajectory::open(path).unwrap();
assert_eq!(trajectory.size, 1);
let frame = trajectory.read().unwrap().unwrap();
assert_eq!(frame.size(), 297);
assert_eq!(
frame.properties.get("name").unwrap().as_string().unwrap(),
"Cpptraj Generated Restart"
);
}
}