Merge pull request #4899 from epretti/dcd-comments

Check comment block length from DCD when appending

wrappers/python/openmm/app/dcdfile.py

@@ -6,7 +6,7 @@ Simbios, the NIH National Center for Physics-Based Simulation of
 Biological Structures at Stanford, funded under the NIH Roadmap for
 Medical Research, grant U54 GM072970. See https://simtk.org.
 
-Portions copyright (c) 2012 Stanford University and the Authors.
+Portions copyright (c) 2012-2025 Stanford University and the Authors.
 Authors: Peter Eastman
 Contributors:
 
@@ -84,12 +84,17 @@ class DCDFile(object):
         if topology.getUnitCellDimensions() is not None:
             boxFlag = 1
         if append:
-            file.seek(8, os.SEEK_SET)
+            file.seek(0, os.SEEK_SET)
+            headerMagic = file.read(8)
+            if headerMagic[4:8] != b'CORD' or struct.unpack('<i', headerMagic[:4])[0] != 84:
+                raise ValueError('Cannot append to file with invalid DCD header')
             self._modelCount = struct.unpack('<i', file.read(4))[0]
-            file.seek(268, os.SEEK_SET)
+            file.seek(92, os.SEEK_SET)
+            commentsBytes = struct.unpack('<i', file.read(4))[0]
+            file.seek(104 + commentsBytes, os.SEEK_SET)
             numAtoms = struct.unpack('<i', file.read(4))[0]
             if numAtoms != topology.getNumAtoms():
-                raise ValueError('Cannot append to a DCD file that contains a different number of atoms')
+                raise ValueError(f'Cannot append from system with {topology.getNumAtoms()} atoms to DCD file with {numAtoms} 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)

wrappers/python/tests/TestDcdFile.py

@@ -5,11 +5,10 @@ import openmm as mm
 from openmm import unit
 from random import random
 import os
+import struct
 
 
 def _read_dcd_header(file):
-    import struct
-    
     with open(file, "r+b") as f:
         f.seek(8, os.SEEK_SET)
         modelCount = struct.unpack("<i", f.read(4))[0]
@@ -129,6 +128,97 @@ class TestDCDFile(unittest.TestCase):
         self.assertEqual(20, currStep)
         os.remove(fname)
 
+    def testAppendAtomCountMismatch(self):
+        """Test that appending to a DCD file with a different number of atoms raises an error."""
+        fname = tempfile.mktemp(suffix='.dcd')
+        pdb = app.PDBFile('systems/alanine-dipeptide-explicit.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.getPlatform('Reference'))
+        atomSubset1 = [atom.index for chain, _ in zip(pdb.topology.chains(), range(1)) for atom in chain.atoms()]
+        atomSubset2 = [atom.index for chain, _ in zip(pdb.topology.chains(), range(2)) for atom in chain.atoms()]
+        dcd = app.DCDReporter(fname, 2, atomSubset=atomSubset1)
+        simulation.reporters.append(dcd)
+        simulation.context.setPositions(pdb.positions)
+        simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
+        simulation.step(10)
+        del simulation
+        del dcd
+
+        # 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.getPlatform('Reference'))
+        simulation.currentStep = 10
+        dcd = app.DCDReporter(fname, 2, append=True, atomSubset=atomSubset2)
+        simulation.reporters.append(dcd)
+        simulation.context.setPositions(pdb.positions)
+        simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
+        with self.assertRaises(ValueError):
+            simulation.step(10)
+
+    def testAppendLongCommentBlock(self):
+        """Test appending to an existing trajectory with a long comment block."""
+        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.getPlatform('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
+        del dcd
+        len1 = os.stat(fname).st_size
+
+        # Some software writes more than 2 80-byte "comment lines" to a DCD
+        # file.  Modify the DCD to simulate this and ensure we can append.
+
+        commentLines = 10
+        with open(fname, "rb") as dcdFile:
+            dcdHeader = dcdFile.read(92)
+            dcdCommentsLength = struct.unpack("<i", dcdFile.read(4))[0]
+            dcdFile.read(dcdCommentsLength + 4)
+            dcdContents = dcdFile.read()
+        with open(fname, "wb") as dcdFile:
+            dcdFile.write(dcdHeader)
+            dcdNewCommentsLength = 4 + 80 * commentLines
+            dcdFile.write(struct.pack("<2i", dcdNewCommentsLength, commentLines))
+            dcdFile.write(bytes(80 * commentLines))
+            dcdFile.write(struct.pack("<i", dcdNewCommentsLength))
+            dcdFile.write(dcdContents)
+
+        # 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.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)
+        len2 = os.stat(fname).st_size
+        self.assertEqual(len2 - len1, dcdNewCommentsLength - dcdCommentsLength + 5 * 3 * 4 * (system.getNumParticles() + 2))
+        del simulation
+        del dcd
+        modelCount, currStep = _read_dcd_header(fname)
+        self.assertEqual(10, modelCount)
+        self.assertEqual(20, currStep)
+        os.remove(fname)
+
 
 if __name__ == '__main__':
     unittest.main()