Adjust plain MD simulation PDB writing & production step counter

.github/workflows/gen-lock-file.yaml

@@ -8,27 +8,27 @@ defaults:
     shell: bash -leo pipefail {0}
 
 jobs:
-  create-conda-lock-file:
+  create-conda-lock-file-and-test-linux:
     runs-on: ubuntu-latest
 
-    steps:
-    - name: Checkout repository
-      uses: actions/checkout@v4
-      with:
-        fetch-depth: 0
-
+    steps:  
     - name: Install conda-lock with Micromamba
       uses: mamba-org/setup-micromamba@v2
       with:
         environment-name: conda-lock
         create-args: >-
           conda-lock
 
-    - name: Get Latest Version
+    # This saves me some time since we only need the latest tag
+    - name: Get latest tag
       id: latest-version
+      env:
+        GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        REPO: ${{ github.repository }}
       run: |
-        LATEST_TAG=$(git describe --tags $(git rev-list --tags --max-count=1))
-        # slice off the v, ie v0.7.2 -> 0.7.2
+        LATEST_TAG=$(curl -s -H "Authorization: token $GH_TOKEN" \
+          https://api.github.com/repos/$REPO/tags \
+          | jq -r '.[0].name')
         VERSION=${LATEST_TAG:1}
         echo $VERSION
         echo "VERSION=$VERSION" >> $GITHUB_OUTPUT
@@ -52,15 +52,49 @@ jobs:
 
     - name: Generate lock files
       run: | 
-        conda lock --with-cuda 11.8 -f environment-to-lock.yaml openfe-conda-lock.yml 
+        conda lock --with-cuda 11.8 -f environment-to-lock.yaml --lockfile openfe-conda-lock.yml 
         cp openfe-conda-lock.yml openfe-${{ steps.latest-version.outputs.VERSION }}-conda-lock.yml
 
-    - name: Upload openfe forge to release
+    - name: Test lock file (linux)
+      run: |
+        conda-lock install -p /home/runner/micromamba/envs/lf-test openfe-conda-lock.yml
+        micromamba activate /home/runner/micromamba/envs/lf-test
+        openfe test
+
+    - name: Upload file as artifact
+      uses: actions/upload-artifact@v4
+      with:
+        name: conda-lock-files
+        path: "*conda-lock.yml"
+
+  test-osx-lock-file:
+    needs: create-conda-lock-file-and-test-linux
+    runs-on: macos-latest
+    steps:
+    - name: Download artifact
+      uses: actions/download-artifact@v4
+      with:
+        name: conda-lock-files
+
+    - name: Install conda-lock with Micromamba
+      uses: mamba-org/setup-micromamba@v2
+      with:
+        environment-name: conda-lock
+        create-args: >-
+          conda-lock
+
+    - name: Test lock file (osx)
+      run: |
+        conda-lock install -p /Users/runner/micromamba/envs/lf-test openfe-conda-lock.yml
+        micromamba activate /Users/runner/micromamba/envs/lf-test
+        openfe test
+
+    - name: Upload lock files to release
       uses: svenstaro/upload-release-action@v2
+      if: startsWith(github.ref, 'refs/tags/')
       with:
         repo_token: ${{ secrets.GITHUB_TOKEN }}
         file: "*conda-lock.yml"
         tag: ${{ github.ref }}
         overwrite: true
         file_glob: true
-      if: startsWith(github.ref, 'refs/tags/')

environment.yml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - cinnabar ~=0.4.0
-  - click <8.2.0
+  - click >=8.2.0
   - coverage
   - duecredit<0.10
   - kartograf>=1.0.0

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,