Can append to DCD files

wrappers/python/simtk/openmm/app/dcdfile.py

@@ -52,8 +52,8 @@ class DCDFile(object):
 
     To use this class, create a DCDFile object, then call writeModel() once for each model in the file."""
 
-    def __init__(self, file, topology, dt, firstStep=0, interval=1):
-        """Create a DCD file and write out the header.
+    def __init__(self, file, topology, dt, firstStep=0, interval=1, append=False):
+        """Create a DCD file and write out the header, or open an existing file to append.
 
         Parameters
         ----------
@@ -68,6 +68,8 @@ class DCDFile(object):
         interval : int=1
             The frequency (measured in time steps) at which states are written
             to the trajectory
+        append : bool=False
+            If True, open an existing DCD file to append to.  If False, create a new file.
         """
         self._file = file
         self._topology = topology
@@ -81,6 +83,14 @@ class DCDFile(object):
         boxFlag = 0
         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 != len(list(topology.atoms())):
+                raise ValueError('Cannot append to a DCD file that contains a different number of atoms')
+        else:
             header = struct.pack('<i4c9if', 84, b'C', b'O', b'R', b'D', 0, firstStep, interval, 0, 0, 0, 0, 0, 0, dt)
             header += struct.pack('<13i', boxFlag, 0, 0, 0, 0, 0, 0, 0, 0, 24, 84, 164, 2)
             header += struct.pack('<80s', b'Created by OpenMM')
@@ -160,3 +170,7 @@ class DCDFile(object):
             data = array.array('f', (10*x[i] for x in positions))
             data.tofile(file)
             file.write(length)
+        try:
+            file.flush()
+        except AttributeError:
+            pass

wrappers/python/simtk/openmm/app/dcdreporter.py

@@ -42,7 +42,7 @@ class DCDReporter(object):
     To use it, create a DCDReporter, then add it to the Simulation's list of reporters.
     """
 
-    def __init__(self, file, reportInterval):
+    def __init__(self, file, reportInterval, append=False):
         """Create a DCDReporter.
 
         Parameters
@@ -51,9 +51,16 @@ class DCDReporter(object):
             The file to write to
         reportInterval : int
             The interval (in time steps) at which to write frames
+        append : bool=False
+            If True, open an existing DCD file to append to.  If False, create a new file.
         """
         self._reportInterval = reportInterval
-        self._out = open(file, 'wb')
+        self._append = append
+        if append:
+            mode = 'a+b'
+        else:
+            mode = 'wb'
+        self._out = open(file, mode)
         self._dcd = None
 
     def describeNextReport(self, simulation):
@@ -87,7 +94,7 @@ class DCDReporter(object):
         """
 
         if self._dcd is None:
-            self._dcd = DCDFile(self._out, simulation.topology, simulation.integrator.getStepSize(), 0, self._reportInterval)
+            self._dcd = DCDFile(self._out, simulation.topology, simulation.integrator.getStepSize(), 0, self._reportInterval, self._append)
         self._dcd.writeModel(state.getPositions(), periodicBoxVectors=state.getPeriodicBoxVectors())
 
     def __del__(self):

wrappers/python/tests/TestDcdFile.py

@@ -29,5 +29,40 @@ class TestDCDFile(unittest.TestCase):
                 dcd.writeModel([mm.Vec3(random(), random(), random()) for j in range(natom)]*unit.angstroms)
         os.remove(fname)
     
+    def testAppend(self):
+        """Test appending to an existing trajectory."""
+        fname = tempfile.mktemp(suffix='.dcd')
+        pdb = app.PDBFile('systems/alanine-dipeptide-implicit.pdb')
+        ff = app.ForceField('amber99sb.xml', 'tip3p.xml')
+        system = ff.createSystem(pdb.topology)
+        
+        # Create a simulation and write some frames to a DCD file.
+        
+        integrator = mm.VerletIntegrator(0.001*unit.picoseconds)
+        simulation = app.Simulation(pdb.topology, system, integrator, mm.Platform.getPlatformByName('Reference'))
+        dcd = app.DCDReporter(fname, 2)
+        simulation.reporters.append(dcd)
+        simulation.context.setPositions(pdb.positions)
+        simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
+        simulation.step(10)
+        self.assertEqual(5, dcd._dcd._modelCount)
+        del simulation
+        len1 = os.stat(fname).st_size
+
+        # Create a new simulation and have it append some more frames.
+
+        integrator = mm.VerletIntegrator(0.001*unit.picoseconds)
+        simulation = app.Simulation(pdb.topology, system, integrator, mm.Platform.getPlatformByName('Reference'))
+        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)
+        len2 = os.stat(fname).st_size
+        self.assertTrue(len2-len1 > 3*4*5*system.getNumParticles())
+        os.remove(fname)
+
+
 if __name__ == '__main__':
     unittest.main()