From 1fbf533e34033694abc7dbfc0a03407f447089fc Mon Sep 17 00:00:00 2001 From: IAlibay Date: Tue, 27 May 2025 00:17:39 +0100 Subject: [PATCH 1/5] Add final frame to XTC file + write final frame to PDB --- .../protocols/openmm_md/plain_md_methods.py | 115 ++++++++++++------ openfe/protocols/openmm_utils/omm_settings.py | 21 ++-- .../openmm/plain_md/test_simulation_slow.py | 10 ++ 3 files changed, 100 insertions(+), 46 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 51613b761..cac69aee0 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/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=False, + ) + + 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,17 +422,10 @@ 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 ) @@ -433,12 +448,15 @@ 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( @@ -462,6 +480,22 @@ def _run_MD(simulation: openmm.app.Simulation, t0 = time.time() simulation.step(prod_steps) t1 = time.time() + + if output_settings.production_trajectory_filename: + state = simulation.context.getState( + getPositions=True, + enforcePeriodicBox=False, + ) + xtc_reporter.report(simulation, state) + + # 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 +734,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}} diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 7bb39b731..5d52ff793 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/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, diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py index ecb7bc799..471208f85 100644 --- a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py +++ b/openfe/tests/protocols/openmm/plain_md/test_simulation_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 @@ -127,6 +129,7 @@ def test_complex_solvent_sim_gpu( "checkpoint.chk", "equil_nvt.pdb", "equil_npt.pdb", + "production.pdb", "minimized.pdb", "simulation.xtc", "simulation.log", @@ -142,3 +145,10 @@ def test_complex_solvent_sim_gpu( 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 + u = mda.Universe(pur.outputs['production_pdb']) + u2 = mda.Universe(pur.outputs['minimized_pdb'], pur.outputs['nc']) + u2.trajectory[-1] + assert_allclose(u.atoms.positions, u2.atoms.positions, rtol=0, atol=1e-2) From bf15f408a61838bf323b9b9c5b2710fccff9e836 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Tue, 27 May 2025 19:28:23 +0100 Subject: [PATCH 2/5] Test something --- .../protocols/openmm_md/plain_md_methods.py | 16 +++++----- .../openmm/plain_md/test_simulation_slow.py | 29 ++++++++++++++++--- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index cac69aee0..cd83c8729 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -288,7 +288,7 @@ def _save_simulation_pdb( """ state = simulation.context.getState( getPositions=True, - enforcePeriodicBox=False, + enforcePeriodicBox=True, ) positions = to_openmm(from_openmm(state.getPositions())) @@ -441,6 +441,8 @@ def _run_MD(simulation: openmm.app.Simulation, "timestep", ) + simulation.context.setStepCount(0) + checkpoint_interval = settings_validation.get_simsteps( sim_length=output_settings.checkpoint_interval, timestep=timestep, @@ -481,12 +483,12 @@ def _run_MD(simulation: openmm.app.Simulation, simulation.step(prod_steps) t1 = time.time() - if output_settings.production_trajectory_filename: - state = simulation.context.getState( - getPositions=True, - enforcePeriodicBox=False, - ) - xtc_reporter.report(simulation, state) + #if output_settings.production_trajectory_filename: + # state = simulation.context.getState( + # getPositions=True, + # enforcePeriodicBox=True, + # ) + # xtc_reporter.report(simulation, state) # Write the final production frame if output_settings.production_structure is not None: diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py index 471208f85..d2cdfc60e 100644 --- a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py +++ b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py @@ -96,6 +96,7 @@ def test_complex_solvent_sim_gpu( 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) @@ -148,7 +149,27 @@ def test_complex_solvent_sim_gpu( assert pur.outputs['production_pdb'] == unit_shared / "production.pdb" # Check the final trajectory frame - u = mda.Universe(pur.outputs['production_pdb']) - u2 = mda.Universe(pur.outputs['minimized_pdb'], pur.outputs['nc']) - u2.trajectory[-1] - assert_allclose(u.atoms.positions, u2.atoms.positions, rtol=0, atol=1e-2) + 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 first frame matches + # assert_allclose( + # first_frame.atoms.positions, + # simulation.atoms.positions, + # rtol=0, + # atol=1e-2, # The PDBs are written at 2d.p. accuracy + # ) + + # 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 + ) + From 4be388a0179be088f9370028c4a3562ef2740617 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Wed, 28 May 2025 10:10:37 +0100 Subject: [PATCH 3/5] Fix PBC writing and reset step counter at start of production phase --- openfe/protocols/openmm_md/plain_md_methods.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index cd83c8729..2a7c0f28e 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -288,7 +288,7 @@ def _save_simulation_pdb( """ state = simulation.context.getState( getPositions=True, - enforcePeriodicBox=True, + enforcePeriodicBox=simulation._usesPBC, ) positions = to_openmm(from_openmm(state.getPositions())) @@ -433,6 +433,9 @@ def _run_MD(simulation: openmm.app.Simulation, 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, @@ -441,8 +444,6 @@ def _run_MD(simulation: openmm.app.Simulation, "timestep", ) - simulation.context.setStepCount(0) - checkpoint_interval = settings_validation.get_simsteps( sim_length=output_settings.checkpoint_interval, timestep=timestep, @@ -459,12 +460,14 @@ def _run_MD(simulation: openmm.app.Simulation, 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), @@ -479,17 +482,11 @@ def _run_MD(simulation: openmm.app.Simulation, density=True, speed=True, )) + t0 = time.time() simulation.step(prod_steps) t1 = time.time() - #if output_settings.production_trajectory_filename: - # state = simulation.context.getState( - # getPositions=True, - # enforcePeriodicBox=True, - # ) - # xtc_reporter.report(simulation, state) - # Write the final production frame if output_settings.production_structure is not None: PlainMDProtocolUnit._save_simulation_pdb( From 15c9ae13708f2a6678ec32f95c0268408f5b9329 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Wed, 28 May 2025 10:40:27 +0100 Subject: [PATCH 4/5] Add news entry --- news/md_outputs.rst | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 news/md_outputs.rst diff --git a/news/md_outputs.rst b/news/md_outputs.rst new file mode 100644 index 000000000..78ad8efb3 --- /dev/null +++ b/news/md_outputs.rst @@ -0,0 +1,26 @@ +**Added:** + +* + +**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:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* From a44d62ce2e63f77ba32bb7ec64e115316550e132 Mon Sep 17 00:00:00 2001 From: IAlibay Date: Wed, 28 May 2025 10:43:53 +0100 Subject: [PATCH 5/5] remove commented out code --- .../protocols/openmm/plain_md/test_simulation_slow.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py index d2cdfc60e..9958339c1 100644 --- a/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py +++ b/openfe/tests/protocols/openmm/plain_md/test_simulation_slow.py @@ -156,14 +156,6 @@ def test_complex_solvent_sim_gpu( # Check we have the right number of frames assert len(simulation.trajectory) == 10 - # # Check that the first frame matches - # assert_allclose( - # first_frame.atoms.positions, - # simulation.atoms.positions, - # rtol=0, - # atol=1e-2, # The PDBs are written at 2d.p. accuracy - # ) - # Check that the final frame matches simulation.trajectory[-1] # fast-forward assert_allclose(