fix dcd, trr and netcdf writers

Author: stefdoerr

moleculekit/readers.py

@@ -2887,14 +2887,15 @@ def NETCDFread(
             "variables in the file were %s" % _handle.variables.keys()
         )
 
+    time = None
     step = None
+    if "step" in _handle.variables:
+        step = _handle.variables["step"][frame_slice]
     if "time" in _handle.variables:
         time = _handle.variables["time"][frame_slice]
-        if len(time) > 1:
+        if len(time) > 1 and step is None:
             timestep = time[1] - time[0]
             step = time / timestep
-    else:
-        time = None
 
     if "cell_lengths" in _handle.variables:
         cell_lengths = _handle.variables["cell_lengths"][frame_slice].T
@@ -2957,7 +2958,7 @@ def DCDread(filename, frame=None, topoloc=None, stride=None, atom_indices=None):
         istart, nsavc, delta = f.read_header()
         # Timestep conversion factor found in OpenMM
         delta = np.round(delta * 0.04888821, decimals=8)
-        steps = np.arange(istart, (nsavc * xyz.shape[0]) + 1, nsavc)
+        steps = np.arange(istart, (nsavc * xyz.shape[0]) + istart, nsavc)
         if stride is not None:
             steps = steps[::stride]
 
@@ -3004,8 +3005,7 @@ def TRRread(filename, frame=None, topoloc=None, stride=None, atom_indices=None):
     box = np.vstack((a_length, b_length, c_length)) * 10  # nm to Angstroms
     boxangles = np.vstack((alpha, beta, gamma))
 
-    step = None
-    if len(time) > 1:
+    if len(time) > 1 and (step is None or len(step) != len(time)):
         timestep = time[1] - time[0]
         step = time / timestep
     time *= 1000  # ps to fs

moleculekit/writers.py

@@ -738,6 +738,16 @@ def NETCDFwrite(mol, filename):
         warn_on_cast=False,
         add_newaxis_on_deficient_ndim=True,
     )
+    step = ensure_type(
+        mol.step,
+        np.int32,
+        1,
+        "step",
+        length=n_frames,
+        can_be_none=True,
+        warn_on_cast=False,
+        add_newaxis_on_deficient_ndim=True,
+    )
     cell_lengths = ensure_type(
         mol.box.T,
         np.float64,
@@ -832,6 +842,8 @@ def NETCDFwrite(mol, filename):
         # Define coordinates and snapshot times.
         frame_times = ncfile.createVariable("time", "f", ("frame",))
         setattr(frame_times, "units", "picosecond")
+        frame_steps = ncfile.createVariable("step", "i", ("frame",))
+        setattr(frame_steps, "units", "step")
 
     frame_coordinates = ncfile.createVariable(
         "coordinates", "f", ("frame", "atom", "spatial")
@@ -852,6 +864,8 @@ def NETCDFwrite(mol, filename):
         ncfile.variables["coordinates"][frame_slice, :, :] = coordinates
         if time is not None and set_time:
             ncfile.variables["time"][frame_slice] = time
+        if step is not None and set_time:
+            ncfile.variables["step"][frame_slice] = step
         if cell_lengths is not None:
             ncfile.variables["cell_lengths"][frame_slice, :] = cell_lengths
         if cell_angles is not None:
@@ -866,6 +880,8 @@ def NETCDFwrite(mol, filename):
     missing = None
     if time is None and "time" in ncfile.variables:
         missing = "time"
+    elif step is None and "step" in ncfile.variables:
+        missing = "step"
     elif cell_angles is None and "cell_angles" in ncfile.variables:
         missing = "cell_angles"
     elif cell_lengths is None and "cell_lengths" in ncfile.variables: