Wrap atoms in box in PDB output

Author: tovrstra

CHANGELOG.md

@@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - By default, run only 100 optimization steps in `build_random_cell`.
+- Wrap atoms back into the cell when writing PDB trajectory files, for nicer visual.
+
 
 ## [0.0.0] - 2024-10-06
 

src/tinyff/trajectory.py

@@ -93,8 +93,12 @@ def dump_single(self, path_pdb: str, atpos: ArrayLike, cell_lengths: ArrayLike):
             self._dump_low(fh, atpos, cell_lengths)
 
     def _dump_low(self, fh: TextIO, atpos: ArrayLike, cell_lengths: ArrayLike):
+        # Process arguments
         atpos = parse_atpos(atpos, len(self.atnums))
         cell_lengths = parse_cell_lengths(cell_lengths)
+        # Wrap atoms in cell for nicer visual
+        atpos -= np.floor(atpos / cell_lengths) * cell_lengths
+        # Actual writing
         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):

tests/test_trajectory.py

@@ -28,13 +28,13 @@
 PDB_REF_DATA = """\
 CRYST1   30.000   30.000   30.000  90.00  90.00  90.00 P 1           1
 HETATM    1 Ar   ATM     1       0.000  20.000   5.000  1.00  1.00          Ar
-HETATM    2 Ar   ATM     1      10.000  33.000  24.400  1.00  1.00          Ar
-HETATM    3 Ar   ATM     1     -10.000  28.000  17.000  1.00  1.00          Ar
+HETATM    2 Ar   ATM     1      10.000   3.000  24.400  1.00  1.00          Ar
+HETATM    3 Ar   ATM     1      20.000  28.000  17.000  1.00  1.00          Ar
 END
 CRYST1   30.000   30.000   30.000  90.00  90.00  90.00 P 1           1
 HETATM    1 Ar   ATM     1       0.100  20.300   0.000  1.00  1.00          Ar
-HETATM    2 Ar   ATM     1      12.000  30.500  32.100  1.00  1.00          Ar
-HETATM    3 Ar   ATM     1     -10.300  35.800  20.300  1.00  1.00          Ar
+HETATM    2 Ar   ATM     1      12.000   0.500   2.100  1.00  1.00          Ar
+HETATM    3 Ar   ATM     1      19.700   5.800  20.300  1.00  1.00          Ar
 END
 """