Repository: openmm/openmm — pull request #4741 (cpp)
Title: Added atomSubset option to DCDReporter and XTCReporter

When saving a trajectory to a DCD or XTC file, you can specify a subset of atoms to write.  Fixes #4564.

Resolved issues:
title: Recommended approach to writing out a subset of atoms with XTCReporter/DCDReporter?
body: I was going through the docs about [XTCReporter](https://docs.openmm.org/latest/api-python/generated/openmm.app.xtcreporter.XTCReporter.html) and [DCDReporter](https://docs.openmm.org/latest/api-python/generated/openmm.app.dcdreporter.DCDReporter.html) and I realized there's not really a well-documented way to use these reporters to save out a subset of atoms (or maybe I'm missing it somewhere else?). 

Is there a general recommended way to achieve this "writing out a subset of atoms" if we're writing out into XTC or DCD files? I'm imagining a basic scenario where it is known which atom indicies are desired to be saved, but the rest do not need to be written out. I'm currently envisioning having to write some kind of `subsetXTCReporter` or the like - would this be the best way to go about this?Just looking for any guidance here!

Addendum: You can do this subsetting of atoms using the [XTCReporter](https://github.com/mdtraj/mdtraj/blob/ccd0c874361e4e31127e3803d3670d08d57f7933/mdtraj/reporters/xtcreporter.py#L42-L68) and [DCDReporter](https://mdtraj.readthedocs.io/en/latest/api/generated/mdtraj.reporters.DCDReporter.html#mdtraj.reporters.DCDReporter) in MDTraj - I'm just wondering if there are alternative approaches now that OpenMM has it's own native XTCReporter

## 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 53aeaafec6..335cdc0b79 100644
--- a/wrappers/python/openmm/app/dcdfile.py
+++ b/wrappers/python/openmm/app/dcdfile.py
@@ -88,14 +88,14 @@ def __init__(self, file, topology, dt, firstStep=0, interval=1, append=False):
             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):
diff --git a/wrappers/python/openmm/app/dcdreporter.py b/wrappers/python/openmm/app/dcdreporter.py
index 09b1d9439a..d09d5c149c 100644
--- a/wrappers/python/openmm/app/dcdreporter.py
+++ b/wrappers/python/openmm/app/dcdreporter.py
@@ -6,7 +6,7 @@
 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 @@
 __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 @@ def __init__(self, file, reportInterval, append=False, enforcePeriodicBox=None):
             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 @@ def report(self, simulation, state):
         """
 
         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()
diff --git a/wrappers/python/openmm/app/xtcfile.py b/wrappers/python/openmm/app/xtcfile.py
index 39fc5226eb..3fb821eb5f 100644
--- a/wrappers/python/openmm/app/xtcfile.py
+++ b/wrappers/python/openmm/app/xtcfile.py
@@ -56,9 +56,9 @@ def __init__(self, fileName, topology, dt, firstStep=0, interval=1, append=False
                 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:
diff --git a/wrappers/python/openmm/app/xtcreporter.py b/wrappers/python/openmm/app/xtcreporter.py
index 38ef3b4d2d..0a05e59ad6 100644
--- a/wrappers/python/openmm/app/xtcreporter.py
+++ b/wrappers/python/openmm/app/xtcreporter.py
@@ -1,5 +1,35 @@
+"""
+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 @@ def __init__(self, file, reportInterval, append=False, enforcePeriodicBox=None):
             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 @@ def report(self, simulation, state):
         """
 
         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())


Test patch:
diff --git a/wrappers/python/tests/TestDcdFile.py b/wrappers/python/tests/TestDcdFile.py
index 73af629eff..1ce1cbdeaf 100644
--- a/wrappers/python/tests/TestDcdFile.py
+++ b/wrappers/python/tests/TestDcdFile.py
@@ -66,6 +66,44 @@ def testAppend(self):
         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()
diff --git a/wrappers/python/tests/TestXtcFile.py b/wrappers/python/tests/TestXtcFile.py
index e9eb954547..a66c6a2a3f 100644
--- a/wrappers/python/tests/TestXtcFile.py
+++ b/wrappers/python/tests/TestXtcFile.py
@@ -194,6 +194,43 @@ def testAppend(self):
             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()