use std::{
collections::HashMap,
fs::{File, OpenOptions},
io::{self, BufReader, BufWriter, ErrorKind, Read, Seek, SeekFrom, Write as _},
path::Path,
process::{Command, Stdio},
time::Instant,
};
use lin_alg::f64::Vec3;
use crate::FrameSlice;
#[derive(Clone, Debug, Default)]
pub struct OutputEnergy {
pub time: f64,
pub temperature: Option<f32>,
pub pressure: Option<f32>,
pub potential_energy: Option<f32>,
pub kinetic_energy: Option<f32>,
pub total_energy: Option<f32>,
pub volume: Option<f32>,
pub density: Option<f32>,
}
pub struct TrrMetadata {
pub num_atoms: usize,
pub num_frames: usize,
pub start_step: f32,
pub save_interval_steps: usize,
pub dt: f32,
pub end_time: f32,
}
impl TrrMetadata {
pub fn read(path: &Path) -> io::Result<Self> {
let start = Instant::now();
println!("Starting TRR frame scan to load metadata...");
const TRR_MAGIC: i32 = 1993;
let mut r = BufReader::new(File::open(path)?);
let mut num_atoms = 0usize;
let mut num_frames = 0usize;
let mut start_step = 0i32;
let mut first_time = 0.0f64;
let mut dt = 0.0f32;
let mut save_interval_steps = 0usize;
let mut end_time = 0.0f32;
loop {
let magic = {
let mut b = [0; 4];
match r.read_exact(&mut b) {
Err(e) if e.kind() == ErrorKind::UnexpectedEof => break,
other => other?,
}
i32::from_be_bytes(b)
};
if magic != TRR_MAGIC {
return Err(io::Error::new(
ErrorKind::InvalidData,
format!("TRR bad magic: expected {TRR_MAGIC}, got {magic}"),
));
}
let ver_len = trr_u32(&mut r)? as usize;
let ver_pad = (4 - (ver_len % 4)) % 4;
r.seek(SeekFrom::Current((ver_len + ver_pad) as i64))?;
let ir_size = trr_i32(&mut r)? as usize;
let e_size = trr_i32(&mut r)? as usize;
let box_size = trr_i32(&mut r)? as usize;
let vir_size = trr_i32(&mut r)? as usize;
let pres_size = trr_i32(&mut r)? as usize;
let top_size = trr_i32(&mut r)? as usize;
let sym_size = trr_i32(&mut r)? as usize;
let x_size = trr_i32(&mut r)? as usize;
let v_size = trr_i32(&mut r)? as usize;
let f_size = trr_i32(&mut r)? as usize;
let natoms = trr_i32(&mut r)? as usize;
let step = trr_i32(&mut r)?;
let _nre = trr_i32(&mut r)?;
let double_prec = if x_size > 0 && natoms > 0 {
x_size == natoms * 3 * 8
} else if box_size > 0 {
box_size == 9 * 8
} else {
false
};
let time_ps: f64 = if double_prec {
trr_f64(&mut r)?
} else {
trr_f32(&mut r)? as f64
};
if double_prec {
trr_f64(&mut r)?;
} else {
trr_f32(&mut r)?;
}
if num_frames == 0 {
num_atoms = natoms;
start_step = step;
first_time = time_ps;
} else if num_frames == 1 {
dt = (time_ps - first_time) as f32;
save_interval_steps = (step - start_step).unsigned_abs() as usize;
}
end_time = time_ps as f32;
num_frames += 1;
let skip = ir_size
+ e_size
+ box_size
+ vir_size
+ pres_size
+ top_size
+ sym_size
+ x_size
+ v_size
+ f_size;
if skip > 0 {
r.seek(SeekFrom::Current(skip as i64))?;
}
}
let elapsed = start.elapsed().as_millis();
println!("Scan complete in {elapsed} ms");
Ok(Self {
num_atoms,
num_frames,
start_step: start_step as f32,
save_interval_steps,
dt,
end_time,
})
}
}
pub fn read_trr(trr: &Path, slice: FrameSlice) -> io::Result<Vec<GromacsFrame>> {
const TRR_MAGIC: i32 = 1993;
let mut r = BufReader::new(File::open(trr)?);
let mut frames = Vec::new();
let mut frame_idx: usize = 0;
loop {
let magic = {
let mut b = [0; 4];
match r.read_exact(&mut b) {
Err(e) if e.kind() == ErrorKind::UnexpectedEof => break,
other => other?,
}
i32::from_be_bytes(b)
};
if magic != TRR_MAGIC {
return Err(io::Error::new(
ErrorKind::InvalidData,
format!("TRR bad magic: expected {TRR_MAGIC}, got {magic}"),
));
}
let ver_len = trr_u32(&mut r)? as usize;
let ver_pad = (4 - (ver_len % 4)) % 4;
r.seek(SeekFrom::Current((ver_len + ver_pad) as i64))?;
let ir_size = trr_i32(&mut r)? as usize;
let e_size = trr_i32(&mut r)? as usize;
let box_size = trr_i32(&mut r)? as usize;
let vir_size = trr_i32(&mut r)? as usize;
let pres_size = trr_i32(&mut r)? as usize;
let top_size = trr_i32(&mut r)? as usize;
let sym_size = trr_i32(&mut r)? as usize;
let x_size = trr_i32(&mut r)? as usize;
let v_size = trr_i32(&mut r)? as usize;
let f_size = trr_i32(&mut r)? as usize;
let natoms = trr_i32(&mut r)? as usize;
let _step = trr_i32(&mut r)?;
let _nre = trr_i32(&mut r)?;
let double_prec = if x_size > 0 && natoms > 0 {
x_size == natoms * 3 * 8
} else if box_size > 0 {
box_size == 9 * 8
} else {
false
};
let time_ps: f64 = if double_prec {
trr_f64(&mut r)?
} else {
trr_f32(&mut r)? as f64
};
let _lambda: f64 = if double_prec {
trr_f64(&mut r)?
} else {
trr_f32(&mut r)? as f64
};
let pre_skip = ir_size + e_size + box_size + vir_size + pres_size + top_size + sym_size;
if pre_skip > 0 {
r.seek(SeekFrom::Current(pre_skip as i64))?;
}
let in_range = match slice {
FrameSlice::Time { start, end } => {
start.map_or(true, |t| time_ps >= t) && end.map_or(true, |t| time_ps <= t)
}
FrameSlice::Index { start, end } => {
start.map_or(true, |s| frame_idx >= s) && end.map_or(true, |e| frame_idx <= e)
}
};
frame_idx += 1;
if !in_range {
r.seek(SeekFrom::Current((x_size + v_size + f_size) as i64))?;
continue;
}
let mut atom_posits = Vec::with_capacity(natoms);
if x_size > 0 {
for _ in 0..natoms {
let (x, y, z) = if double_prec {
(trr_f64(&mut r)?, trr_f64(&mut r)?, trr_f64(&mut r)?)
} else {
(
trr_f32(&mut r)? as f64,
trr_f32(&mut r)? as f64,
trr_f32(&mut r)? as f64,
)
};
atom_posits.push(Vec3 { x, y, z });
}
}
let mut atom_velocities = Vec::with_capacity(if v_size > 0 { natoms } else { 0 });
if v_size > 0 {
for _ in 0..natoms {
let (x, y, z) = if double_prec {
(trr_f64(&mut r)?, trr_f64(&mut r)?, trr_f64(&mut r)?)
} else {
(
trr_f32(&mut r)? as f64,
trr_f32(&mut r)? as f64,
trr_f32(&mut r)? as f64,
)
};
atom_velocities.push(Vec3 { x, y, z });
}
}
if f_size > 0 {
r.seek(SeekFrom::Current(f_size as i64))?;
}
frames.push(GromacsFrame {
time: time_ps,
atom_posits,
atom_velocities,
energy: None,
});
}
Ok(frames)
}
pub fn write_trr(path: &Path, frames: &[GromacsFrame]) -> io::Result<()> {
const VERSION: &[u8] = b"GMX_trn_file";
let mut w = BufWriter::new(OpenOptions::new().create(true).append(true).open(&*path)?);
for frame in frames {
let natoms = frame.atom_posits.len();
let has_vel = frame.atom_velocities.len() == natoms && natoms > 0;
let x_size = (natoms * 3 * 4) as i32;
let v_size = if has_vel { (natoms * 3 * 4) as i32 } else { 0 };
w.write_all(&1993_i32.to_be_bytes())?;
w.write_all(&(VERSION.len() as u32).to_be_bytes())?;
w.write_all(VERSION)?;
for val in [
0i32,
0,
0,
0,
0,
0,
0,
x_size,
v_size,
0,
natoms as i32,
0,
0,
] {
w.write_all(&val.to_be_bytes())?;
}
w.write_all(&(frame.time as f32).to_be_bytes())?;
w.write_all(&0_f32.to_be_bytes())?;
for p in &frame.atom_posits {
w.write_all(&(p.x as f32).to_be_bytes())?;
w.write_all(&(p.y as f32).to_be_bytes())?;
w.write_all(&(p.z as f32).to_be_bytes())?;
}
if has_vel {
for v in &frame.atom_velocities {
w.write_all(&(v.x as f32).to_be_bytes())?;
w.write_all(&(v.y as f32).to_be_bytes())?;
w.write_all(&(v.z as f32).to_be_bytes())?;
}
}
}
Ok(())
}
impl OutputEnergy {
pub fn from_edr(path: &Path) -> io::Result<Vec<Self>> {
let dir = path.parent().unwrap_or(Path::new("."));
let edr_name = path
.file_name()
.and_then(|n| n.to_str())
.ok_or_else(|| io::Error::new(ErrorKind::InvalidInput, "invalid EDR path"))?;
const XVG_OUT: &str = "energy_out.xvg";
let selection = b"1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 \
21 22 23 24 25 26 27 28 29 30 0\n";
let mut child = match Command::new("gmx")
.current_dir(dir)
.args(["energy", "-f", edr_name, "-o", XVG_OUT])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.spawn()
{
Ok(c) => c,
Err(e) if e.kind() == ErrorKind::NotFound => {
eprintln!("`gmx` not found; skipping energy parsing");
return Ok(Vec::new());
}
Err(e) => return Err(e),
};
if let Some(mut stdin) = child.stdin.take() {
let _ = stdin.write_all(selection);
}
let out = child.wait_with_output()?;
if !out.status.success() {
let msg = String::from_utf8_lossy(&out.stderr);
eprintln!("gmx energy exited non-zero: {msg}");
}
let xvg_path = dir.join(XVG_OUT);
if !xvg_path.exists() {
return Ok(Vec::new());
}
let text = std::fs::read_to_string(&xvg_path)?;
parse_xvg(&text)
}
}
#[derive(Clone, Debug)]
pub struct GromacsFrame {
pub time: f64,
pub atom_posits: Vec<Vec3>,
pub atom_velocities: Vec<Vec3>,
pub energy: Option<OutputEnergy>,
}
#[derive(Clone, Debug, Default)]
pub struct GromacsOutput {
pub log_text: String,
pub trajectory: Vec<GromacsFrame>,
pub solute_atom_count: usize,
}
impl GromacsOutput {
pub fn new(
log_text: String,
mut trajectory: Vec<GromacsFrame>,
energies: Vec<OutputEnergy>,
solute_atom_count: usize,
) -> io::Result<Self> {
attach_energies(&mut trajectory, energies);
Ok(Self {
log_text,
trajectory,
solute_atom_count,
})
}
}
pub fn parse_multi_gro(text: &str) -> io::Result<Vec<GromacsFrame>> {
let mut frames = Vec::new();
let mut lines = text.lines().peekable();
while lines.peek().is_some() {
let title = match lines.next() {
Some(l) => l,
None => break,
};
let time_ps = extract_time(title);
let natoms_str = lines.next().ok_or_else(|| {
io::Error::new(ErrorKind::UnexpectedEof, "Expected atom count after title")
})?;
let natoms: usize = natoms_str
.trim()
.split_whitespace()
.next()
.ok_or_else(|| io::Error::new(ErrorKind::InvalidData, "Empty atom count line"))?
.parse()
.map_err(|_| io::Error::new(ErrorKind::InvalidData, "Could not parse atom count"))?;
let mut posits = Vec::with_capacity(natoms);
for _ in 0..natoms {
let line = lines.next().ok_or_else(|| {
io::Error::new(ErrorKind::UnexpectedEof, "Unexpected end of atom block")
})?;
let x_nm = parse_col(line, 20, 28)?;
let y_nm = parse_col(line, 28, 36)?;
let z_nm = parse_col(line, 36, 44)?;
posits.push(Vec3 {
x: x_nm,
y: y_nm,
z: z_nm,
});
}
lines.next();
frames.push(GromacsFrame {
time: time_ps,
atom_posits: posits,
atom_velocities: Vec::new(),
energy: None,
});
}
Ok(frames)
}
fn attach_energies(frames: &mut Vec<GromacsFrame>, energies: Vec<OutputEnergy>) {
if energies.is_empty() {
return;
}
let map: HashMap<i64, OutputEnergy> = energies
.into_iter()
.map(|e| ((e.time * 1_000.0).round() as i64, e))
.collect();
for frame in frames.iter_mut() {
let key = (frame.time * 1_000.0).round() as i64;
frame.energy = map.get(&key).cloned();
}
}
fn parse_xvg(text: &str) -> io::Result<Vec<OutputEnergy>> {
let mut col_names: Vec<String> = Vec::new();
let mut rows: Vec<Vec<f64>> = Vec::new();
for line in text.lines() {
let trimmed = line.trim();
if trimmed.is_empty() || trimmed.starts_with('#') {
continue;
}
if trimmed.starts_with('@') {
if let Some(name) = parse_xvg_legend(trimmed) {
col_names.push(name);
}
continue;
}
let vals: Vec<f64> = trimmed
.split_whitespace()
.filter_map(|t| t.parse().ok())
.collect();
if !vals.is_empty() {
rows.push(vals);
}
}
let energies = rows
.into_iter()
.map(|vals| {
let time = vals.first().copied().unwrap_or(0.0);
let mut e = OutputEnergy {
time,
..OutputEnergy::default()
};
for (col, name) in col_names.iter().enumerate() {
let val = vals.get(col + 1).copied().map(|v| v as f32);
match name.as_str() {
"Temperature" => e.temperature = val,
"Pressure" => e.pressure = val,
"Potential" => e.potential_energy = val,
n if n.starts_with("Kinetic") => e.kinetic_energy = val,
n if n.starts_with("Total") => e.total_energy = val,
"Volume" => e.volume = val,
"Density" => e.density = val,
_ => {}
}
}
e
})
.collect();
Ok(energies)
}
fn parse_xvg_legend(line: &str) -> Option<String> {
let rest = line.strip_prefix('@')?.trim();
if !rest.starts_with('s') {
return None;
}
let after_index = rest.trim_start_matches(|c: char| c == 's' || c.is_ascii_digit());
let rest = after_index.trim();
let rest = rest.strip_prefix("legend")?.trim();
let name = rest.trim_matches('"');
if name.is_empty() {
None
} else {
Some(name.to_string())
}
}
fn extract_time(title: &str) -> f64 {
for sep in ["t=", "t ="] {
if let Some(pos) = title.find(sep) {
let rest = title[pos + sep.len()..].trim();
if let Some(tok) = rest.split_whitespace().next() {
if let Ok(v) = tok.parse::<f64>() {
return v;
}
}
}
}
0.0
}
fn parse_col(line: &str, start: usize, end: usize) -> io::Result<f64> {
let len = line.len();
if start >= len {
return Ok(0.0);
}
let end = end.min(len);
let s = line[start..end].trim();
if s.is_empty() {
return Ok(0.0);
}
s.parse::<f64>()
.map_err(|_| io::Error::new(ErrorKind::InvalidData, format!("Cannot parse '{s}' as f64")))
}
fn trr_i32(r: &mut impl Read) -> io::Result<i32> {
let mut b = [0u8; 4];
r.read_exact(&mut b)?;
Ok(i32::from_be_bytes(b))
}
fn trr_u32(r: &mut impl Read) -> io::Result<u32> {
let mut b = [0u8; 4];
r.read_exact(&mut b)?;
Ok(u32::from_be_bytes(b))
}
fn trr_f32(r: &mut impl Read) -> io::Result<f32> {
let mut b = [0u8; 4];
r.read_exact(&mut b)?;
Ok(f32::from_be_bytes(b))
}
fn trr_f64(r: &mut impl Read) -> io::Result<f64> {
let mut b = [0u8; 8];
r.read_exact(&mut b)?;
Ok(f64::from_be_bytes(b))
}