oxiphysics_io/
particle_formats.rs

1#![allow(clippy::needless_range_loop)]
2// Copyright 2026 COOLJAPAN OU (Team KitaSan)
3// SPDX-License-Identifier: Apache-2.0
4
5//! Particle trajectory I/O formats.
6//!
7//! Provides readers and writers for common particle simulation formats:
8//! - **XYZ** multi-frame text trajectory
9//! - **LAMMPS dump** (`ITEM: ATOMS id type x y z [vx vy vz]`)
10//! - **GROMACS GRO** (nm units, box vectors)
11//! - **Binary frame** format (compact f32 arrays with magic header)
12//! - **DCD** stub reader (header parsing + coordinate frames)
13//!
14//! [`ParticleTrajectory`] supports frame indexing, slicing, sub-sampling,
15//! and analysis (MSD, RMSF, centre-of-mass drift).
16
17#![allow(dead_code)]
18
19use std::fs::{File, OpenOptions};
20use std::io::{self, BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
21
22// ─────────────────────────────────────────────────────────────────────────────
23// ParticleFrame
24// ─────────────────────────────────────────────────────────────────────────────
25
26/// A single snapshot of a particle system at a given time-step.
27#[derive(Debug, Clone, Default)]
28pub struct ParticleFrame {
29    /// Integer time-step index.
30    pub timestep: u64,
31    /// Number of particles (redundant with `positions.len()` but convenient).
32    pub n_particles: usize,
33    /// Positions of each particle `[x, y, z]` in simulation units.
34    pub positions: Vec<[f32; 3]>,
35    /// Velocities of each particle `[vx, vy, vz]`, if present.
36    pub velocities: Option<Vec<[f32; 3]>>,
37    /// Integer type/species identifier for each particle.
38    pub types: Vec<u8>,
39}
40
41impl ParticleFrame {
42    /// Create a new frame with positions and types.
43    ///
44    /// `velocities` is set to `None`; fill it manually if needed.
45    pub fn new(timestep: u64, positions: Vec<[f32; 3]>, types: Vec<u8>) -> Self {
46        let n = positions.len();
47        Self {
48            timestep,
49            n_particles: n,
50            positions,
51            velocities: None,
52            types,
53        }
54    }
55
56    /// Create a frame where all particles have type `0` and no velocities.
57    pub fn from_positions(timestep: u64, positions: Vec<[f32; 3]>) -> Self {
58        let n = positions.len();
59        let types = vec![0u8; n];
60        Self::new(timestep, positions, types)
61    }
62
63    /// Centre of mass (unweighted, equal masses assumed).
64    ///
65    /// Returns `[0.0; 3]` for an empty frame.
66    pub fn center_of_mass(&self) -> [f32; 3] {
67        if self.positions.is_empty() {
68            return [0.0; 3];
69        }
70        let n = self.positions.len() as f32;
71        let mut com = [0.0f32; 3];
72        for p in &self.positions {
73            com[0] += p[0];
74            com[1] += p[1];
75            com[2] += p[2];
76        }
77        [com[0] / n, com[1] / n, com[2] / n]
78    }
79
80    /// Axis-aligned bounding box `(min, max)`.  Returns `([0;3],[0;3])` if empty.
81    pub fn bounding_box(&self) -> ([f32; 3], [f32; 3]) {
82        if self.positions.is_empty() {
83            return ([0.0; 3], [0.0; 3]);
84        }
85        let mut lo = self.positions[0];
86        let mut hi = self.positions[0];
87        for &p in &self.positions[1..] {
88            for d in 0..3 {
89                if p[d] < lo[d] {
90                    lo[d] = p[d];
91                }
92                if p[d] > hi[d] {
93                    hi[d] = p[d];
94                }
95            }
96        }
97        (lo, hi)
98    }
99}
100
101// ─────────────────────────────────────────────────────────────────────────────
102// XyzWriter
103// ─────────────────────────────────────────────────────────────────────────────
104
105/// Write XYZ trajectory files (multi-frame, one frame per call).
106///
107/// Each frame is written as:
108/// ```text
109/// `N`
110/// Timestep=`ts` `comment`
111/// `type` `x` `y` `z`
112/// ...
113/// ```
114#[derive(Debug, Clone)]
115pub struct XyzWriter {
116    /// Output path.
117    pub path: std::path::PathBuf,
118}
119
120impl XyzWriter {
121    /// Create a new `XyzWriter` targeting `path`.
122    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
123        Self { path: path.into() }
124    }
125
126    /// Overwrite the file with a single frame.
127    pub fn write_frame(&self, frame: &ParticleFrame, comment: &str) -> io::Result<()> {
128        let f = File::create(&self.path)?;
129        let mut w = BufWriter::new(f);
130        write_xyz_frame_to(&mut w, frame, comment)
131    }
132
133    /// Append a frame to the file (creates if not present).
134    pub fn append_frame(&self, frame: &ParticleFrame, comment: &str) -> io::Result<()> {
135        let f = std::fs::OpenOptions::new()
136            .create(true)
137            .append(true)
138            .open(&self.path)?;
139        let mut w = BufWriter::new(f);
140        write_xyz_frame_to(&mut w, frame, comment)
141    }
142
143    /// Write a complete trajectory (overwrites).
144    pub fn write_trajectory(&self, traj: &ParticleTrajectory) -> io::Result<()> {
145        let f = File::create(&self.path)?;
146        let mut w = BufWriter::new(f);
147        for frame in &traj.frames {
148            let comment = format!("Timestep={}", frame.timestep);
149            write_xyz_frame_to(&mut w, frame, &comment)?;
150        }
151        Ok(())
152    }
153}
154
155/// Write one XYZ frame to any `Write` implementor.
156fn write_xyz_frame_to<W: Write>(w: &mut W, frame: &ParticleFrame, comment: &str) -> io::Result<()> {
157    writeln!(w, "{}", frame.n_particles)?;
158    writeln!(w, "Timestep={} {}", frame.timestep, comment)?;
159    for (i, p) in frame.positions.iter().enumerate() {
160        let t = frame.types.get(i).copied().unwrap_or(0);
161        writeln!(w, "{} {:.8} {:.8} {:.8}", t, p[0], p[1], p[2])?;
162    }
163    Ok(())
164}
165
166// ─────────────────────────────────────────────────────────────────────────────
167// XyzReader
168// ─────────────────────────────────────────────────────────────────────────────
169
170/// Read XYZ trajectory files, returning one or all frames.
171#[derive(Debug, Clone)]
172pub struct XyzReader {
173    /// Input path.
174    pub path: std::path::PathBuf,
175}
176
177impl XyzReader {
178    /// Create a new `XyzReader` for `path`.
179    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
180        Self { path: path.into() }
181    }
182
183    /// Read the first frame from the file.
184    pub fn read_frame(&self) -> io::Result<ParticleFrame> {
185        let f = File::open(&self.path)?;
186        let mut r = BufReader::new(f);
187        parse_xyz_frame(&mut r)
188    }
189
190    /// Read all frames from the file.
191    pub fn read_all(&self) -> io::Result<Vec<ParticleFrame>> {
192        let f = File::open(&self.path)?;
193        let mut r = BufReader::new(f);
194        let mut frames = Vec::new();
195        loop {
196            match parse_xyz_frame(&mut r) {
197                Ok(fr) => frames.push(fr),
198                Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
199                Err(e) => return Err(e),
200            }
201        }
202        Ok(frames)
203    }
204}
205
206/// Parse one XYZ frame from a `BufRead` source.
207fn parse_xyz_frame<R: BufRead>(r: &mut R) -> io::Result<ParticleFrame> {
208    let mut line = String::new();
209
210    // Line 1: atom count.
211    r.read_line(&mut line)?;
212    if line.is_empty() {
213        return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "EOF"));
214    }
215    let n: usize = line
216        .trim()
217        .parse()
218        .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "expected atom count"))?;
219    line.clear();
220
221    // Line 2: comment — extract timestep if present.
222    r.read_line(&mut line)?;
223    let timestep = extract_timestep_from_comment(&line);
224    line.clear();
225
226    // Lines 3..(3+n): atom records.
227    let mut positions = Vec::with_capacity(n);
228    let mut types = Vec::with_capacity(n);
229    for _ in 0..n {
230        r.read_line(&mut line)?;
231        let parts: Vec<&str> = line.split_whitespace().collect();
232        if parts.len() < 4 {
233            return Err(io::Error::new(
234                io::ErrorKind::InvalidData,
235                "short atom line",
236            ));
237        }
238        let t: u8 = parts[0].parse().unwrap_or(0);
239        let x: f32 = parts[1]
240            .parse()
241            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "bad x"))?;
242        let y: f32 = parts[2]
243            .parse()
244            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "bad y"))?;
245        let z: f32 = parts[3]
246            .parse()
247            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "bad z"))?;
248        positions.push([x, y, z]);
249        types.push(t);
250        line.clear();
251    }
252
253    Ok(ParticleFrame {
254        timestep,
255        n_particles: n,
256        positions,
257        velocities: None,
258        types,
259    })
260}
261
262/// Extract `Timestep=`value` from an XYZ comment line.  Returns 0 on failure.
263fn extract_timestep_from_comment(comment: &str) -> u64 {
264    for token in comment.split_whitespace() {
265        if let Some(rest) = token.strip_prefix("Timestep=")
266            && let Ok(v) = rest.parse::<u64>()
267        {
268            return v;
269        }
270    }
271    0
272}
273
274// ─────────────────────────────────────────────────────────────────────────────
275// LammpsDumpWriter
276// ─────────────────────────────────────────────────────────────────────────────
277
278/// Write LAMMPS dump files in `ITEM: ATOMS id type x y z \[vx vy vz\]` format.
279#[derive(Debug, Clone)]
280pub struct LammpsDumpWriter {
281    /// Output path.
282    pub path: std::path::PathBuf,
283    /// Whether to write velocities when available.
284    pub write_velocities: bool,
285}
286
287impl LammpsDumpWriter {
288    /// Create a new `LammpsDumpWriter`.  Velocities are included if present.
289    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
290        Self {
291            path: path.into(),
292            write_velocities: true,
293        }
294    }
295
296    /// Append one frame to the dump file.
297    ///
298    /// `box_lo` and `box_hi` are the simulation box boundaries `\[x, y, z\]`.
299    pub fn write_frame(
300        &self,
301        frame: &ParticleFrame,
302        box_lo: [f32; 3],
303        box_hi: [f32; 3],
304    ) -> io::Result<()> {
305        let f = std::fs::OpenOptions::new()
306            .create(true)
307            .append(true)
308            .open(&self.path)?;
309        let mut w = BufWriter::new(f);
310        write_lammps_dump_frame(&mut w, frame, box_lo, box_hi, self.write_velocities)
311    }
312}
313
314/// Write one LAMMPS dump frame to any `Write` implementor.
315fn write_lammps_dump_frame<W: Write>(
316    w: &mut W,
317    frame: &ParticleFrame,
318    box_lo: [f32; 3],
319    box_hi: [f32; 3],
320    with_vel: bool,
321) -> io::Result<()> {
322    writeln!(w, "ITEM: TIMESTEP")?;
323    writeln!(w, "{}", frame.timestep)?;
324    writeln!(w, "ITEM: NUMBER OF ATOMS")?;
325    writeln!(w, "{}", frame.n_particles)?;
326    writeln!(w, "ITEM: BOX BOUNDS pp pp pp")?;
327    for d in 0..3 {
328        writeln!(w, "{:.6} {:.6}", box_lo[d], box_hi[d])?;
329    }
330    let has_vel = with_vel && frame.velocities.is_some();
331    if has_vel {
332        writeln!(w, "ITEM: ATOMS id type x y z vx vy vz")?;
333    } else {
334        writeln!(w, "ITEM: ATOMS id type x y z")?;
335    }
336    let vels = frame.velocities.as_deref();
337    for (i, p) in frame.positions.iter().enumerate() {
338        let t = frame.types.get(i).copied().unwrap_or(0);
339        if has_vel {
340            let v = vels.map(|vs| vs[i]).unwrap_or([0.0; 3]);
341            writeln!(
342                w,
343                "{} {} {:.8} {:.8} {:.8} {:.8} {:.8} {:.8}",
344                i + 1,
345                t,
346                p[0],
347                p[1],
348                p[2],
349                v[0],
350                v[1],
351                v[2]
352            )?;
353        } else {
354            writeln!(w, "{} {} {:.8} {:.8} {:.8}", i + 1, t, p[0], p[1], p[2])?;
355        }
356    }
357    Ok(())
358}
359
360// ─────────────────────────────────────────────────────────────────────────────
361// LammpsDumpReader
362// ─────────────────────────────────────────────────────────────────────────────
363
364/// Read LAMMPS dump files, handling `ITEM: ATOMS` headers and multiple timesteps.
365#[derive(Debug, Clone)]
366pub struct LammpsDumpReader {
367    /// Input path.
368    pub path: std::path::PathBuf,
369}
370
371impl LammpsDumpReader {
372    /// Create a new `LammpsDumpReader` for `path`.
373    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
374        Self { path: path.into() }
375    }
376
377    /// Read all frames from the dump file.
378    pub fn read_all(&self) -> io::Result<Vec<ParticleFrame>> {
379        let f = File::open(&self.path)?;
380        let r = BufReader::new(f);
381        parse_lammps_dump_all(r)
382    }
383}
384
385/// Parse all LAMMPS dump frames from a `BufRead` source.
386fn parse_lammps_dump_all<R: BufRead>(mut r: R) -> io::Result<Vec<ParticleFrame>> {
387    let mut frames = Vec::new();
388    let mut lines = Vec::<String>::new();
389    // Collect all lines.
390    {
391        let mut line = String::new();
392        while r.read_line(&mut line)? > 0 {
393            lines.push(line.trim_end().to_string());
394            line.clear();
395        }
396    }
397
398    let mut i = 0;
399    while i < lines.len() {
400        if lines[i].starts_with("ITEM: TIMESTEP") {
401            i += 1;
402            let timestep: u64 = lines
403                .get(i)
404                .and_then(|l| l.trim().parse().ok())
405                .unwrap_or(0);
406            i += 1; // "ITEM: NUMBER OF ATOMS"
407            i += 1;
408            let n_atoms: usize = lines
409                .get(i)
410                .and_then(|l| l.trim().parse().ok())
411                .unwrap_or(0);
412            i += 1; // "ITEM: BOX BOUNDS ..."
413            i += 4; // skip 3 bound lines + the item line itself
414
415            // "ITEM: ATOMS ..."
416            let header = lines.get(i).cloned().unwrap_or_default();
417            let cols: Vec<&str> = header.split_whitespace().collect();
418            // Find column indices for id, type, x, y, z, vx, vy, vz.
419            let idx_id = col_idx(&cols, "id");
420            let idx_type = col_idx(&cols, "type");
421            let idx_x = col_idx(&cols, "x");
422            let idx_y = col_idx(&cols, "y");
423            let idx_z = col_idx(&cols, "z");
424            let idx_vx = col_idx(&cols, "vx");
425            let idx_vy = col_idx(&cols, "vy");
426            let idx_vz = col_idx(&cols, "vz");
427            let has_vel = idx_vx.is_some() && idx_vy.is_some() && idx_vz.is_some();
428            // Skip ITEM: ATOMS line — offsets start at 2 after ITEM: ATOMS header.
429            // Columns in header are offset by 2 ("ITEM:" "ATOMS" are first two).
430            let data_offset = 2usize;
431            i += 1;
432
433            let mut positions = Vec::with_capacity(n_atoms);
434            let mut types = Vec::with_capacity(n_atoms);
435            let mut velocities: Vec<[f32; 3]> = Vec::with_capacity(n_atoms);
436
437            for _ in 0..n_atoms {
438                if i >= lines.len() {
439                    break;
440                }
441                let parts: Vec<&str> = lines[i].split_whitespace().collect();
442
443                let get = |opt_col: Option<usize>| -> f32 {
444                    opt_col
445                        .and_then(|c| {
446                            let real_c = c.saturating_sub(data_offset);
447                            parts.get(real_c).and_then(|s| s.parse().ok())
448                        })
449                        .unwrap_or(0.0)
450                };
451                let get_u8 = |opt_col: Option<usize>| -> u8 {
452                    opt_col
453                        .and_then(|c| {
454                            let real_c = c.saturating_sub(data_offset);
455                            parts.get(real_c).and_then(|s| s.parse::<u8>().ok())
456                        })
457                        .unwrap_or(0)
458                };
459                let _ = idx_id; // id unused but parsed
460                let t = get_u8(idx_type);
461                let x = get(idx_x);
462                let y = get(idx_y);
463                let z = get(idx_z);
464                positions.push([x, y, z]);
465                types.push(t);
466                if has_vel {
467                    let vx = get(idx_vx);
468                    let vy = get(idx_vy);
469                    let vz = get(idx_vz);
470                    velocities.push([vx, vy, vz]);
471                }
472                i += 1;
473            }
474
475            let vel_opt = if has_vel && !velocities.is_empty() {
476                Some(velocities)
477            } else {
478                None
479            };
480
481            frames.push(ParticleFrame {
482                timestep,
483                n_particles: positions.len(),
484                positions,
485                velocities: vel_opt,
486                types,
487            });
488        } else {
489            i += 1;
490        }
491    }
492    Ok(frames)
493}
494
495/// Find the index of a column name in the LAMMPS ITEM:ATOMS header tokens.
496/// The header starts with "ITEM:" "ATOMS", so data columns start at index 2.
497fn col_idx(cols: &[&str], name: &str) -> Option<usize> {
498    cols.iter().position(|&c| c == name)
499}
500
501// ─────────────────────────────────────────────────────────────────────────────
502// ParticleTrajectory
503// ─────────────────────────────────────────────────────────────────────────────
504
505/// An in-memory trajectory: an ordered sequence of [`ParticleFrame`]s.
506#[derive(Debug, Clone, Default)]
507pub struct ParticleTrajectory {
508    /// Ordered frames (earliest first).
509    pub frames: Vec<ParticleFrame>,
510}
511
512impl ParticleTrajectory {
513    /// Create an empty trajectory.
514    pub fn new() -> Self {
515        Self { frames: Vec::new() }
516    }
517
518    /// Number of frames.
519    pub fn n_frames(&self) -> usize {
520        self.frames.len()
521    }
522
523    /// Number of particles (from the first frame).  Returns 0 if empty.
524    pub fn n_particles(&self) -> usize {
525        self.frames.first().map(|f| f.n_particles).unwrap_or(0)
526    }
527
528    /// Get frame `i` by index.
529    pub fn get(&self, i: usize) -> Option<&ParticleFrame> {
530        self.frames.get(i)
531    }
532
533    /// Return a sub-trajectory containing frames in `\[start, end)`.
534    pub fn slice(&self, start: usize, end: usize) -> Self {
535        let end = end.min(self.frames.len());
536        Self {
537            frames: self.frames[start..end].to_vec(),
538        }
539    }
540
541    /// Return every `step`-th frame (stride sub-sampling).
542    pub fn subsample(&self, step: usize) -> Self {
543        if step == 0 {
544            return Self::new();
545        }
546        Self {
547            frames: self.frames.iter().step_by(step).cloned().collect(),
548        }
549    }
550
551    /// Append a frame to the trajectory.
552    pub fn push(&mut self, frame: ParticleFrame) {
553        self.frames.push(frame);
554    }
555}
556
557// ─────────────────────────────────────────────────────────────────────────────
558// TrajectoryStats
559// ─────────────────────────────────────────────────────────────────────────────
560
561/// Compute analysis statistics over a [`ParticleTrajectory`].
562#[derive(Debug, Clone, Default)]
563pub struct TrajectoryStats {
564    /// Mean square displacement per lag step (length = n_frames).
565    pub msd: Vec<f64>,
566    /// Root mean square fluctuation per particle (length = n_particles).
567    pub rmsf: Vec<f64>,
568    /// Centre-of-mass position per frame (length = n_frames).
569    pub com_drift: Vec<[f64; 3]>,
570}
571
572impl TrajectoryStats {
573    /// Compute all statistics from a trajectory.
574    pub fn compute(traj: &ParticleTrajectory) -> Self {
575        Self {
576            msd: compute_msd(traj),
577            rmsf: compute_rmsf(traj),
578            com_drift: compute_com_drift(traj),
579        }
580    }
581}
582
583/// Compute MSD: `msd\[lag\] = <|r(t+lag) - r(t)|^2>` averaged over all
584/// particles and all valid origins.
585fn compute_msd(traj: &ParticleTrajectory) -> Vec<f64> {
586    let nf = traj.n_frames();
587    let np = traj.n_particles();
588    if nf == 0 || np == 0 {
589        return vec![];
590    }
591    let mut result = vec![0.0f64; nf];
592    for lag in 0..nf {
593        let mut sum = 0.0f64;
594        let mut count = 0usize;
595        for t in 0..(nf - lag) {
596            let f0 = &traj.frames[t];
597            let f1 = &traj.frames[t + lag];
598            let n = f0.positions.len().min(f1.positions.len());
599            for i in 0..n {
600                let r0 = f0.positions[i];
601                let r1 = f1.positions[i];
602                let dr2 =
603                    (r1[0] - r0[0]).powi(2) + (r1[1] - r0[1]).powi(2) + (r1[2] - r0[2]).powi(2);
604                sum += dr2 as f64;
605                count += 1;
606            }
607        }
608        result[lag] = if count > 0 { sum / count as f64 } else { 0.0 };
609    }
610    result
611}
612
613/// Compute RMSF: root-mean-square fluctuation of each particle about its
614/// time-averaged position.
615fn compute_rmsf(traj: &ParticleTrajectory) -> Vec<f64> {
616    let nf = traj.n_frames();
617    let np = traj.n_particles();
618    if nf == 0 || np == 0 {
619        return vec![];
620    }
621    // Compute mean position for each particle.
622    let mut mean = vec![[0.0f64; 3]; np];
623    for frame in &traj.frames {
624        for (i, p) in frame.positions.iter().enumerate().take(np) {
625            mean[i][0] += p[0] as f64;
626            mean[i][1] += p[1] as f64;
627            mean[i][2] += p[2] as f64;
628        }
629    }
630    for m in &mut mean {
631        m[0] /= nf as f64;
632        m[1] /= nf as f64;
633        m[2] /= nf as f64;
634    }
635    // Compute variance.
636    let mut var = vec![0.0f64; np];
637    for frame in &traj.frames {
638        for (i, p) in frame.positions.iter().enumerate().take(np) {
639            let dx = p[0] as f64 - mean[i][0];
640            let dy = p[1] as f64 - mean[i][1];
641            let dz = p[2] as f64 - mean[i][2];
642            var[i] += dx * dx + dy * dy + dz * dz;
643        }
644    }
645    var.iter().map(|&v| (v / nf as f64).sqrt()).collect()
646}
647
648/// Compute centre-of-mass position for each frame.
649fn compute_com_drift(traj: &ParticleTrajectory) -> Vec<[f64; 3]> {
650    traj.frames
651        .iter()
652        .map(|f| {
653            if f.positions.is_empty() {
654                [0.0; 3]
655            } else {
656                let n = f.positions.len() as f64;
657                let mut com = [0.0f64; 3];
658                for p in &f.positions {
659                    com[0] += p[0] as f64;
660                    com[1] += p[1] as f64;
661                    com[2] += p[2] as f64;
662                }
663                [com[0] / n, com[1] / n, com[2] / n]
664            }
665        })
666        .collect()
667}
668
669// ─────────────────────────────────────────────────────────────────────────────
670// BinaryFrameWriter / BinaryFrameReader
671// ─────────────────────────────────────────────────────────────────────────────
672
673/// Magic bytes used to identify the binary frame format.
674const BINARY_MAGIC: u32 = 0x4F58_4950; // "OXIP"
675
676/// Write compact binary frames: `MAGIC(4) | timestep(8) | n(4) | x0y0z0... | has_vel(1) | vx0vy0vz0...`.
677#[derive(Debug, Clone)]
678pub struct BinaryFrameWriter {
679    /// Output path.
680    pub path: std::path::PathBuf,
681}
682
683impl BinaryFrameWriter {
684    /// Create a new `BinaryFrameWriter` targeting `path`.
685    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
686        Self { path: path.into() }
687    }
688
689    /// Append one frame in binary format.
690    pub fn write_frame(&self, frame: &ParticleFrame) -> io::Result<()> {
691        let f = std::fs::OpenOptions::new()
692            .create(true)
693            .append(true)
694            .open(&self.path)?;
695        let mut w = BufWriter::new(f);
696        write_binary_frame(&mut w, frame)
697    }
698}
699
700/// Write one binary frame to any `Write` implementor.
701fn write_binary_frame<W: Write>(w: &mut W, frame: &ParticleFrame) -> io::Result<()> {
702    w.write_all(&BINARY_MAGIC.to_le_bytes())?;
703    w.write_all(&frame.timestep.to_le_bytes())?;
704    w.write_all(&(frame.n_particles as u32).to_le_bytes())?;
705    for p in &frame.positions {
706        w.write_all(&p[0].to_le_bytes())?;
707        w.write_all(&p[1].to_le_bytes())?;
708        w.write_all(&p[2].to_le_bytes())?;
709    }
710    // types
711    for &t in &frame.types {
712        w.write_all(&[t])?;
713    }
714    // velocities flag + data
715    let has_vel = frame.velocities.is_some();
716    w.write_all(&[has_vel as u8])?;
717    if let Some(vels) = &frame.velocities {
718        for v in vels {
719            w.write_all(&v[0].to_le_bytes())?;
720            w.write_all(&v[1].to_le_bytes())?;
721            w.write_all(&v[2].to_le_bytes())?;
722        }
723    }
724    Ok(())
725}
726
727/// Read compact binary frames written by [`BinaryFrameWriter`].
728#[derive(Debug, Clone)]
729pub struct BinaryFrameReader {
730    /// Input path.
731    pub path: std::path::PathBuf,
732}
733
734impl BinaryFrameReader {
735    /// Create a new `BinaryFrameReader` for `path`.
736    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
737        Self { path: path.into() }
738    }
739
740    /// Read all frames from the binary file.
741    pub fn read_all(&self) -> io::Result<Vec<ParticleFrame>> {
742        let f = File::open(&self.path)?;
743        let mut r = BufReader::new(f);
744        let mut frames = Vec::new();
745        loop {
746            match read_binary_frame(&mut r) {
747                Ok(fr) => frames.push(fr),
748                Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
749                Err(e) => return Err(e),
750            }
751        }
752        Ok(frames)
753    }
754}
755
756/// Read one binary frame from a `Read` implementor.
757fn read_binary_frame<R: Read>(r: &mut R) -> io::Result<ParticleFrame> {
758    let magic = read_u32_le(r)?;
759    if magic != BINARY_MAGIC {
760        return Err(io::Error::new(io::ErrorKind::InvalidData, "bad magic"));
761    }
762    let timestep = read_u64_le(r)?;
763    let n = read_u32_le(r)? as usize;
764
765    let mut positions = Vec::with_capacity(n);
766    for _ in 0..n {
767        let x = read_f32_le(r)?;
768        let y = read_f32_le(r)?;
769        let z = read_f32_le(r)?;
770        positions.push([x, y, z]);
771    }
772    let mut types = vec![0u8; n];
773    r.read_exact(&mut types)?;
774
775    let mut has_vel_buf = [0u8; 1];
776    r.read_exact(&mut has_vel_buf)?;
777    let velocities = if has_vel_buf[0] != 0 {
778        let mut vels = Vec::with_capacity(n);
779        for _ in 0..n {
780            let vx = read_f32_le(r)?;
781            let vy = read_f32_le(r)?;
782            let vz = read_f32_le(r)?;
783            vels.push([vx, vy, vz]);
784        }
785        Some(vels)
786    } else {
787        None
788    };
789
790    Ok(ParticleFrame {
791        timestep,
792        n_particles: n,
793        positions,
794        velocities,
795        types,
796    })
797}
798
799// ── binary I/O helpers ────────────────────────────────────────────────────────
800
801fn read_u32_le<R: Read>(r: &mut R) -> io::Result<u32> {
802    let mut buf = [0u8; 4];
803    r.read_exact(&mut buf)?;
804    Ok(u32::from_le_bytes(buf))
805}
806
807fn read_u64_le<R: Read>(r: &mut R) -> io::Result<u64> {
808    let mut buf = [0u8; 8];
809    r.read_exact(&mut buf)?;
810    Ok(u64::from_le_bytes(buf))
811}
812
813fn read_f32_le<R: Read>(r: &mut R) -> io::Result<f32> {
814    let mut buf = [0u8; 4];
815    r.read_exact(&mut buf)?;
816    Ok(f32::from_le_bytes(buf))
817}
818
819// ─────────────────────────────────────────────────────────────────────────────
820// GroWriter / GroReader
821// ─────────────────────────────────────────────────────────────────────────────
822
823/// Write GROMACS `.gro` format snapshots (coordinates in nm).
824///
825/// Positions are stored as-is; the caller is responsible for unit conversion.
826#[derive(Debug, Clone)]
827pub struct GroWriter {
828    /// Output path.
829    pub path: std::path::PathBuf,
830}
831
832impl GroWriter {
833    /// Create a new `GroWriter` targeting `path`.
834    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
835        Self { path: path.into() }
836    }
837
838    /// Write a single frame in GRO format with optional box vectors.
839    ///
840    /// `box_vecs` should be `\[lx, ly, lz\]` in nm.
841    pub fn write_frame(&self, frame: &ParticleFrame, box_vecs: [f32; 3]) -> io::Result<()> {
842        let f = File::create(&self.path)?;
843        let mut w = BufWriter::new(f);
844        writeln!(w, "Generated by oxiphysics-io  t={}", frame.timestep)?;
845        writeln!(w, "{}", frame.n_particles)?;
846        for (i, p) in frame.positions.iter().enumerate() {
847            let t = frame.types.get(i).copied().unwrap_or(0);
848            // GRO: %5d%-5s%5s%5d%8.3f%8.3f%8.3f
849            writeln!(
850                w,
851                "{:5}UNK  {:>4}{:5}{:8.3}{:8.3}{:8.3}",
852                (i + 1) % 100_000,
853                t,
854                (i + 1) % 100_000,
855                p[0],
856                p[1],
857                p[2]
858            )?;
859        }
860        writeln!(
861            w,
862            "{:.5}  {:.5}  {:.5}",
863            box_vecs[0], box_vecs[1], box_vecs[2]
864        )?;
865        Ok(())
866    }
867}
868
869/// Read GROMACS `.gro` format snapshots.
870#[derive(Debug, Clone)]
871pub struct GroReader {
872    /// Input path.
873    pub path: std::path::PathBuf,
874}
875
876impl GroReader {
877    /// Create a new `GroReader` for `path`.
878    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
879        Self { path: path.into() }
880    }
881
882    /// Read the first frame from the GRO file.
883    pub fn read_frame(&self) -> io::Result<ParticleFrame> {
884        let f = File::open(&self.path)?;
885        let mut r = BufReader::new(f);
886        parse_gro_frame(&mut r)
887    }
888}
889
890/// Parse one GRO frame from a `BufRead` source.
891fn parse_gro_frame<R: BufRead>(r: &mut R) -> io::Result<ParticleFrame> {
892    let mut line = String::new();
893    // Title line.
894    r.read_line(&mut line)?;
895    if line.is_empty() {
896        return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "EOF"));
897    }
898    line.clear();
899    // Atom count.
900    r.read_line(&mut line)?;
901    let n: usize = line
902        .trim()
903        .parse()
904        .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "expected atom count"))?;
905    line.clear();
906
907    let mut positions = Vec::with_capacity(n);
908    let mut types = Vec::with_capacity(n);
909    for _ in 0..n {
910        r.read_line(&mut line)?;
911        // GRO fixed-width: fields at specific columns. Use whitespace split as fallback.
912        let parts: Vec<&str> = line.split_whitespace().collect();
913        // fields: resid resname atomname atomid x y z [vx vy vz]
914        if parts.len() < 5 {
915            return Err(io::Error::new(io::ErrorKind::InvalidData, "short GRO line"));
916        }
917        // atom type from atomname (try parse as u8, else 0)
918        let t: u8 = parts[1].parse().unwrap_or(0);
919        let x: f32 = parts[3]
920            .parse()
921            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "bad x"))?;
922        let y: f32 = parts[4]
923            .parse()
924            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "bad y"))?;
925        let z: f32 = if parts.len() > 5 {
926            parts[5].parse().unwrap_or(0.0)
927        } else {
928            0.0
929        };
930        positions.push([x, y, z]);
931        types.push(t);
932        line.clear();
933    }
934    // Box line (ignore for now).
935    Ok(ParticleFrame {
936        timestep: 0,
937        n_particles: n,
938        positions,
939        velocities: None,
940        types,
941    })
942}
943
944// ─────────────────────────────────────────────────────────────────────────────
945// DcdReader / DcdWriter
946// ─────────────────────────────────────────────────────────────────────────────
947
948/// Header information parsed from a DCD file.
949#[derive(Debug, Clone, Default)]
950pub struct DcdHeader {
951    /// Number of frames stored in the file.
952    pub n_frames: u32,
953    /// First timestep number.
954    pub first_step: u32,
955    /// Number of steps between saved frames.
956    pub step_interval: u32,
957    /// Number of atoms.
958    pub n_atoms: u32,
959    /// Title strings from the header.
960    pub titles: Vec<String>,
961}
962
963/// Stub DCD binary reader: parses the header and reads coordinate frames.
964///
965/// Supports the standard CHARMM/NAMD little-endian DCD format.
966#[derive(Debug)]
967pub struct DcdReader {
968    /// Input path.
969    pub path: std::path::PathBuf,
970}
971
972impl DcdReader {
973    /// Create a new `DcdReader` for `path`.
974    pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
975        Self { path: path.into() }
976    }
977
978    /// Parse only the DCD file header.
979    pub fn read_header(&self) -> io::Result<DcdHeader> {
980        let f = File::open(&self.path)?;
981        let mut r = BufReader::new(f);
982        parse_dcd_header(&mut r)
983    }
984
985    /// Read all coordinate frames.  Returns `(header, frames)`.
986    pub fn read_all(&self) -> io::Result<(DcdHeader, Vec<ParticleFrame>)> {
987        let f = File::open(&self.path)?;
988        let mut r = BufReader::new(f);
989        let header = parse_dcd_header(&mut r)?;
990        let n = header.n_atoms as usize;
991        let mut frames = Vec::with_capacity(header.n_frames as usize);
992        for step in 0..header.n_frames {
993            match parse_dcd_frame(&mut r, n, step as u64) {
994                Ok(fr) => frames.push(fr),
995                Err(e) if e.kind() == io::ErrorKind::UnexpectedEof => break,
996                Err(e) => return Err(e),
997            }
998        }
999        Ok((header, frames))
1000    }
1001}
1002
1003/// Parse DCD header from a `Read + Seek` source.
1004fn parse_dcd_header<R: Read + Seek>(r: &mut R) -> io::Result<DcdHeader> {
1005    // Block 1: 4-byte Fortran record length, "CORD", 9 ints, delta, 10 ints.
1006    let rec_len = read_u32_le(r)?;
1007    if rec_len < 4 {
1008        return Err(io::Error::new(
1009            io::ErrorKind::InvalidData,
1010            "DCD block1 too short",
1011        ));
1012    }
1013    let mut magic = [0u8; 4];
1014    r.read_exact(&mut magic)?;
1015    // Expect "CORD" or "VELO"
1016    let n_frames = read_u32_le(r)?;
1017    let first_step = read_u32_le(r)?;
1018    let step_interval = read_u32_le(r)?;
1019    // Skip remaining ints + delta + padding to end of block.
1020    let remaining = rec_len as usize - 4 - 4 * 3;
1021    r.seek(SeekFrom::Current(remaining as i64))?;
1022    let _end_len = read_u32_le(r)?; // closing Fortran record
1023
1024    // Block 2: title block.
1025    let title_block_len = read_u32_le(r)? as usize;
1026    let n_titles = read_u32_le(r)? as usize;
1027    let mut titles = Vec::new();
1028    for _ in 0..n_titles {
1029        let mut buf = vec![0u8; 80];
1030        r.read_exact(&mut buf)?;
1031        titles.push(String::from_utf8_lossy(&buf).trim_end().to_string());
1032    }
1033    let consumed = 4 + n_titles * 80;
1034    if consumed < title_block_len {
1035        r.seek(SeekFrom::Current((title_block_len - consumed) as i64))?;
1036    }
1037    let _end_title = read_u32_le(r)?;
1038
1039    // Block 3: n_atoms.
1040    let _len3 = read_u32_le(r)?;
1041    let n_atoms = read_u32_le(r)?;
1042    let _end3 = read_u32_le(r)?;
1043
1044    Ok(DcdHeader {
1045        n_frames,
1046        first_step,
1047        step_interval,
1048        n_atoms,
1049        titles,
1050    })
1051}
1052
1053/// Parse one DCD coordinate frame.
1054fn parse_dcd_frame<R: Read>(r: &mut R, n_atoms: usize, step: u64) -> io::Result<ParticleFrame> {
1055    let read_f32_block = |r: &mut R| -> io::Result<Vec<f32>> {
1056        let len = read_u32_le(r)? as usize;
1057        let count = len / 4;
1058        let mut vals = Vec::with_capacity(count);
1059        for _ in 0..count {
1060            vals.push(read_f32_le(r)?);
1061        }
1062        let _end = read_u32_le(r)?;
1063        Ok(vals)
1064    };
1065
1066    let xs = read_f32_block(r)?;
1067    let ys = read_f32_block(r)?;
1068    let zs = read_f32_block(r)?;
1069
1070    let n = n_atoms.min(xs.len()).min(ys.len()).min(zs.len());
1071    let positions: Vec<[f32; 3]> = (0..n).map(|i| [xs[i], ys[i], zs[i]]).collect();
1072    let types = vec![0u8; n];
1073
1074    Ok(ParticleFrame {
1075        timestep: step,
1076        n_particles: n,
1077        positions,
1078        velocities: None,
1079        types,
1080    })
1081}
1082
1083/// Write a slice of `f32` values as a Fortran-style binary record.
1084///
1085/// The record is bracketed by a leading and trailing 4-byte little-endian
1086/// integer whose value equals `data.len() * 4` (the byte count of the payload).
1087fn write_f32_fortran_block<W: Write>(w: &mut W, data: &[f32]) -> io::Result<()> {
1088    let byte_len = (data.len() * 4) as u32;
1089    w.write_all(&byte_len.to_le_bytes())?;
1090    for &v in data {
1091        w.write_all(&v.to_le_bytes())?;
1092    }
1093    w.write_all(&byte_len.to_le_bytes())?;
1094    Ok(())
1095}
1096
1097/// DCD trajectory writer — produces files in the standard CHARMM/NAMD binary
1098/// DCD format that can be read by [`DcdReader`] and common molecular dynamics
1099/// visualisation tools (VMD, NAMD, etc.).
1100///
1101/// # Usage
1102///
1103/// ```rust,no_run
1104/// use oxiphysics_io::DcdWriter;
1105///
1106/// let mut w = DcdWriter::new("/tmp/traj.dcd", 10, 2.0e-3).unwrap();
1107/// // write frames …
1108/// w.finalize().unwrap();  // patches the frame count in the header
1109/// ```
1110#[derive(Debug)]
1111pub struct DcdWriter {
1112    /// Output path.
1113    pub path: std::path::PathBuf,
1114    /// Number of atoms per frame.
1115    pub n_atoms: u32,
1116    /// Integration time-step (ps).
1117    pub delta: f32,
1118    frames_written: u32,
1119}
1120
1121impl DcdWriter {
1122    /// Create a new `DcdWriter`, truncating / creating the output file and
1123    /// writing an initial DCD header.
1124    pub fn new(path: impl Into<std::path::PathBuf>, n_atoms: u32, delta: f32) -> io::Result<Self> {
1125        let path = path.into();
1126        let mut writer = Self {
1127            path,
1128            n_atoms,
1129            delta,
1130            frames_written: 0,
1131        };
1132        writer.write_header(0)?;
1133        Ok(writer)
1134    }
1135
1136    /// Append one coordinate frame to the DCD file.
1137    ///
1138    /// The frame must contain at least `n_atoms` positions; extra positions are
1139    /// silently ignored, missing ones are padded with `0.0`.
1140    pub fn write_frame(&mut self, frame: &ParticleFrame) -> io::Result<()> {
1141        use std::io::Write as _;
1142        let n = self.n_atoms as usize;
1143
1144        let mut xs = Vec::with_capacity(n);
1145        let mut ys = Vec::with_capacity(n);
1146        let mut zs = Vec::with_capacity(n);
1147        for i in 0..n {
1148            if i < frame.positions.len() {
1149                xs.push(frame.positions[i][0]);
1150                ys.push(frame.positions[i][1]);
1151                zs.push(frame.positions[i][2]);
1152            } else {
1153                xs.push(0.0_f32);
1154                ys.push(0.0_f32);
1155                zs.push(0.0_f32);
1156            }
1157        }
1158
1159        let file = OpenOptions::new().append(true).open(&self.path)?;
1160        let mut w = BufWriter::new(file);
1161        write_f32_fortran_block(&mut w, &xs)?;
1162        write_f32_fortran_block(&mut w, &ys)?;
1163        write_f32_fortran_block(&mut w, &zs)?;
1164        w.flush()?;
1165        self.frames_written += 1;
1166        Ok(())
1167    }
1168
1169    /// Patch the frame count stored in the DCD header (byte offset 8) with the
1170    /// actual number of frames written, then flush the file.
1171    ///
1172    /// Call this once after all frames have been written.
1173    pub fn finalize(&self) -> io::Result<()> {
1174        use std::io::{Seek, SeekFrom, Write as _};
1175        let mut file = OpenOptions::new().write(true).open(&self.path)?;
1176        // The frame count is stored at byte 8 (after the 4-byte Fortran record
1177        // length and the 4-byte "CORD" magic).
1178        file.seek(SeekFrom::Start(8))?;
1179        file.write_all(&self.frames_written.to_le_bytes())?;
1180        file.flush()
1181    }
1182
1183    /// Write the DCD file header.
1184    ///
1185    /// `n_frames` is the preliminary frame count; it is updated by
1186    /// [`finalize`](DcdWriter::finalize).
1187    fn write_header(&mut self, n_frames: u32) -> io::Result<()> {
1188        use std::io::Write as _;
1189        let mut f = BufWriter::new(File::create(&self.path)?);
1190
1191        // ── Block 1 (84 bytes payload) ────────────────────────────────────────
1192        // Fortran record: 84 bytes.
1193        let block1_payload: u32 = 84;
1194        f.write_all(&block1_payload.to_le_bytes())?; // leading sentinel
1195        f.write_all(b"CORD")?; // magic
1196        f.write_all(&n_frames.to_le_bytes())?; // nframes  (offset 8)
1197        f.write_all(&0u32.to_le_bytes())?; // first step
1198        f.write_all(&1u32.to_le_bytes())?; // step interval
1199        f.write_all(&0u32.to_le_bytes())?; // last step (unknown)
1200        f.write_all(&0u32.to_le_bytes())?; // reserved
1201        f.write_all(&0u32.to_le_bytes())?; // reserved
1202        f.write_all(&0u32.to_le_bytes())?; // reserved
1203        f.write_all(&0u32.to_le_bytes())?; // n_free_atoms
1204        f.write_all(&self.delta.to_le_bytes())?; // timestep (f32)
1205        // 10 padding integers (unit-cell flag + 9 zeros)
1206        for _ in 0..10u32 {
1207            f.write_all(&0u32.to_le_bytes())?;
1208        }
1209        // 4-byte CHARMM version word
1210        f.write_all(&0u32.to_le_bytes())?;
1211        f.write_all(&block1_payload.to_le_bytes())?; // trailing sentinel
1212
1213        // ── Block 2: title block ──────────────────────────────────────────────
1214        // 1 title of 80 bytes.
1215        let title_payload: u32 = 4 + 80; // n_titles (u32) + 1 × 80-byte title
1216        f.write_all(&title_payload.to_le_bytes())?;
1217        f.write_all(&1u32.to_le_bytes())?; // n_titles = 1
1218        let mut title_buf = [b' '; 80];
1219        let label = b"Created by oxiphysics DcdWriter";
1220        let copy_len = label.len().min(80);
1221        title_buf[..copy_len].copy_from_slice(&label[..copy_len]);
1222        f.write_all(&title_buf)?;
1223        f.write_all(&title_payload.to_le_bytes())?;
1224
1225        // ── Block 3: n_atoms ─────────────────────────────────────────────────
1226        f.write_all(&4u32.to_le_bytes())?; // payload = 4 bytes
1227        f.write_all(&self.n_atoms.to_le_bytes())?;
1228        f.write_all(&4u32.to_le_bytes())?;
1229
1230        f.flush()
1231    }
1232}
1233
1234// ─────────────────────────────────────────────────────────────────────────────
1235// Tests
1236// ─────────────────────────────────────────────────────────────────────────────
1237
1238#[cfg(test)]
1239mod tests {
1240    use super::*;
1241    use std::io::Cursor;
1242
1243    // ── helpers ───────────────────────────────────────────────────────────────
1244
1245    fn make_frame(n: usize) -> ParticleFrame {
1246        let positions: Vec<[f32; 3]> = (0..n).map(|i| [i as f32, 0.0, 0.0]).collect();
1247        let types = vec![1u8; n];
1248        ParticleFrame::new(0, positions, types)
1249    }
1250
1251    fn make_traj(n_frames: usize, n_particles: usize) -> ParticleTrajectory {
1252        let mut traj = ParticleTrajectory::new();
1253        for t in 0..n_frames {
1254            let positions: Vec<[f32; 3]> = (0..n_particles).map(|_| [t as f32, 0.0, 0.0]).collect();
1255            traj.push(ParticleFrame::new(
1256                t as u64,
1257                positions,
1258                vec![0u8; n_particles],
1259            ));
1260        }
1261        traj
1262    }
1263
1264    // ── ParticleFrame ─────────────────────────────────────────────────────────
1265
1266    #[test]
1267    fn test_frame_new_sets_n_particles() {
1268        let f = make_frame(5);
1269        assert_eq!(f.n_particles, 5);
1270    }
1271
1272    #[test]
1273    fn test_frame_from_positions() {
1274        let pos = vec![[1.0f32, 0.0, 0.0], [2.0, 0.0, 0.0]];
1275        let f = ParticleFrame::from_positions(10, pos);
1276        assert_eq!(f.n_particles, 2);
1277        assert_eq!(f.timestep, 10);
1278        assert_eq!(f.types, vec![0u8, 0]);
1279    }
1280
1281    #[test]
1282    fn test_frame_center_of_mass_empty() {
1283        let f = make_frame(0);
1284        assert_eq!(f.center_of_mass(), [0.0; 3]);
1285    }
1286
1287    #[test]
1288    fn test_frame_center_of_mass_single() {
1289        let f = ParticleFrame::from_positions(0, vec![[2.0f32, 4.0, 6.0]]);
1290        let com = f.center_of_mass();
1291        assert!((com[0] - 2.0).abs() < 1e-6);
1292        assert!((com[1] - 4.0).abs() < 1e-6);
1293        assert!((com[2] - 6.0).abs() < 1e-6);
1294    }
1295
1296    #[test]
1297    fn test_frame_center_of_mass_symmetric() {
1298        let pos = vec![[-1.0f32, 0.0, 0.0], [1.0, 0.0, 0.0]];
1299        let f = ParticleFrame::from_positions(0, pos);
1300        let com = f.center_of_mass();
1301        assert!(com[0].abs() < 1e-6);
1302    }
1303
1304    #[test]
1305    fn test_frame_bounding_box_empty() {
1306        let f = make_frame(0);
1307        let (lo, hi) = f.bounding_box();
1308        assert_eq!(lo, [0.0; 3]);
1309        assert_eq!(hi, [0.0; 3]);
1310    }
1311
1312    #[test]
1313    fn test_frame_bounding_box_basic() {
1314        let pos = vec![[1.0f32, -1.0, 0.0], [-1.0, 2.0, 3.0]];
1315        let f = ParticleFrame::from_positions(0, pos);
1316        let (lo, hi) = f.bounding_box();
1317        assert!((lo[0] - (-1.0)).abs() < 1e-6);
1318        assert!((hi[1] - 2.0).abs() < 1e-6);
1319        assert!((hi[2] - 3.0).abs() < 1e-6);
1320    }
1321
1322    // ── XYZ round-trip ────────────────────────────────────────────────────────
1323
1324    #[test]
1325    fn test_xyz_write_frame_internal() {
1326        let frame = make_frame(3);
1327        let mut buf = Vec::new();
1328        write_xyz_frame_to(&mut buf, &frame, "test").unwrap();
1329        let s = String::from_utf8(buf).unwrap();
1330        assert!(s.starts_with("3\n"));
1331        assert!(s.contains("test"));
1332    }
1333
1334    #[test]
1335    fn test_xyz_roundtrip_file() {
1336        let path = std::env::temp_dir().join("oxiphysics_pf_xyz_rt.xyz");
1337        let _ = std::fs::remove_file(&path);
1338        let pos = vec![[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]];
1339        let frame = ParticleFrame::new(42, pos.clone(), vec![1u8, 2]);
1340        XyzWriter::new(&path).write_frame(&frame, "test").unwrap();
1341        let loaded = XyzReader::new(&path).read_frame().unwrap();
1342        assert_eq!(loaded.n_particles, 2);
1343        assert_eq!(loaded.timestep, 42);
1344        assert!((loaded.positions[0][0] - 1.0).abs() < 1e-4);
1345        assert!((loaded.positions[1][2] - 6.0).abs() < 1e-4);
1346        let _ = std::fs::remove_file(&path);
1347    }
1348
1349    #[test]
1350    fn test_xyz_append_and_read_all() {
1351        let path = std::env::temp_dir().join("oxiphysics_pf_xyz_all.xyz");
1352        let _ = std::fs::remove_file(&path);
1353        let w = XyzWriter::new(&path);
1354        for ts in 0..4u64 {
1355            let f = ParticleFrame::new(ts, vec![[ts as f32, 0.0, 0.0]], vec![0u8]);
1356            w.append_frame(&f, "").unwrap();
1357        }
1358        let frames = XyzReader::new(&path).read_all().unwrap();
1359        assert_eq!(frames.len(), 4);
1360        let _ = std::fs::remove_file(&path);
1361    }
1362
1363    #[test]
1364    fn test_xyz_write_trajectory() {
1365        let path = std::env::temp_dir().join("oxiphysics_pf_xyz_traj.xyz");
1366        let _ = std::fs::remove_file(&path);
1367        let mut traj = ParticleTrajectory::new();
1368        for ts in 0..3u64 {
1369            traj.push(ParticleFrame::new(ts, vec![[0.0f32; 3]], vec![0u8]));
1370        }
1371        XyzWriter::new(&path).write_trajectory(&traj).unwrap();
1372        let frames = XyzReader::new(&path).read_all().unwrap();
1373        assert_eq!(frames.len(), 3);
1374        let _ = std::fs::remove_file(&path);
1375    }
1376
1377    #[test]
1378    fn test_extract_timestep_from_comment() {
1379        assert_eq!(extract_timestep_from_comment("Timestep=100 other"), 100);
1380        assert_eq!(extract_timestep_from_comment("no timestep here"), 0);
1381    }
1382
1383    // ── LAMMPS dump ───────────────────────────────────────────────────────────
1384
1385    #[test]
1386    fn test_lammps_dump_write_internal() {
1387        let frame = make_frame(2);
1388        let mut buf = Vec::new();
1389        write_lammps_dump_frame(&mut buf, &frame, [0.0; 3], [10.0; 3], false).unwrap();
1390        let s = String::from_utf8(buf).unwrap();
1391        assert!(s.contains("ITEM: TIMESTEP"));
1392        assert!(s.contains("ITEM: ATOMS id type x y z"));
1393    }
1394
1395    #[test]
1396    fn test_lammps_dump_write_with_velocities() {
1397        let mut frame = make_frame(2);
1398        frame.velocities = Some(vec![[1.0f32, 0.0, 0.0], [0.0, 1.0, 0.0]]);
1399        let mut buf = Vec::new();
1400        write_lammps_dump_frame(&mut buf, &frame, [0.0; 3], [10.0; 3], true).unwrap();
1401        let s = String::from_utf8(buf).unwrap();
1402        assert!(s.contains("vx vy vz"));
1403    }
1404
1405    #[test]
1406    fn test_lammps_dump_roundtrip_file() {
1407        let path = std::env::temp_dir().join("oxiphysics_pf_lmp.lammpstrj");
1408        let _ = std::fs::remove_file(&path);
1409        let w = LammpsDumpWriter::new(&path);
1410        for ts in 0..3u64 {
1411            let f = ParticleFrame::new(ts, vec![[ts as f32, 0.0, 0.0]], vec![1u8]);
1412            w.write_frame(&f, [0.0; 3], [10.0; 3]).unwrap();
1413        }
1414        let frames = LammpsDumpReader::new(&path).read_all().unwrap();
1415        assert_eq!(frames.len(), 3);
1416        assert_eq!(frames[1].timestep, 1);
1417        let _ = std::fs::remove_file(&path);
1418    }
1419
1420    #[test]
1421    fn test_lammps_dump_parse_inline() {
1422        let dump = "\
1423ITEM: TIMESTEP\n\
14245\n\
1425ITEM: NUMBER OF ATOMS\n\
14262\n\
1427ITEM: BOX BOUNDS pp pp pp\n\
14280.0 10.0\n\
14290.0 10.0\n\
14300.0 10.0\n\
1431ITEM: ATOMS id type x y z\n\
14321 1 1.0 2.0 3.0\n\
14332 2 4.0 5.0 6.0\n";
1434        let r = BufReader::new(Cursor::new(dump));
1435        let frames = parse_lammps_dump_all(r).unwrap();
1436        assert_eq!(frames.len(), 1);
1437        assert_eq!(frames[0].timestep, 5);
1438        assert_eq!(frames[0].n_particles, 2);
1439    }
1440
1441    // ── ParticleTrajectory ────────────────────────────────────────────────────
1442
1443    #[test]
1444    fn test_trajectory_push_and_get() {
1445        let mut traj = ParticleTrajectory::new();
1446        traj.push(make_frame(3));
1447        assert_eq!(traj.n_frames(), 1);
1448        assert_eq!(traj.n_particles(), 3);
1449        assert!(traj.get(0).is_some());
1450        assert!(traj.get(1).is_none());
1451    }
1452
1453    #[test]
1454    fn test_trajectory_slice() {
1455        let traj = make_traj(10, 2);
1456        let sub = traj.slice(2, 5);
1457        assert_eq!(sub.n_frames(), 3);
1458    }
1459
1460    #[test]
1461    fn test_trajectory_slice_beyond_end() {
1462        let traj = make_traj(5, 2);
1463        let sub = traj.slice(3, 100);
1464        assert_eq!(sub.n_frames(), 2);
1465    }
1466
1467    #[test]
1468    fn test_trajectory_subsample() {
1469        let traj = make_traj(10, 1);
1470        let sub = traj.subsample(2);
1471        assert_eq!(sub.n_frames(), 5);
1472    }
1473
1474    #[test]
1475    fn test_trajectory_subsample_zero_step() {
1476        let traj = make_traj(10, 1);
1477        let sub = traj.subsample(0);
1478        assert_eq!(sub.n_frames(), 0);
1479    }
1480
1481    // ── TrajectoryStats ───────────────────────────────────────────────────────
1482
1483    #[test]
1484    fn test_stats_msd_stationary() {
1485        let mut traj = ParticleTrajectory::new();
1486        for _ in 0..5 {
1487            traj.push(ParticleFrame::from_positions(0, vec![[0.0f32; 3]]));
1488        }
1489        let stats = TrajectoryStats::compute(&traj);
1490        for &m in &stats.msd {
1491            assert!(m.abs() < 1e-8, "msd = {m}");
1492        }
1493    }
1494
1495    #[test]
1496    fn test_stats_msd_drift() {
1497        let mut traj = ParticleTrajectory::new();
1498        for i in 0..5usize {
1499            traj.push(ParticleFrame::from_positions(
1500                i as u64,
1501                vec![[i as f32, 0.0, 0.0]],
1502            ));
1503        }
1504        let stats = TrajectoryStats::compute(&traj);
1505        // msd[1] should be 1, msd[2] = 4, etc.
1506        assert!(
1507            (stats.msd[1] - 1.0).abs() < 1e-6,
1508            "msd[1] = {}",
1509            stats.msd[1]
1510        );
1511    }
1512
1513    #[test]
1514    fn test_stats_rmsf_stationary() {
1515        let mut traj = ParticleTrajectory::new();
1516        for _ in 0..10 {
1517            traj.push(ParticleFrame::from_positions(0, vec![[1.0f32, 2.0, 3.0]]));
1518        }
1519        let stats = TrajectoryStats::compute(&traj);
1520        assert!(stats.rmsf[0].abs() < 1e-8);
1521    }
1522
1523    #[test]
1524    fn test_stats_com_drift_single_frame() {
1525        let mut traj = ParticleTrajectory::new();
1526        traj.push(ParticleFrame::from_positions(0, vec![[2.0f32, 4.0, 6.0]]));
1527        let stats = TrajectoryStats::compute(&traj);
1528        let com = stats.com_drift[0];
1529        assert!((com[0] - 2.0).abs() < 1e-6);
1530    }
1531
1532    #[test]
1533    fn test_stats_empty_trajectory() {
1534        let traj = ParticleTrajectory::new();
1535        let stats = TrajectoryStats::compute(&traj);
1536        assert!(stats.msd.is_empty());
1537        assert!(stats.rmsf.is_empty());
1538        assert!(stats.com_drift.is_empty());
1539    }
1540
1541    // ── BinaryFrameWriter / BinaryFrameReader ─────────────────────────────────
1542
1543    #[test]
1544    fn test_binary_frame_roundtrip_internal() {
1545        let frame = make_frame(4);
1546        let mut buf = Vec::new();
1547        write_binary_frame(&mut buf, &frame).unwrap();
1548        let mut cursor = Cursor::new(buf);
1549        let loaded = read_binary_frame(&mut cursor).unwrap();
1550        assert_eq!(loaded.n_particles, 4);
1551        for i in 0..4 {
1552            assert!((loaded.positions[i][0] - i as f32).abs() < 1e-6);
1553        }
1554    }
1555
1556    #[test]
1557    fn test_binary_frame_with_velocities() {
1558        let mut frame = make_frame(2);
1559        frame.velocities = Some(vec![[1.0f32, 0.0, 0.0], [0.0, 1.0, 0.0]]);
1560        let mut buf = Vec::new();
1561        write_binary_frame(&mut buf, &frame).unwrap();
1562        let mut cursor = Cursor::new(buf);
1563        let loaded = read_binary_frame(&mut cursor).unwrap();
1564        assert!(loaded.velocities.is_some());
1565        let vels = loaded.velocities.unwrap();
1566        assert!((vels[0][0] - 1.0).abs() < 1e-6);
1567    }
1568
1569    #[test]
1570    fn test_binary_frame_file_roundtrip() {
1571        let path = std::env::temp_dir().join("oxiphysics_pf_bin_rt.bin");
1572        let _ = std::fs::remove_file(&path);
1573        let w = BinaryFrameWriter::new(&path);
1574        for ts in 0..3u64 {
1575            let f = ParticleFrame::new(ts, vec![[ts as f32, 0.0, 0.0]], vec![0u8]);
1576            w.write_frame(&f).unwrap();
1577        }
1578        let frames = BinaryFrameReader::new(&path).read_all().unwrap();
1579        assert_eq!(frames.len(), 3);
1580        assert_eq!(frames[2].timestep, 2);
1581        let _ = std::fs::remove_file(&path);
1582    }
1583
1584    #[test]
1585    fn test_binary_frame_bad_magic() {
1586        let buf = vec![0xFFu8, 0xFF, 0xFF, 0xFF];
1587        let mut cursor = Cursor::new(buf);
1588        let result = read_binary_frame(&mut cursor);
1589        assert!(result.is_err());
1590    }
1591
1592    // ── GroWriter / GroReader ─────────────────────────────────────────────────
1593
1594    #[test]
1595    fn test_gro_write_creates_file() {
1596        let path = std::env::temp_dir().join("oxiphysics_pf_gro.gro");
1597        let frame = make_frame(3);
1598        GroWriter::new(&path)
1599            .write_frame(&frame, [10.0; 3])
1600            .unwrap();
1601        assert!(path.exists());
1602        let _ = std::fs::remove_file(&path);
1603    }
1604
1605    #[test]
1606    fn test_gro_write_contains_box() {
1607        let path = std::env::temp_dir().join("oxiphysics_pf_gro_box.gro");
1608        let frame = make_frame(1);
1609        GroWriter::new(&path)
1610            .write_frame(&frame, [5.0f32, 5.0, 5.0])
1611            .unwrap();
1612        let content = std::fs::read_to_string(&path).unwrap();
1613        assert!(content.contains("5.00000"));
1614        let _ = std::fs::remove_file(&path);
1615    }
1616
1617    #[test]
1618    fn test_gro_roundtrip_positions() {
1619        let path = std::env::temp_dir().join("oxiphysics_pf_gro_rt.gro");
1620        let pos = vec![[1.5f32, 2.5, 3.5], [4.5, 5.5, 6.5]];
1621        let frame = ParticleFrame::new(0, pos.clone(), vec![0u8, 0]);
1622        GroWriter::new(&path)
1623            .write_frame(&frame, [10.0; 3])
1624            .unwrap();
1625        let loaded = GroReader::new(&path).read_frame().unwrap();
1626        assert_eq!(loaded.n_particles, 2);
1627        // GRO uses %.3f precision for coordinates
1628        assert!((loaded.positions[0][0] - 1.5).abs() < 0.01);
1629        let _ = std::fs::remove_file(&path);
1630    }
1631
1632    // ── DcdWriter / DcdReader round-trip tests ────────────────────────────────
1633
1634    #[test]
1635    fn test_dcd_writer_creates_file() {
1636        let path = std::env::temp_dir().join("oxiphysics_pf_dcd_create.dcd");
1637        let _ = std::fs::remove_file(&path);
1638        let mut w = DcdWriter::new(&path, 3, 2.0e-3_f32).unwrap();
1639        let frame = make_frame(3);
1640        w.write_frame(&frame).unwrap();
1641        w.finalize().unwrap();
1642        assert!(path.exists());
1643        let _ = std::fs::remove_file(&path);
1644    }
1645
1646    #[test]
1647    fn test_dcd_roundtrip_single_frame() {
1648        let path = std::env::temp_dir().join("oxiphysics_pf_dcd_rt1.dcd");
1649        let _ = std::fs::remove_file(&path);
1650        let frame = make_frame(4);
1651        let mut w = DcdWriter::new(&path, 4, 1.0e-3_f32).unwrap();
1652        w.write_frame(&frame).unwrap();
1653        w.finalize().unwrap();
1654
1655        let (header, frames) = DcdReader::new(&path).read_all().unwrap();
1656        assert_eq!(header.n_frames, 1);
1657        assert_eq!(header.n_atoms, 4);
1658        assert_eq!(frames.len(), 1);
1659        for i in 0..4 {
1660            assert!(
1661                (frames[0].positions[i][0] - i as f32).abs() < 1e-5,
1662                "position mismatch at atom {i}"
1663            );
1664        }
1665        let _ = std::fs::remove_file(&path);
1666    }
1667
1668    #[test]
1669    fn test_dcd_roundtrip_multi_frame() {
1670        let path = std::env::temp_dir().join("oxiphysics_pf_dcd_rt3.dcd");
1671        let _ = std::fs::remove_file(&path);
1672        let n_atoms = 5u32;
1673        let mut w = DcdWriter::new(&path, n_atoms, 2.0e-3_f32).unwrap();
1674        for ts in 0..3u64 {
1675            let pos: Vec<[f32; 3]> = (0..n_atoms as usize)
1676                .map(|i| [ts as f32 + i as f32, 0.0, 0.0])
1677                .collect();
1678            let types = vec![0u8; n_atoms as usize];
1679            let frame = ParticleFrame::new(ts, pos, types);
1680            w.write_frame(&frame).unwrap();
1681        }
1682        w.finalize().unwrap();
1683
1684        let (header, frames) = DcdReader::new(&path).read_all().unwrap();
1685        assert_eq!(header.n_frames, 3);
1686        assert_eq!(frames.len(), 3);
1687        // Verify frame 2, atom 0: x == 2.0
1688        assert!((frames[2].positions[0][0] - 2.0).abs() < 1e-5);
1689        let _ = std::fs::remove_file(&path);
1690    }
1691}
oxiphysics_io/particle_formats.rs

oxiphysics_io/
particle_formats.rs
