Adjust plain MD simulation PDB writing & production step counter

news/md_outputs.rst

@@ -0,0 +1,26 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* Plain MD simulations now write out PDBs with periodicity shifts consistent
+  with the simulation XTC file (PR #1302).
+* The simulation log file written during plain MD simulations now starts from
+  zero rather than including the steps from the equilibration.
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

openfe/protocols/openmm_md/plain_md_methods.py

@@ -25,6 +25,7 @@
 import time
 import mdtraj
 from mdtraj.reporters import XTCReporter
+import numpy.typing as npt
 from openfe.utils import without_oechem_backend, log_system_probe
 from gufe import (
     settings,
@@ -267,6 +268,40 @@ def __init__(
             generation=generation
         )
 
+    @staticmethod
+    def _save_simulation_pdb(
+        simulation: openmm.app.Simulation,
+        selection_indices: npt.NDArray,
+        filepath: pathlib.Path
+    ) -> None:
+        """
+        Helper method to write out PDB for the current state of the simulation.
+
+        Parameters
+        ----------
+        simulation : openmm.app.Simulation
+          The simulation to save the PDB for.
+        output_indices : str
+          A string defining which parts of the system to save an output for.
+        filepath : pathlib.Path
+          The full path of the file to be written.
+        """
+        state = simulation.context.getState(
+            getPositions=True,
+            enforcePeriodicBox=simulation._usesPBC,
+        )
+
+        positions = to_openmm(from_openmm(state.getPositions()))
+
+        # Store subset of atoms, specified in input, as PDB file
+        mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology)
+        traj = mdtraj.Trajectory(
+            positions[selection_indices, :],
+            mdtraj_top.subset(selection_indices),
+        )
+
+        traj.save_pdb(filepath)
+
     @staticmethod
     def _run_MD(simulation: openmm.app.Simulation,
                 positions: omm_unit.Quantity,
@@ -326,25 +361,19 @@ def _run_MD(simulation: openmm.app.Simulation,
             maxIterations=simulation_settings.minimization_steps
         )
 
-        # Get the sub selection of the system to save coords for
-        selection_indices = mdtraj.Topology.from_openmm(
-            simulation.topology).select(output_settings.output_indices)
-
-        positions = to_openmm(from_openmm(
-            simulation.context.getState(getPositions=True,
-                                        enforcePeriodicBox=False
-                                        ).getPositions()))
-        # Store subset of atoms, specified in input, as PDB file
-        mdtraj_top = mdtraj.Topology.from_openmm(simulation.topology)
-        traj = mdtraj.Trajectory(
-            positions[selection_indices, :],
-            mdtraj_top.subset(selection_indices),
-        )
+        # Get the sub selection of the system for coord saving
+        # will be using this going forward
+        mdtop = mdtraj.Topology.from_openmm(simulation.topology)
+        selection_indices = mdtop.select(output_settings.output_indices)
 
+        # Save a PDB of the minimized structure
         if output_settings.minimized_structure:
-            traj.save_pdb(
+            PlainMDProtocolUnit._save_simulation_pdb(
+                simulation,
+                selection_indices,
                 shared_basepath / output_settings.minimized_structure
             )
+
         # equilibrate
         # NVT equilibration
         if equil_steps_nvt:
@@ -366,18 +395,11 @@ def _run_MD(simulation: openmm.app.Simulation,
                 logger.info(
                     f"Completed NVT equilibration in {t1 - t0} seconds")
 
-            # Save last frame NVT equilibration
-            positions = to_openmm(
-                from_openmm(simulation.context.getState(
-                    getPositions=True, enforcePeriodicBox=False
-                ).getPositions()))
-
-            traj = mdtraj.Trajectory(
-                positions[selection_indices, :],
-                mdtraj_top.subset(selection_indices),
-            )
+            # Optionally save last frame NVT equilibration
             if output_settings.equil_nvt_structure is not None:
-                traj.save_pdb(
+                PlainMDProtocolUnit._save_simulation_pdb(
+                    simulation,
+                    selection_indices,
                     shared_basepath / output_settings.equil_nvt_structure
                 )
 
@@ -400,24 +422,20 @@ def _run_MD(simulation: openmm.app.Simulation,
                 f"Completed NPT equilibration in {t1 - t0} seconds")
 
         # Save last frame NPT equilibration
-        positions = to_openmm(
-            from_openmm(simulation.context.getState(
-                getPositions=True, enforcePeriodicBox=False
-            ).getPositions()))
-
-        traj = mdtraj.Trajectory(
-            positions[selection_indices, :],
-            mdtraj_top.subset(selection_indices),
-        )
         if output_settings.equil_npt_structure is not None:
-            traj.save_pdb(
+            PlainMDProtocolUnit._save_simulation_pdb(
+                simulation,
+                selection_indices,
                 shared_basepath / output_settings.equil_npt_structure
             )
 
         # production
         if verbose:
             logger.info("running production phase")
 
+        # For the production simulation, we reset the step counter to 0
+        simulation.context.setStepCount(0)
+
         # Setup the reporters
         write_interval = settings_validation.divmod_time_and_check(
             output_settings.trajectory_write_interval,
@@ -433,18 +451,23 @@ def _run_MD(simulation: openmm.app.Simulation,
         )
 
         if output_settings.production_trajectory_filename:
-            simulation.reporters.append(XTCReporter(
+            xtc_reporter = XTCReporter(
                 file=str(
                     shared_basepath /
-                    output_settings.production_trajectory_filename),
+                    output_settings.production_trajectory_filename
+                ),
                 reportInterval=write_interval,
-                atomSubset=selection_indices))
+                atomSubset=selection_indices
+            )
+            simulation.reporters.append(xtc_reporter)
+
         if output_settings.checkpoint_storage_filename:
             simulation.reporters.append(openmm.app.CheckpointReporter(
                 file=str(
                     shared_basepath /
                     output_settings.checkpoint_storage_filename),
                 reportInterval=checkpoint_interval))
+
         if output_settings.log_output:
             simulation.reporters.append(openmm.app.StateDataReporter(
                 str(shared_basepath / output_settings.log_output),
@@ -459,9 +482,19 @@ def _run_MD(simulation: openmm.app.Simulation,
                 density=True,
                 speed=True,
             ))
+
         t0 = time.time()
         simulation.step(prod_steps)
         t1 = time.time()
+
+        # Write the final production frame
+        if output_settings.production_structure is not None:
+            PlainMDProtocolUnit._save_simulation_pdb(
+                simulation,
+                selection_indices,
+                shared_basepath / output_settings.production_structure
+            )
+
         if verbose:
             logger.info(f"Completed simulation in {t1 - t0} seconds")
 
@@ -700,6 +733,9 @@ def run(self, *, dry=False, verbose=True,
             if output_settings.equil_npt_structure:
                 output['npt_equil_pdb'] = shared_basepath / output_settings.equil_npt_structure
 
+            if output_settings.production_structure:
+                output['production_pdb'] = shared_basepath / output_settings.production_structure
+
             return output
         else:
             return {'debug': {'system': stateA_system}}

openfe/protocols/openmm_utils/omm_settings.py

@@ -673,17 +673,24 @@ class Config:
     Frequency to write the xtc file. Default 5000 * unit.timestep.
     """
     preminimized_structure: Optional[str] = 'system.pdb'
-    """Path to the pdb file of the full pre-minimized system. 
+    """Filename for the pdb file of the full pre-minimized system. 
     Default 'system.pdb'."""
     minimized_structure: Optional[str] = 'minimized.pdb'
-    """Path to the pdb file of the system after minimization. 
-    Only the specified atom subset is saved. Default 'minimized.pdb'."""
+    """Filename for the pdb file of the system after minimization. 
+    Only the specified atom subset defined by ``output_indices`` is saved. 
+    Default 'minimized.pdb'."""
     equil_nvt_structure: Optional[str] = 'equil_nvt.pdb'
-    """Path to the pdb file of the system after NVT equilibration. 
-    Only the specified atom subset is saved. Default 'equil_nvt.pdb'."""
+    """Filename for the pdb file of the system after NVT equilibration. 
+    Only the specified atom subset defined by ``output_indices`` is saved. 
+    Default 'equil_nvt.pdb'."""
     equil_npt_structure: Optional[str] = 'equil_npt.pdb'
-    """Path to the pdb file of the system after NPT equilibration. 
-    Only the specified atom subset is saved. Default 'equil_npt.pdb'."""
+    """Filename for the pdb file of the system after NPT equilibration. 
+    Only the specified atom subset defined by ``output_indices`` is saved. 
+    Default 'equil_npt.pdb'."""
+    production_structure: Optional[str] = 'production.pdb'
+    """Filename of the pdb file of the system after the production simulation. 
+    Only the specified atom subset defined by ``output_indices`` is saved. 
+    Default to 'production.pdb'."""
     log_output: Optional[str] = 'simulation.log'
     """
     Filename for writing the log of the MD simulation, including timesteps,

openfe/tests/protocols/openmm_md/test_plain_md_slow.py

@@ -1,5 +1,7 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/openfe
+import MDAnalysis as mda
+from numpy.testing import assert_allclose
 import pathlib
 import pytest
 from openff.units import unit
@@ -94,6 +96,7 @@
     settings.simulation_settings.equilibration_length = 50 * unit.picosecond
     settings.simulation_settings.production_length = 100 * unit.picosecond
     settings.output_settings.checkpoint_interval = 10 * unit.picosecond
+    settings.output_settings.trajectory_write_interval = 10 * unit.picosecond
     settings.engine_settings.compute_platform = platform
 
     prot = openmm_md.PlainMDProtocol(settings)
@@ -127,6 +130,7 @@
         "checkpoint.chk",
         "equil_nvt.pdb",
         "equil_npt.pdb",
+        "production.pdb",
         "minimized.pdb",
         "simulation.xtc",
         "simulation.log",
@@ -142,3 +146,22 @@
     assert pur.outputs['last_checkpoint'] == unit_shared / "checkpoint.chk"
     assert pur.outputs['nvt_equil_pdb'] == unit_shared / "equil_nvt.pdb"
     assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb"
+    assert pur.outputs['production_pdb'] == unit_shared / "production.pdb"
+
+    # Check the final trajectory frame
+    first_frame = mda.Universe(pur.outputs['npt_equil_pdb'])
+    final_frame = mda.Universe(pur.outputs['production_pdb'])
+    simulation = mda.Universe(pur.outputs['minimized_pdb'], pur.outputs['nc'])
+
+    # Check we have the right number of frames
+    assert len(simulation.trajectory) == 10
+
+    # Check that the final frame matches
+    simulation.trajectory[-1]  # fast-forward
+    assert_allclose(
+        final_frame.atoms.positions,
+        simulation.atoms.positions,
+        rtol=0,
+        atol=1e-2
+    )
+