Repository: openmm/openmm — pull request #4875 (cpp)
Title: Attempting to fix XTC/DCD restarting

Fixes #4873 by writing directly the simulation.currentStep into the XTC and DCD files.

I'm making here some assumptions that the first step written in a file will be at `step=interval` but I think this is a valid assumption?

Also assuming that you would only append to an XTC/DCD file if you were resuming a previous simulation (i.e. via a checkpoint) and not creating a new Simulation from scratch which would set the currentStep to 0

Resolved issues:
title: Wrong step in DCD/XTCreporter when restarting simulations
body: Hi, I just noticed what I assume is a bug in the XTCreporter and DCDreporter class

https://github.com/openmm/openmm/blob/master/wrappers/python/openmm/app/xtcfile.py#L146

If you look at how the step is calculated

```py
step = (self._modelCount - 1) * self._interval + self._firstStep
```

When appending, the `_modelCount` variable is overwritten here: https://github.com/openmm/openmm/blob/master/wrappers/python/openmm/app/xtcfile.py#L57

And the `_firstStep` value is passed here: https://github.com/openmm/openmm/blob/master/wrappers/python/openmm/app/xtcreporter.py#L110 directly from the `simulation.currentStep`

This leads to the following behaviour if I add a `print("STEP", step, "TIME", time, self._modelCount, self._interval, self._firstStep, self._dt)` right before writing the XTC frame
```
STEP 50 TIME 0.2 1 50 50 0.004
STEP 100 TIME 0.4 2 50 50 0.004
STEP 150 TIME 0.6 3 50 50 0.004
STEP 200 TIME 0.8 4 50 50 0.004
STEP 250 TIME 1.0 5 50 50 0.004
RESTART
STEP 550 TIME 2.2 6 50 300 0.004
STEP 600 TIME 2.4 7 50 300 0.004
STEP 650 TIME 2.6 8 50 300 0.004
STEP 700 TIME 2.8 9 50 300 0.004
STEP 750 TIME 3.0 10 50 300 0.004
```
As you can see, after the restart, the `_firstStep` is 300 since that indeed is the currentStep of the simulation, but also the `_modelCount` is 6. This means that it calculates wrongly the step as 550 instead of 300 which it should.

I took a look at the DCDreporter and it does nearly the same, but with an offset of +50 compared to XTCreporter, so after the restart it calculates the step as 600 instead of 300.

## IMPORTANT INSTRUCTIONS

1. If you edit a python file, you MUST run 'make PythonInstall' inside the build directory
2. If you edit a c++ file, you MUST run 'make -j4 && make install" inside the build directory

Hints:

Fix patch:
diff --git a/wrappers/python/openmm/app/dcdfile.py b/wrappers/python/openmm/app/dcdfile.py
index 335cdc0b79..c886527da9 100644
--- a/wrappers/python/openmm/app/dcdfile.py
+++ b/wrappers/python/openmm/app/dcdfile.py
@@ -75,7 +75,6 @@ def __init__(self, file, topology, dt, firstStep=0, interval=1, append=False):
         self._topology = topology
         self._firstStep = firstStep
         self._interval = interval
-        self._modelCount = 0
         if is_quantity(dt):
             dt = dt.value_in_unit(picoseconds)
         dt /= 0.04888821
@@ -84,8 +83,6 @@ def __init__(self, file, topology, dt, firstStep=0, interval=1, append=False):
         if topology.getUnitCellDimensions() is not None:
             boxFlag = 1
         if append:
-            file.seek(8, os.SEEK_SET)
-            self._modelCount = struct.unpack('<i', file.read(4))[0]
             file.seek(268, os.SEEK_SET)
             numAtoms = struct.unpack('<i', file.read(4))[0]
             if numAtoms != topology.getNumAtoms():
@@ -98,7 +95,7 @@ def __init__(self, file, topology, dt, firstStep=0, interval=1, append=False):
             header += struct.pack('<4i', 164, 4, topology.getNumAtoms(), 4)
             file.write(header)
 
-    def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None):
+    def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None, currentStep=0):
         """Write out a model to the DCD file.
 
         The periodic box can be specified either by the unit cell dimensions
@@ -116,6 +113,8 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
             The dimensions of the crystallographic unit cell.
         periodicBoxVectors : tuple of Vec3=None
             The vectors defining the periodic box.
+        currentStep : int=0
+            The current step of the simulation.
         """
         if self._topology.getNumAtoms() != len(positions):
             raise ValueError('The number of positions must match the number of atoms')
@@ -129,23 +128,21 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
             raise ValueError('Particle position is infinite.  For more information, see https://github.com/openmm/openmm/wiki/Frequently-Asked-Questions#nan')
         file = self._file
 
-        self._modelCount += 1
-        if self._interval > 1 and self._firstStep+self._modelCount*self._interval > 1<<31:
+        if self._interval > 1 and currentStep + self._interval > 1<<31:
             # This will exceed the range of a 32 bit integer.  To avoid crashing or producing a corrupt file,
             # update the header to say the trajectory consisted of a smaller number of larger steps (so the
             # total trajectory length remains correct).
-            self._firstStep //= self._interval
             self._dt *= self._interval
             self._interval = 1
             file.seek(0, os.SEEK_SET)
-            file.write(struct.pack('<i4c9if', 84, b'C', b'O', b'R', b'D', 0, self._firstStep, self._interval, 0, 0, 0, 0, 0, 0, self._dt))
+            file.write(struct.pack('<i4c9if', 84, b'C', b'O', b'R', b'D', 0, 1, self._interval, 0, 0, 0, 0, 0, 0, self._dt))
 
         # Update the header.
 
         file.seek(8, os.SEEK_SET)
-        file.write(struct.pack('<i', self._modelCount))
+        file.write(struct.pack('<i', currentStep // self._interval))
         file.seek(20, os.SEEK_SET)
-        file.write(struct.pack('<i', self._firstStep+self._modelCount*self._interval))
+        file.write(struct.pack('<i', currentStep))
 
         # Write the data.
 
diff --git a/wrappers/python/openmm/app/dcdreporter.py b/wrappers/python/openmm/app/dcdreporter.py
index d09d5c149c..4626e1f225 100644
--- a/wrappers/python/openmm/app/dcdreporter.py
+++ b/wrappers/python/openmm/app/dcdreporter.py
@@ -115,7 +115,7 @@ def report(self, simulation, state):
         positions = state.getPositions(asNumpy=True)
         if self._atomSubset is not None:
             positions = [positions[i] for i in self._atomSubset]
-        self._dcd.writeModel(positions, periodicBoxVectors=state.getPeriodicBoxVectors())
+        self._dcd.writeModel(positions, periodicBoxVectors=state.getPeriodicBoxVectors(), currentStep=simulation.currentStep)
 
     def __del__(self):
         self._out.close()
diff --git a/wrappers/python/openmm/app/xtcfile.py b/wrappers/python/openmm/app/xtcfile.py
index 3fb821eb5f..1852c72265 100644
--- a/wrappers/python/openmm/app/xtcfile.py
+++ b/wrappers/python/openmm/app/xtcfile.py
@@ -22,7 +22,7 @@ class XTCFile(object):
     To use this class, create a XTCFile object, then call writeModel() once for each model in the file.
     """
 
-    def __init__(self, fileName, topology, dt, firstStep=0, interval=1, append=False):
+    def __init__(self, fileName, topology, dt, interval=1, append=False):
         """Create a XTC file, or open an existing file to append.
 
         Parameters
@@ -33,8 +33,6 @@ def __init__(self, fileName, topology, dt, firstStep=0, interval=1, append=False
             The Topology defining the molecular system being written
         dt : time
             The time step used in the trajectory
-        firstStep : int=0
-            The index of the first step in the trajectory
         interval : int=1
             The frequency (measured in time steps) at which states are written
             to the trajectory
@@ -45,16 +43,13 @@ def __init__(self, fileName, topology, dt, firstStep=0, interval=1, append=False
             raise TypeError("fileName must be a string")
         self._filename = fileName
         self._topology = topology
-        self._firstStep = firstStep
         self._interval = interval
-        self._modelCount = 0
         if is_quantity(dt):
             dt = dt.value_in_unit(picoseconds)
         self._dt = dt
         if append:
             if not os.path.isfile(self._filename):
                 raise FileNotFoundError(f"The file '{self._filename}' does not exist.")
-            self._modelCount = get_xtc_nframes(self._filename.encode("utf-8"))
             natoms = get_xtc_natoms(self._filename.encode("utf-8"))
             if natoms != topology.getNumAtoms():
                 raise ValueError(
@@ -67,7 +62,7 @@ def __init__(self, fileName, topology, dt, firstStep=0, interval=1, append=False
     def _getNumFrames(self):
         return get_xtc_nframes(self._filename.encode("utf-8"))
 
-    def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None):
+    def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None, currentStep=0):
         """Write out a model to the XTC file.
 
         The periodic box can be specified either by the unit cell dimensions
@@ -85,6 +80,8 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
             The dimensions of the crystallographic unit cell.
         periodicBoxVectors : tuple of Vec3=None
             The vectors defining the periodic box.
+        currentStep : int=0
+            The current step in the trajectory
         """
         if self._topology.getNumAtoms() != len(positions):
             raise ValueError("The number of positions must match the number of atoms")
@@ -103,15 +100,14 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
         if (np.abs(positions * 1000) > (2 ** 31 - 1)).any():
             raise ValueError("Particle position is too large for XTC format")
 
-        self._modelCount += 1
         if (
             self._interval > 1
-            and self._firstStep + self._modelCount * self._interval > 1 << 31
+            and currentStep + self._interval > 1 << 31
         ):
             # This will exceed the range of a 32 bit integer.  To avoid crashing or producing a corrupt file,
             # update the file to say the trajectory consisted of a smaller number of larger steps (so the
             # total trajectory length remains correct).
-            self._firstStep //= self._interval
+            currentStep //= self._interval
             self._dt *= self._interval
             self._interval = 1
             with tempfile.TemporaryDirectory() as temp:
@@ -119,7 +115,7 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
                 xtc_rewrite_with_new_timestep(
                     self._filename.encode("utf-8"),
                     fname.encode("utf-8"),
-                    self._firstStep,
+                    1,
                     self._interval,
                     self._dt,
                 )
@@ -143,12 +139,12 @@ def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None
             )
         else:
             boxVectors = np.zeros((3, 3)).astype(np.float32)
-        step = (self._modelCount - 1) * self._interval + self._firstStep
-        time = step * self._dt
+
+        time = currentStep * self._dt
         xtc_write_frame(
             self._filename.encode("utf-8"),
             np.array(positions, dtype=np.float32),
             boxVectors,
             np.float32(time),
-            np.int32(step),
+            np.int32(currentStep),
         )
diff --git a/wrappers/python/openmm/app/xtcreporter.py b/wrappers/python/openmm/app/xtcreporter.py
index 0a05e59ad6..9642e54666 100644
--- a/wrappers/python/openmm/app/xtcreporter.py
+++ b/wrappers/python/openmm/app/xtcreporter.py
@@ -107,11 +107,10 @@ def report(self, simulation, state):
                 self._fileName,
                 topology,
                 simulation.integrator.getStepSize(),
-                simulation.currentStep,
                 self._reportInterval,
                 self._append,
             )
         positions = state.getPositions(asNumpy=True)
         if self._atomSubset is not None:
             positions = [positions[i] for i in self._atomSubset]
-        self._xtc.writeModel(positions, periodicBoxVectors=state.getPeriodicBoxVectors())
+        self._xtc.writeModel(positions, periodicBoxVectors=state.getPeriodicBoxVectors(), currentStep=simulation.currentStep)


Test patch:
diff --git a/wrappers/python/tests/TestDcdFile.py b/wrappers/python/tests/TestDcdFile.py
index 1ce1cbdeaf..a2bd280bce 100644
--- a/wrappers/python/tests/TestDcdFile.py
+++ b/wrappers/python/tests/TestDcdFile.py
@@ -6,6 +6,13 @@
 from random import random
 import os
 
+def _readModelCount(file):
+    import struct
+
+    with open(file, "r+b") as f:
+        f.seek(8, os.SEEK_SET)
+        return struct.unpack('<i', f.read(4))[0]
+
 class TestDCDFile(unittest.TestCase):
     def test_dcd(self):
         """ Test the DCD file """
@@ -15,7 +22,7 @@ def test_dcd(self):
         with open(fname, 'wb') as f:
             dcd = app.DCDFile(f, pdbfile.topology, 0.001)
             for i in range(5):
-                dcd.writeModel([mm.Vec3(random(), random(), random()) for j in range(natom)]*unit.angstroms)
+                dcd.writeModel([mm.Vec3(random(), random(), random()) for j in range(natom)]*unit.angstroms, currentStep=i)
         os.remove(fname)
     
     def testLongTrajectory(self):
@@ -26,7 +33,7 @@ def testLongTrajectory(self):
         with open(fname, 'wb') as f:
             dcd = app.DCDFile(f, pdbfile.topology, 0.001, interval=1000000000)
             for i in range(5):
-                dcd.writeModel([mm.Vec3(random(), random(), random()) for j in range(natom)]*unit.angstroms)
+                dcd.writeModel([mm.Vec3(random(), random(), random()) for j in range(natom)]*unit.angstroms, currentStep=i)
         os.remove(fname)
     
     def testAppend(self):
@@ -45,7 +52,7 @@ def testAppend(self):
         simulation.context.setPositions(pdb.positions)
         simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
         simulation.step(10)
-        self.assertEqual(5, dcd._dcd._modelCount)
+        self.assertEqual(5, _readModelCount(fname))
         del simulation
         del dcd
         len1 = os.stat(fname).st_size
@@ -54,12 +61,13 @@ def testAppend(self):
 
         integrator = mm.VerletIntegrator(0.001*unit.picoseconds)
         simulation = app.Simulation(pdb.topology, system, integrator, mm.Platform.getPlatform('Reference'))
+        simulation.currentStep = 10
         dcd = app.DCDReporter(fname, 2, append=True)
         simulation.reporters.append(dcd)
         simulation.context.setPositions(pdb.positions)
         simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
         simulation.step(10)
-        self.assertEqual(10, dcd._dcd._modelCount)
+        self.assertEqual(10, _readModelCount(fname))
         len2 = os.stat(fname).st_size
         self.assertTrue(len2-len1 > 3*4*5*system.getNumParticles())
         del simulation
@@ -83,7 +91,7 @@ def testAtomSubset(self):
         simulation.context.setPositions(pdb.positions)
         simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
         simulation.step(10)
-        self.assertEqual(5, dcd._dcd._modelCount)
+        self.assertEqual(5, _readModelCount(fname))
         del simulation
         del dcd
         len1 = os.stat(fname).st_size
@@ -92,12 +100,13 @@ def testAtomSubset(self):
 
         integrator = mm.VerletIntegrator(0.001*unit.picoseconds)
         simulation = app.Simulation(pdb.topology, system, integrator, mm.Platform.getPlatform('Reference'))
+        simulation.currentStep = 10
         dcd = app.DCDReporter(fname, 2, append=True, atomSubset=atomSubset)
         simulation.reporters.append(dcd)
         simulation.context.setPositions(pdb.positions)
         simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
         simulation.step(10)
-        self.assertEqual(10, dcd._dcd._modelCount)
+        self.assertEqual(10, _readModelCount(fname))
         len2 = os.stat(fname).st_size
         self.assertTrue(len2-len1 > 3*4*5*len(atomSubset))
         del simulation
diff --git a/wrappers/python/tests/TestXtcFile.py b/wrappers/python/tests/TestXtcFile.py
index dbbb93cc7a..aa59b446fb 100644
--- a/wrappers/python/tests/TestXtcFile.py
+++ b/wrappers/python/tests/TestXtcFile.py
@@ -34,7 +34,7 @@ def test_xtc_triclinic(self):
                     mm.Vec3(random(), random(), random()) * unit.nanometers,
                 )
                 box.append(np.array([[vec.x, vec.y, vec.z] for vec in box_i]))
-                xtc.writeModel(coords[i], periodicBoxVectors=box_i)
+                xtc.writeModel(coords[i], periodicBoxVectors=box_i, currentStep=i)
             # The  XTCFile class  does not  provide a  way to  read the
             # trajectory back, but the underlying XTC library does
             coords_read, box_read, time, step = read_xtc(fname.encode("utf-8"))
@@ -72,7 +72,7 @@ def test_xtc_cubic(self):
                 )
                 box_i = mm.Vec3(random(), random(), random()) * unit.nanometers
                 box.append(np.diag(box_i[:3].value_in_unit(unit.nanometers)))
-                xtc.writeModel(coords[i], unitCellDimensions=box_i)
+                xtc.writeModel(coords[i], unitCellDimensions=box_i, currentStep=i)
             #The  XTCFile class  does not  provide a  way to  read the
             #trajectory back, but the underlying XTC library does
             coords_read, box_read, time, step = read_xtc(fname.encode("utf-8"))
@@ -108,7 +108,7 @@ def test_xtc_box_from_topology(self):
                     [mm.Vec3(random(), random(), random()) for j in range(natom)]
                     * unit.nanometers
                 )
-                xtc.writeModel(coords[i])
+                xtc.writeModel(coords[i], currentStep=i)
             # The  XTCFile class  does not  provide a  way to  read the
             # trajectory back, but the underlying XTC library does
             coords_read, box_read, time, step = read_xtc(fname.encode("utf-8"))
@@ -158,7 +158,7 @@ def test_xtc_small(self):
                     mm.Vec3(random(), random(), random()) * unit.nanometers,
                 )
                 box.append(np.array([[vec.x, vec.y, vec.z] for vec in box_i]))
-                xtc.writeModel(coords[i], periodicBoxVectors=box_i)
+                xtc.writeModel(coords[i], periodicBoxVectors=box_i, currentStep=i)
             # The  XTCFile class  does not  provide a  way to  read the
             # trajectory back, but the underlying XTC library does
             coords_read, box_read, time, step = read_xtc(fname.encode("utf-8"))
@@ -187,7 +187,8 @@ def testLongTrajectory(self):
             for i in range(5):
                 xtc.writeModel(
                     [mm.Vec3(random(), random(), random()) for j in range(natom)]
-                    * unit.angstroms
+                    * unit.angstroms,
+                    currentStep=i
                 )
 
     def testAppend(self):
@@ -212,7 +213,6 @@ def testAppend(self):
             simulation.context.setPositions(pdb.positions)
             simulation.context.setVelocitiesToTemperature(300 * unit.kelvin)
             simulation.step(10)
-            self.assertEqual(5, xtc._xtc._modelCount)
             self.assertEqual(5, xtc._xtc._getNumFrames())
             del simulation
             del xtc
@@ -226,12 +226,12 @@ def testAppend(self):
                 integrator,
                 mm.Platform.getPlatform("Reference"),
             )
+            simulation.currentStep = 10
             xtc = app.XTCReporter(fname, 2, append=True)
             simulation.reporters.append(xtc)
             simulation.context.setPositions(pdb.positions)
             simulation.context.setVelocitiesToTemperature(300 * unit.kelvin)
             simulation.step(10)
-            self.assertEqual(10, xtc._xtc._modelCount)
             self.assertEqual(10, xtc._xtc._getNumFrames())
             del simulation
             del xtc
@@ -259,7 +259,6 @@ def testAtomSubset(self):
             simulation.context.setPositions(pdb.positions)
             simulation.context.setVelocitiesToTemperature(300 * unit.kelvin)
             simulation.step(10)
-            self.assertEqual(5, xtc._xtc._modelCount)
             self.assertEqual(5, xtc._xtc._getNumFrames())
 
             # The  XTCFile class  does not  provide a  way to  read the