patching large time numbers for trajectory writers to not overflow the low precision they use

Author: stefdoerr

moleculekit/writers.py

@@ -31,6 +31,37 @@
 # fmt: on
 
 
+def _format_large_time(mol, ext):
+    step = mol.step
+    time = mol.time
+    nframes = mol.numFrames
+
+    if len(step) == 1:
+        trajfreq = step[0]
+    else:
+        trajfreq = step[1] - step[0]
+
+    if step[-1] != step[-1].astype(np.uint32):
+        logger.warning(
+            f"Molecule.step contains values too large to be written to a {ext} file. They will be renumbered starting from 1."
+        )
+        step = np.arange(1, nframes + 1, dtype=np.uint32)
+    else:
+        step = step.astype(np.uint32)
+
+    time = time / 1e3  # convert from fs to ps
+    if time[-1] != time[-1].astype(np.float32):
+        logger.warning(
+            f"Molecule.time contains values too large to be written to a {ext} file. They will be renumbered starting from 0."
+        )
+        if trajfreq == 0:
+            raise AssertionError("The trajectory step should not be 0")
+        timestep = (mol.fstep / trajfreq) / 1e-6
+        time = (mol.time - ((mol.step[0] - trajfreq) * timestep)) / 1e3
+
+    return step, time
+
+
 def _format_pdb_name(name, resname, element=None):
     name = name[:4]
     first_col = f"{name:<4}"
@@ -263,9 +294,10 @@ def XTCwrite(mol, filename):
     if os.path.isfile(filename):
         os.unlink(filename)
 
+    step, time = _format_large_time(mol, "XTC")
+
     box = box.astype(np.float32) * 0.1
-    step = step.astype(np.int32)
-    time = time.astype(np.float32) / 1e3  # Convert from fs to ps
+    time = time.astype(np.float32)
     coords = coords.astype(np.float32) * 0.1  # Convert from A to nm
     if not box.flags["C_CONTIGUOUS"]:
         box = np.ascontiguousarray(box)
@@ -671,6 +703,16 @@ def DCDwrite(mol, filename):
         nsavc = int(mol.step[1] - mol.step[0])
         fstep = mol.fstep * 1000  # ns to ps
         delta = fstep / nsavc / 0.04888821  # Conversion factor found in OpenMM
+        if mol.step[0] != mol.step[0].astype(np.int32):
+            logger.warning(
+                "Molecule.step contains values too large to be written to DCD file. They will be renumbered starting from 1."
+            )
+            istart //= nsavc
+            delta *= nsavc
+            nsavc = 1
+        # If it's still too large just start from 1
+        if istart != np.array(istart).astype(np.int32):
+            istart = 1
     except Exception:
         istart = 0
         nsavc = 1
@@ -699,8 +741,7 @@ def TRRwrite(mol, filename):
     from moleculekit.trr import TRRTrajectoryFile
 
     xyz = np.transpose(mol.coords, (2, 0, 1)) / 10  # Convert Angstrom to nm
-    time = mol.time / 1000  # Convert fs to ps
-    step = mol.step
+    step, time = _format_large_time(mol, "TRR")
     boxvectors = np.transpose(mol.boxvectors, (2, 0, 1)) / 10  # Angstrom to nm
     with TRRTrajectoryFile(filename, "w") as fh:
         fh.write(xyz, time=time, step=step, box=boxvectors, lambd=None)
@@ -728,8 +769,10 @@ def NETCDFwrite(mol, filename):
     )
     n_frames, n_atoms = coordinates.shape[0], coordinates.shape[1]
 
+    step, time = _format_large_time(mol, "NETCDF")
+
     time = ensure_type(
-        mol.time / 1000,  # Convert from fs to ps
+        time,  # In ps
         np.float32,
         1,
         "time",
@@ -739,7 +782,7 @@ def NETCDFwrite(mol, filename):
         add_newaxis_on_deficient_ndim=True,
     )
     step = ensure_type(
-        mol.step,
+        step,
         np.int32,
         1,
         "step",
@@ -984,7 +1027,7 @@ def MDTRAJwrite(mol, filename):
             traj = Trajectory(
                 xyz=np.transpose(mol.coords, (2, 0, 1)) / 10,  # Ang to nm
                 topology=traj.topology,
-                time=time,
+                time=time.astype(np.float32),
                 unitcell_lengths=box,
                 unitcell_angles=boxangles,
             )

tests/test_writers.py

@@ -185,18 +185,41 @@ def _test_cif_mol2_atom_renaming():
         assert filelines == reflines, f"Failed comparison of {reffile2} {tmpfile}"
 
 
-@pytest.mark.parametrize("ext", ("netcdf", "trr", "binpos", "dcd", "xyz", "xyz.gz"))
-def _test_traj_writers(ext):
+@pytest.mark.parametrize(
+    "ext", ("xtc", "netcdf", "trr", "binpos", "dcd", "xyz", "xyz.gz")
+)
+@pytest.mark.parametrize("maxtime", [1e9, 1e15])
+def _test_traj_writers(ext, maxtime):
     from moleculekit.molecule import Molecule
     import tempfile
 
     mol = Molecule(os.path.join(curr_dir, "test_readers", "1N09", "structure.prmtop"))
     mol.read(os.path.join(curr_dir, "test_readers", "1N09", "output.dcd"))
+    # 1e9 fs = 1us. Test if the trajectories can write steps of 100ps over 1us trajectories
+    trajfreq = 25000
+    timestep = 4
+    timefreq = trajfreq * timestep
+    mol.time[:] = np.arange(
+        maxtime,
+        maxtime + mol.numFrames * timefreq,
+        timefreq,
+        dtype=Molecule._dtypes["time"],
+    )
+    mol.step[:] = np.arange(
+        maxtime / timestep,
+        maxtime / timestep + (mol.numFrames * trajfreq),
+        trajfreq,
+        dtype=Molecule._dtypes["step"],
+    )
 
     with tempfile.TemporaryDirectory() as tmpdir:
         mol.write(os.path.join(tmpdir, f"output.{ext}"))
         molc = Molecule(os.path.join(tmpdir, f"output.{ext}"))
 
+    if maxtime > 1e10:
+        mol.time -= (mol.step[0] - trajfreq) * timestep
+        mol.step[:] = range(1, mol.numFrames + 1)
+
     if ext == "binpos":
         assert np.allclose(mol.coords, molc.coords, atol=1e-6)
     elif ext in ("xyz", "xyz.gz"):
@@ -207,12 +230,14 @@ def _test_traj_writers(ext):
             fieldPrecision={"coords": 2e-5},
         )
     else:
+        coor_prec = 3e-6 if ext != "xtc" else 1e-2
+        assert abs(mol.fstep - molc.fstep) < 1e-2
         assert mol_equal(
             mol,
             molc,
             checkFields=Molecule._traj_fields,
             exceptFields=("fileloc"),
-            fieldPrecision={"coords": 3e-6, "box": 3e-6},
+            fieldPrecision={"coords": coor_prec, "box": 3e-6, "time": 1e-3},
         )