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:** + +* diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 51613b761..2a7c0f28e 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=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,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 ) @@ -418,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, @@ -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}} 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_md/test_plain_md_slow.py b/openfe/tests/protocols/openmm_md/test_plain_md_slow.py index ecb7bc799..9958339c1 100644 --- a/openfe/tests/protocols/openmm_md/test_plain_md_slow.py +++ b/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 @@ 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) @@ -127,6 +130,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 +146,22 @@ 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 + 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 + ) +