1#![allow(clippy::needless_range_loop)]
2#![allow(dead_code)]
18
19use std::fs::File;
20use std::io::{self, BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write};
21
22#[derive(Debug, Clone, Default)]
28pub struct ParticleFrame {
29 pub timestep: u64,
31 pub n_particles: usize,
33 pub positions: Vec<[f32; 3]>,
35 pub velocities: Option<Vec<[f32; 3]>>,
37 pub types: Vec<u8>,
39}
40
41impl ParticleFrame {
42 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 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 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 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#[derive(Debug, Clone)]
115pub struct XyzWriter {
116 pub path: std::path::PathBuf,
118}
119
120impl XyzWriter {
121 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
123 Self { path: path.into() }
124 }
125
126 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 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 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
155fn 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#[derive(Debug, Clone)]
172pub struct XyzReader {
173 pub path: std::path::PathBuf,
175}
176
177impl XyzReader {
178 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
180 Self { path: path.into() }
181 }
182
183 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 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
206fn parse_xyz_frame<R: BufRead>(r: &mut R) -> io::Result<ParticleFrame> {
208 let mut line = String::new();
209
210 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 r.read_line(&mut line)?;
223 let timestep = extract_timestep_from_comment(&line);
224 line.clear();
225
226 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
262fn 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#[derive(Debug, Clone)]
280pub struct LammpsDumpWriter {
281 pub path: std::path::PathBuf,
283 pub write_velocities: bool,
285}
286
287impl LammpsDumpWriter {
288 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
290 Self {
291 path: path.into(),
292 write_velocities: true,
293 }
294 }
295
296 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
314fn 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#[derive(Debug, Clone)]
366pub struct LammpsDumpReader {
367 pub path: std::path::PathBuf,
369}
370
371impl LammpsDumpReader {
372 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
374 Self { path: path.into() }
375 }
376
377 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
385fn 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 {
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; 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; i += 4; let header = lines.get(i).cloned().unwrap_or_default();
417 let cols: Vec<&str> = header.split_whitespace().collect();
418 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 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; 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
495fn col_idx(cols: &[&str], name: &str) -> Option<usize> {
498 cols.iter().position(|&c| c == name)
499}
500
501#[derive(Debug, Clone, Default)]
507pub struct ParticleTrajectory {
508 pub frames: Vec<ParticleFrame>,
510}
511
512impl ParticleTrajectory {
513 pub fn new() -> Self {
515 Self { frames: Vec::new() }
516 }
517
518 pub fn n_frames(&self) -> usize {
520 self.frames.len()
521 }
522
523 pub fn n_particles(&self) -> usize {
525 self.frames.first().map(|f| f.n_particles).unwrap_or(0)
526 }
527
528 pub fn get(&self, i: usize) -> Option<&ParticleFrame> {
530 self.frames.get(i)
531 }
532
533 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 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 pub fn push(&mut self, frame: ParticleFrame) {
553 self.frames.push(frame);
554 }
555}
556
557#[derive(Debug, Clone, Default)]
563pub struct TrajectoryStats {
564 pub msd: Vec<f64>,
566 pub rmsf: Vec<f64>,
568 pub com_drift: Vec<[f64; 3]>,
570}
571
572impl TrajectoryStats {
573 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
583fn 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
613fn 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 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 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
648fn 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
669const BINARY_MAGIC: u32 = 0x4F58_4950; #[derive(Debug, Clone)]
678pub struct BinaryFrameWriter {
679 pub path: std::path::PathBuf,
681}
682
683impl BinaryFrameWriter {
684 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
686 Self { path: path.into() }
687 }
688
689 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
700fn 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 for &t in &frame.types {
712 w.write_all(&[t])?;
713 }
714 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#[derive(Debug, Clone)]
729pub struct BinaryFrameReader {
730 pub path: std::path::PathBuf,
732}
733
734impl BinaryFrameReader {
735 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
737 Self { path: path.into() }
738 }
739
740 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
756fn 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
799fn 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#[derive(Debug, Clone)]
827pub struct GroWriter {
828 pub path: std::path::PathBuf,
830}
831
832impl GroWriter {
833 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
835 Self { path: path.into() }
836 }
837
838 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 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#[derive(Debug, Clone)]
871pub struct GroReader {
872 pub path: std::path::PathBuf,
874}
875
876impl GroReader {
877 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
879 Self { path: path.into() }
880 }
881
882 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
890fn parse_gro_frame<R: BufRead>(r: &mut R) -> io::Result<ParticleFrame> {
892 let mut line = String::new();
893 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 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 let parts: Vec<&str> = line.split_whitespace().collect();
913 if parts.len() < 5 {
915 return Err(io::Error::new(io::ErrorKind::InvalidData, "short GRO line"));
916 }
917 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 Ok(ParticleFrame {
936 timestep: 0,
937 n_particles: n,
938 positions,
939 velocities: None,
940 types,
941 })
942}
943
944#[derive(Debug, Clone, Default)]
950pub struct DcdHeader {
951 pub n_frames: u32,
953 pub first_step: u32,
955 pub step_interval: u32,
957 pub n_atoms: u32,
959 pub titles: Vec<String>,
961}
962
963#[derive(Debug)]
967pub struct DcdReader {
968 pub path: std::path::PathBuf,
970}
971
972impl DcdReader {
973 pub fn new(path: impl Into<std::path::PathBuf>) -> Self {
975 Self { path: path.into() }
976 }
977
978 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 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
1003fn parse_dcd_header<R: Read + Seek>(r: &mut R) -> io::Result<DcdHeader> {
1005 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 let n_frames = read_u32_le(r)?;
1017 let first_step = read_u32_le(r)?;
1018 let step_interval = read_u32_le(r)?;
1019 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)?; 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 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
1053fn 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#[cfg(test)]
1088mod tests {
1089 use super::*;
1090 use std::io::Cursor;
1091
1092 fn make_frame(n: usize) -> ParticleFrame {
1095 let positions: Vec<[f32; 3]> = (0..n).map(|i| [i as f32, 0.0, 0.0]).collect();
1096 let types = vec![1u8; n];
1097 ParticleFrame::new(0, positions, types)
1098 }
1099
1100 fn make_traj(n_frames: usize, n_particles: usize) -> ParticleTrajectory {
1101 let mut traj = ParticleTrajectory::new();
1102 for t in 0..n_frames {
1103 let positions: Vec<[f32; 3]> = (0..n_particles).map(|_| [t as f32, 0.0, 0.0]).collect();
1104 traj.push(ParticleFrame::new(
1105 t as u64,
1106 positions,
1107 vec![0u8; n_particles],
1108 ));
1109 }
1110 traj
1111 }
1112
1113 #[test]
1116 fn test_frame_new_sets_n_particles() {
1117 let f = make_frame(5);
1118 assert_eq!(f.n_particles, 5);
1119 }
1120
1121 #[test]
1122 fn test_frame_from_positions() {
1123 let pos = vec![[1.0f32, 0.0, 0.0], [2.0, 0.0, 0.0]];
1124 let f = ParticleFrame::from_positions(10, pos);
1125 assert_eq!(f.n_particles, 2);
1126 assert_eq!(f.timestep, 10);
1127 assert_eq!(f.types, vec![0u8, 0]);
1128 }
1129
1130 #[test]
1131 fn test_frame_center_of_mass_empty() {
1132 let f = make_frame(0);
1133 assert_eq!(f.center_of_mass(), [0.0; 3]);
1134 }
1135
1136 #[test]
1137 fn test_frame_center_of_mass_single() {
1138 let f = ParticleFrame::from_positions(0, vec![[2.0f32, 4.0, 6.0]]);
1139 let com = f.center_of_mass();
1140 assert!((com[0] - 2.0).abs() < 1e-6);
1141 assert!((com[1] - 4.0).abs() < 1e-6);
1142 assert!((com[2] - 6.0).abs() < 1e-6);
1143 }
1144
1145 #[test]
1146 fn test_frame_center_of_mass_symmetric() {
1147 let pos = vec![[-1.0f32, 0.0, 0.0], [1.0, 0.0, 0.0]];
1148 let f = ParticleFrame::from_positions(0, pos);
1149 let com = f.center_of_mass();
1150 assert!(com[0].abs() < 1e-6);
1151 }
1152
1153 #[test]
1154 fn test_frame_bounding_box_empty() {
1155 let f = make_frame(0);
1156 let (lo, hi) = f.bounding_box();
1157 assert_eq!(lo, [0.0; 3]);
1158 assert_eq!(hi, [0.0; 3]);
1159 }
1160
1161 #[test]
1162 fn test_frame_bounding_box_basic() {
1163 let pos = vec![[1.0f32, -1.0, 0.0], [-1.0, 2.0, 3.0]];
1164 let f = ParticleFrame::from_positions(0, pos);
1165 let (lo, hi) = f.bounding_box();
1166 assert!((lo[0] - (-1.0)).abs() < 1e-6);
1167 assert!((hi[1] - 2.0).abs() < 1e-6);
1168 assert!((hi[2] - 3.0).abs() < 1e-6);
1169 }
1170
1171 #[test]
1174 fn test_xyz_write_frame_internal() {
1175 let frame = make_frame(3);
1176 let mut buf = Vec::new();
1177 write_xyz_frame_to(&mut buf, &frame, "test").unwrap();
1178 let s = String::from_utf8(buf).unwrap();
1179 assert!(s.starts_with("3\n"));
1180 assert!(s.contains("test"));
1181 }
1182
1183 #[test]
1184 fn test_xyz_roundtrip_file() {
1185 let path = std::env::temp_dir().join("oxiphysics_pf_xyz_rt.xyz");
1186 let _ = std::fs::remove_file(&path);
1187 let pos = vec![[1.0f32, 2.0, 3.0], [4.0, 5.0, 6.0]];
1188 let frame = ParticleFrame::new(42, pos.clone(), vec![1u8, 2]);
1189 XyzWriter::new(&path).write_frame(&frame, "test").unwrap();
1190 let loaded = XyzReader::new(&path).read_frame().unwrap();
1191 assert_eq!(loaded.n_particles, 2);
1192 assert_eq!(loaded.timestep, 42);
1193 assert!((loaded.positions[0][0] - 1.0).abs() < 1e-4);
1194 assert!((loaded.positions[1][2] - 6.0).abs() < 1e-4);
1195 let _ = std::fs::remove_file(&path);
1196 }
1197
1198 #[test]
1199 fn test_xyz_append_and_read_all() {
1200 let path = std::env::temp_dir().join("oxiphysics_pf_xyz_all.xyz");
1201 let _ = std::fs::remove_file(&path);
1202 let w = XyzWriter::new(&path);
1203 for ts in 0..4u64 {
1204 let f = ParticleFrame::new(ts, vec![[ts as f32, 0.0, 0.0]], vec![0u8]);
1205 w.append_frame(&f, "").unwrap();
1206 }
1207 let frames = XyzReader::new(&path).read_all().unwrap();
1208 assert_eq!(frames.len(), 4);
1209 let _ = std::fs::remove_file(&path);
1210 }
1211
1212 #[test]
1213 fn test_xyz_write_trajectory() {
1214 let path = std::env::temp_dir().join("oxiphysics_pf_xyz_traj.xyz");
1215 let _ = std::fs::remove_file(&path);
1216 let mut traj = ParticleTrajectory::new();
1217 for ts in 0..3u64 {
1218 traj.push(ParticleFrame::new(ts, vec![[0.0f32; 3]], vec![0u8]));
1219 }
1220 XyzWriter::new(&path).write_trajectory(&traj).unwrap();
1221 let frames = XyzReader::new(&path).read_all().unwrap();
1222 assert_eq!(frames.len(), 3);
1223 let _ = std::fs::remove_file(&path);
1224 }
1225
1226 #[test]
1227 fn test_extract_timestep_from_comment() {
1228 assert_eq!(extract_timestep_from_comment("Timestep=100 other"), 100);
1229 assert_eq!(extract_timestep_from_comment("no timestep here"), 0);
1230 }
1231
1232 #[test]
1235 fn test_lammps_dump_write_internal() {
1236 let frame = make_frame(2);
1237 let mut buf = Vec::new();
1238 write_lammps_dump_frame(&mut buf, &frame, [0.0; 3], [10.0; 3], false).unwrap();
1239 let s = String::from_utf8(buf).unwrap();
1240 assert!(s.contains("ITEM: TIMESTEP"));
1241 assert!(s.contains("ITEM: ATOMS id type x y z"));
1242 }
1243
1244 #[test]
1245 fn test_lammps_dump_write_with_velocities() {
1246 let mut frame = make_frame(2);
1247 frame.velocities = Some(vec![[1.0f32, 0.0, 0.0], [0.0, 1.0, 0.0]]);
1248 let mut buf = Vec::new();
1249 write_lammps_dump_frame(&mut buf, &frame, [0.0; 3], [10.0; 3], true).unwrap();
1250 let s = String::from_utf8(buf).unwrap();
1251 assert!(s.contains("vx vy vz"));
1252 }
1253
1254 #[test]
1255 fn test_lammps_dump_roundtrip_file() {
1256 let path = std::env::temp_dir().join("oxiphysics_pf_lmp.lammpstrj");
1257 let _ = std::fs::remove_file(&path);
1258 let w = LammpsDumpWriter::new(&path);
1259 for ts in 0..3u64 {
1260 let f = ParticleFrame::new(ts, vec![[ts as f32, 0.0, 0.0]], vec![1u8]);
1261 w.write_frame(&f, [0.0; 3], [10.0; 3]).unwrap();
1262 }
1263 let frames = LammpsDumpReader::new(&path).read_all().unwrap();
1264 assert_eq!(frames.len(), 3);
1265 assert_eq!(frames[1].timestep, 1);
1266 let _ = std::fs::remove_file(&path);
1267 }
1268
1269 #[test]
1270 fn test_lammps_dump_parse_inline() {
1271 let dump = "\
1272ITEM: TIMESTEP\n\
12735\n\
1274ITEM: NUMBER OF ATOMS\n\
12752\n\
1276ITEM: BOX BOUNDS pp pp pp\n\
12770.0 10.0\n\
12780.0 10.0\n\
12790.0 10.0\n\
1280ITEM: ATOMS id type x y z\n\
12811 1 1.0 2.0 3.0\n\
12822 2 4.0 5.0 6.0\n";
1283 let r = BufReader::new(Cursor::new(dump));
1284 let frames = parse_lammps_dump_all(r).unwrap();
1285 assert_eq!(frames.len(), 1);
1286 assert_eq!(frames[0].timestep, 5);
1287 assert_eq!(frames[0].n_particles, 2);
1288 }
1289
1290 #[test]
1293 fn test_trajectory_push_and_get() {
1294 let mut traj = ParticleTrajectory::new();
1295 traj.push(make_frame(3));
1296 assert_eq!(traj.n_frames(), 1);
1297 assert_eq!(traj.n_particles(), 3);
1298 assert!(traj.get(0).is_some());
1299 assert!(traj.get(1).is_none());
1300 }
1301
1302 #[test]
1303 fn test_trajectory_slice() {
1304 let traj = make_traj(10, 2);
1305 let sub = traj.slice(2, 5);
1306 assert_eq!(sub.n_frames(), 3);
1307 }
1308
1309 #[test]
1310 fn test_trajectory_slice_beyond_end() {
1311 let traj = make_traj(5, 2);
1312 let sub = traj.slice(3, 100);
1313 assert_eq!(sub.n_frames(), 2);
1314 }
1315
1316 #[test]
1317 fn test_trajectory_subsample() {
1318 let traj = make_traj(10, 1);
1319 let sub = traj.subsample(2);
1320 assert_eq!(sub.n_frames(), 5);
1321 }
1322
1323 #[test]
1324 fn test_trajectory_subsample_zero_step() {
1325 let traj = make_traj(10, 1);
1326 let sub = traj.subsample(0);
1327 assert_eq!(sub.n_frames(), 0);
1328 }
1329
1330 #[test]
1333 fn test_stats_msd_stationary() {
1334 let mut traj = ParticleTrajectory::new();
1335 for _ in 0..5 {
1336 traj.push(ParticleFrame::from_positions(0, vec![[0.0f32; 3]]));
1337 }
1338 let stats = TrajectoryStats::compute(&traj);
1339 for &m in &stats.msd {
1340 assert!(m.abs() < 1e-8, "msd = {m}");
1341 }
1342 }
1343
1344 #[test]
1345 fn test_stats_msd_drift() {
1346 let mut traj = ParticleTrajectory::new();
1347 for i in 0..5usize {
1348 traj.push(ParticleFrame::from_positions(
1349 i as u64,
1350 vec![[i as f32, 0.0, 0.0]],
1351 ));
1352 }
1353 let stats = TrajectoryStats::compute(&traj);
1354 assert!(
1356 (stats.msd[1] - 1.0).abs() < 1e-6,
1357 "msd[1] = {}",
1358 stats.msd[1]
1359 );
1360 }
1361
1362 #[test]
1363 fn test_stats_rmsf_stationary() {
1364 let mut traj = ParticleTrajectory::new();
1365 for _ in 0..10 {
1366 traj.push(ParticleFrame::from_positions(0, vec![[1.0f32, 2.0, 3.0]]));
1367 }
1368 let stats = TrajectoryStats::compute(&traj);
1369 assert!(stats.rmsf[0].abs() < 1e-8);
1370 }
1371
1372 #[test]
1373 fn test_stats_com_drift_single_frame() {
1374 let mut traj = ParticleTrajectory::new();
1375 traj.push(ParticleFrame::from_positions(0, vec![[2.0f32, 4.0, 6.0]]));
1376 let stats = TrajectoryStats::compute(&traj);
1377 let com = stats.com_drift[0];
1378 assert!((com[0] - 2.0).abs() < 1e-6);
1379 }
1380
1381 #[test]
1382 fn test_stats_empty_trajectory() {
1383 let traj = ParticleTrajectory::new();
1384 let stats = TrajectoryStats::compute(&traj);
1385 assert!(stats.msd.is_empty());
1386 assert!(stats.rmsf.is_empty());
1387 assert!(stats.com_drift.is_empty());
1388 }
1389
1390 #[test]
1393 fn test_binary_frame_roundtrip_internal() {
1394 let frame = make_frame(4);
1395 let mut buf = Vec::new();
1396 write_binary_frame(&mut buf, &frame).unwrap();
1397 let mut cursor = Cursor::new(buf);
1398 let loaded = read_binary_frame(&mut cursor).unwrap();
1399 assert_eq!(loaded.n_particles, 4);
1400 for i in 0..4 {
1401 assert!((loaded.positions[i][0] - i as f32).abs() < 1e-6);
1402 }
1403 }
1404
1405 #[test]
1406 fn test_binary_frame_with_velocities() {
1407 let mut frame = make_frame(2);
1408 frame.velocities = Some(vec![[1.0f32, 0.0, 0.0], [0.0, 1.0, 0.0]]);
1409 let mut buf = Vec::new();
1410 write_binary_frame(&mut buf, &frame).unwrap();
1411 let mut cursor = Cursor::new(buf);
1412 let loaded = read_binary_frame(&mut cursor).unwrap();
1413 assert!(loaded.velocities.is_some());
1414 let vels = loaded.velocities.unwrap();
1415 assert!((vels[0][0] - 1.0).abs() < 1e-6);
1416 }
1417
1418 #[test]
1419 fn test_binary_frame_file_roundtrip() {
1420 let path = std::env::temp_dir().join("oxiphysics_pf_bin_rt.bin");
1421 let _ = std::fs::remove_file(&path);
1422 let w = BinaryFrameWriter::new(&path);
1423 for ts in 0..3u64 {
1424 let f = ParticleFrame::new(ts, vec![[ts as f32, 0.0, 0.0]], vec![0u8]);
1425 w.write_frame(&f).unwrap();
1426 }
1427 let frames = BinaryFrameReader::new(&path).read_all().unwrap();
1428 assert_eq!(frames.len(), 3);
1429 assert_eq!(frames[2].timestep, 2);
1430 let _ = std::fs::remove_file(&path);
1431 }
1432
1433 #[test]
1434 fn test_binary_frame_bad_magic() {
1435 let buf = vec![0xFFu8, 0xFF, 0xFF, 0xFF];
1436 let mut cursor = Cursor::new(buf);
1437 let result = read_binary_frame(&mut cursor);
1438 assert!(result.is_err());
1439 }
1440
1441 #[test]
1444 fn test_gro_write_creates_file() {
1445 let path = std::env::temp_dir().join("oxiphysics_pf_gro.gro");
1446 let frame = make_frame(3);
1447 GroWriter::new(&path)
1448 .write_frame(&frame, [10.0; 3])
1449 .unwrap();
1450 assert!(path.exists());
1451 let _ = std::fs::remove_file(&path);
1452 }
1453
1454 #[test]
1455 fn test_gro_write_contains_box() {
1456 let path = std::env::temp_dir().join("oxiphysics_pf_gro_box.gro");
1457 let frame = make_frame(1);
1458 GroWriter::new(&path)
1459 .write_frame(&frame, [5.0f32, 5.0, 5.0])
1460 .unwrap();
1461 let content = std::fs::read_to_string(&path).unwrap();
1462 assert!(content.contains("5.00000"));
1463 let _ = std::fs::remove_file(&path);
1464 }
1465
1466 #[test]
1467 fn test_gro_roundtrip_positions() {
1468 let path = std::env::temp_dir().join("oxiphysics_pf_gro_rt.gro");
1469 let pos = vec![[1.5f32, 2.5, 3.5], [4.5, 5.5, 6.5]];
1470 let frame = ParticleFrame::new(0, pos.clone(), vec![0u8, 0]);
1471 GroWriter::new(&path)
1472 .write_frame(&frame, [10.0; 3])
1473 .unwrap();
1474 let loaded = GroReader::new(&path).read_frame().unwrap();
1475 assert_eq!(loaded.n_particles, 2);
1476 assert!((loaded.positions[0][0] - 1.5).abs() < 0.01);
1478 let _ = std::fs::remove_file(&path);
1479 }
1480}