Added atomSubset option to DCDReporter and XTCReporter (#4741)

wrappers/python/openmm/app/dcdfile.py

@@ -88,14 +88,14 @@ class DCDFile(object):
             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())):
+            if numAtoms != topology.getNumAtoms():
                 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')
             header += struct.pack('<80s', b'Created '+time.asctime(time.localtime(time.time())).encode('ascii'))
-            header += struct.pack('<4i', 164, 4, len(list(topology.atoms())), 4)
+            header += struct.pack('<4i', 164, 4, topology.getNumAtoms(), 4)
             file.write(header)
 
     def writeModel(self, positions, unitCellDimensions=None, periodicBoxVectors=None):

wrappers/python/openmm/app/dcdreporter.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-2024 Stanford University and the Authors.
 Authors: Peter Eastman
 Contributors:
 
@@ -32,8 +32,7 @@ from __future__ import absolute_import
 __author__ = "Peter Eastman"
 __version__ = "1.0"
 
-from openmm.app import DCDFile
-from openmm.unit import nanometer
+from openmm.app import DCDFile, Topology
 
 class DCDReporter(object):
     """DCDReporter outputs a series of frames from a Simulation to a DCD file.
@@ -41,7 +40,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, append=False, enforcePeriodicBox=None):
+    def __init__(self, file, reportInterval, append=False, enforcePeriodicBox=None, atomSubset=None):
         """Create a DCDReporter.
 
         Parameters
@@ -57,10 +56,13 @@ class DCDReporter(object):
             lies in the same periodic box.  If None (the default), it will automatically decide whether
             to translate molecules based on whether the system being simulated uses periodic boundary
             conditions.
+        atomSubset: list
+            Atom indices (zero indexed) of the particles to output.  If None (the default), all particles will be output.
         """
         self._reportInterval = reportInterval
         self._append = append
         self._enforcePeriodicBox = enforcePeriodicBox
+        self._atomSubset = atomSubset
         if append:
             mode = 'r+b'
         else:
@@ -96,11 +98,24 @@ class DCDReporter(object):
         """
 
         if self._dcd is None:
+            if self._atomSubset is None:
+                topology = simulation.topology
+            else:
+                topology = Topology()
+                topology.setPeriodicBoxVectors(simulation.topology.getPeriodicBoxVectors())
+                atoms = list(simulation.topology.atoms())
+                chain = topology.addChain()
+                residue = topology.addResidue('', chain)
+                for i in self._atomSubset:
+                    topology.addAtom(atoms[i].name, atoms[i].element, residue)
             self._dcd = DCDFile(
-                self._out, simulation.topology, simulation.integrator.getStepSize(),
+                self._out, topology, simulation.integrator.getStepSize(),
                 simulation.currentStep, self._reportInterval, self._append
             )
-        self._dcd.writeModel(state.getPositions(asNumpy=True), periodicBoxVectors=state.getPeriodicBoxVectors())
+        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())
 
     def __del__(self):
         self._out.close()

wrappers/python/openmm/app/xtcfile.py

@@ -56,9 +56,9 @@ class XTCFile(object):
                 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 != len(list(topology.atoms())):
+            if natoms != topology.getNumAtoms():
                 raise ValueError(
-                    f"The number of atoms in the topology ({len(list(topology.atoms()))}) does not match the number of atoms in the XTC file ({natoms})"
+                    f"The number of atoms in the topology ({topology.getNumAtoms()}) does not match the number of atoms in the XTC file ({natoms})"
                 )
         else:
             if os.path.isfile(self._filename) and os.path.getsize(self._filename) > 0:

wrappers/python/openmm/app/xtcreporter.py

+"""
+xtcreporter.py: Outputs simulation trajectories in XTC format
+
+This is part of the OpenMM molecular simulation toolkit originating from
+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) 2024 Stanford University and the Authors.
+Authors: Raul P. Pelaez
+Contributors: Peter Eastman
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
 __author__ = "Raul P. Pelaez"
-from openmm.app import XTCFile
+from openmm.app import Topology, XTCFile
 
 class XTCReporter(object):
     """XTCReporter outputs a series of frames from a Simulation to a XTC file.
@@ -7,7 +37,7 @@ class XTCReporter(object):
     To use it, create a XTCReporter, then add it to the Simulation's list of reporters.
     """
 
-    def __init__(self, file, reportInterval, append=False, enforcePeriodicBox=None):
+    def __init__(self, file, reportInterval, append=False, enforcePeriodicBox=None, atomSubset=None):
         """Create a XTCReporter.
 
         Parameters
@@ -23,10 +53,13 @@ class XTCReporter(object):
             lies in the same periodic box.  If None (the default), it will automatically decide whether
             to translate molecules based on whether the system being simulated uses periodic boundary
             conditions.
+        atomSubset: list
+            Atom indices (zero indexed) of the particles to output.  If None (the default), all particles will be output.
         """
         self._reportInterval = reportInterval
         self._append = append
         self._enforcePeriodicBox = enforcePeriodicBox
+        self._atomSubset = atomSubset
         self._fileName = file
         self._xtc = None
         if not append:
@@ -60,14 +93,25 @@ class XTCReporter(object):
         """
 
         if self._xtc is None:
+            if self._atomSubset is None:
+                topology = simulation.topology
+            else:
+                topology = Topology()
+                topology.setPeriodicBoxVectors(simulation.topology.getPeriodicBoxVectors())
+                atoms = list(simulation.topology.atoms())
+                chain = topology.addChain()
+                residue = topology.addResidue('', chain)
+                for i in self._atomSubset:
+                    topology.addAtom(atoms[i].name, atoms[i].element, residue)
             self._xtc = XTCFile(
                 self._fileName,
-                simulation.topology,
+                topology,
                 simulation.integrator.getStepSize(),
                 simulation.currentStep,
                 self._reportInterval,
                 self._append,
             )
-        self._xtc.writeModel(
-            state.getPositions(asNumpy=True), periodicBoxVectors=state.getPeriodicBoxVectors()
-        )
+        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())

wrappers/python/tests/TestDcdFile.py

@@ -66,6 +66,44 @@ class TestDCDFile(unittest.TestCase):
         del dcd
         os.remove(fname)
 
+    def testAtomSubset(self):
+        """Test writing a DCD file containing a subset of atoms"""
+        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'))
+        atomSubset = [atom.index for atom in next(pdb.topology.chains()).atoms()]
+        dcd = app.DCDReporter(fname, 2, atomSubset=atomSubset)
+        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
+
+        # 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'))
+        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)
+        len2 = os.stat(fname).st_size
+        self.assertTrue(len2-len1 > 3*4*5*len(atomSubset))
+        del simulation
+        del dcd
+        os.remove(fname)
+
 
 if __name__ == '__main__':
     unittest.main()

wrappers/python/tests/TestXtcFile.py

@@ -194,6 +194,43 @@ class TestXtcFile(unittest.TestCase):
             del simulation
             del xtc
 
+    def testAtomSubset(self):
+        """Test writing an XTC file containing a subset of atoms"""
+        with tempfile.TemporaryDirectory() as temp:
+            fname = os.path.join(temp, 'traj.xtc')
+            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 XTC file.
+
+            integrator = mm.VerletIntegrator(1e-10 * unit.picoseconds)
+            simulation = app.Simulation(
+                pdb.topology,
+                system,
+                integrator,
+                mm.Platform.getPlatform("Reference"),
+            )
+            atomSubset = [atom.index for atom in next(pdb.topology.chains()).atoms()]
+            xtc = app.XTCReporter(fname, 2, atomSubset=atomSubset)
+            simulation.reporters.append(xtc)
+            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
+            # trajectory back, but the underlying XTC library does
+            coords_read, box_read, time, step = read_xtc(fname.encode("utf-8"))
+            self.assertEqual(coords_read.shape, (22, 3, 5))
+            self.assertEqual(box_read.shape, (3, 3, 5))
+            self.assertEqual(len(time), 5)
+            self.assertEqual(len(step), 5)
+            coords = [pdb.positions[i].value_in_unit(unit.nanometers) for i in atomSubset]
+            self.assertTrue(np.allclose(coords_read[:,:,0], coords, atol=1e-3))
+            self.assertTrue(np.allclose(box_read[:,:,0], pdb.topology.getPeriodicBoxVectors().value_in_unit(unit.nanometers), atol=1e-3))
+
 
 if __name__ == "__main__":
     unittest.main()