Add PDBWriter.dump_single

Author: tovrstra

CHANGELOG.md

@@ -7,10 +7,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add a `stride` option to the trajectory writers.
+- Add method `dump_single` method to `PDBWriter` to write one-off file with a single snapshot.
+
 ### Changed
 
 - By default, run only 100 optimization steps in `build_random_cell`.
-- Add a `stride` option to the trajectory writers.
 
 ## [0.0.0] - 2024-10-06
 

src/tinyff/trajectory.py

@@ -27,6 +27,7 @@
 import os
 from functools import partial
 from glob import glob
+from typing import TextIO
 
 import attrs
 import numpy as np
@@ -83,19 +84,27 @@ def dump(self, atpos: ArrayLike, cell_lengths: ArrayLike):
 
     def dump_each(self, atpos: ArrayLike, cell_lengths: ArrayLike):
         """Write a snapshot to the PDB file without considering `self.stride`."""
+        with open(self.path_pdb, "a") as fh:
+            self._dump_low(fh, atpos, cell_lengths)
+
+    def dump_single(self, path_pdb: str, atpos: ArrayLike, cell_lengths: ArrayLike):
+        """Dump a single snapshot with the same settings to a different file."""
+        with open(path_pdb, "w") as fh:
+            self._dump_low(fh, atpos, cell_lengths)
+
+    def _dump_low(self, fh: TextIO, atpos: ArrayLike, cell_lengths: ArrayLike):
         atpos = parse_atpos(atpos, len(self.atnums))
         cell_lengths = parse_cell_lengths(cell_lengths)
-        with open(self.path_pdb, "a") as fh:
-            a, b, c = cell_lengths * self.to_angstrom
-            print(f"CRYST1{a:9.3f}{b:9.3f}{c:9.3f}  90.00  90.00  90.00 P 1           1", file=fh)
-            for i, (x, y, z) in enumerate(atpos * self.to_angstrom):
-                symbol = SYMBOLS[self.atnums[i] - 1]
-                print(
-                    f"HETATM{i+1:5d} {symbol:2s}   ATM     1    {x:8.3f}{y:8.3f}{z:8.3f}"
-                    f"  1.00  1.00          {symbol:2s}",
-                    file=fh,
-                )
-            print("END", file=fh)
+        a, b, c = cell_lengths * self.to_angstrom
+        print(f"CRYST1{a:9.3f}{b:9.3f}{c:9.3f}  90.00  90.00  90.00 P 1           1", file=fh)
+        for i, (x, y, z) in enumerate(atpos * self.to_angstrom):
+            symbol = SYMBOLS[self.atnums[i] - 1]
+            print(
+                f"HETATM{i+1:5d} {symbol:2s}   ATM     1    {x:8.3f}{y:8.3f}{z:8.3f}"
+                f"  1.00  1.00          {symbol:2s}",
+                file=fh,
+            )
+        print("END", file=fh)
 
 
 @attrs.define

tests/test_trajectory.py

@@ -79,6 +79,26 @@ def test_pdb_writer_stride(tmpdir):
         assert fh.read().count("CRYST1") == 3
 
 
+PDB_REF_SINGLE = """\
+CRYST1   10.000   10.000   10.000  90.00  90.00  90.00 P 1           1
+HETATM    1 Ar   ATM     1       3.000   4.000   5.000  1.00  1.00          Ar
+HETATM    2 Ar   ATM     1       6.000   7.000   3.142  1.00  1.00          Ar
+END
+"""
+
+
+def test_pdb_writer_single(tmpdir):
+    cell_length = 10.0
+    atpos = np.array([[3.0, 4.0, 5.0], [6.0, 7.0, 3.1415]])
+    path_pdb = os.path.join(tmpdir, "test.pdb")
+    path_other = os.path.join(tmpdir, "other.pdb")
+    pdb_writer = PDBWriter(path_pdb, to_angstrom=1.0, atnums=[18] * len(atpos), stride=5)
+    pdb_writer.dump_single(path_other, atpos, cell_length)
+
+    with open(path_other) as fh:
+        assert fh.read() == PDB_REF_SINGLE
+
+
 def test_npy_traj(tmpdir):
     traj_atpos = np.array([[[1.2, 1.3], [4.9, 3.1]], [[0.7, 8.1], [-7.9, 0.5]]])
     traj_pressure = np.array([5.0, 4.0])