diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py index df5bcc06c..d08ad3a71 100644 --- a/openfe/protocols/openmm_afe/base.py +++ b/openfe/protocols/openmm_afe/base.py @@ -168,7 +168,7 @@ def _pre_equilibrate( Parameters ---------- system : openmm.System - An OpenMM System to equilibrate. + The OpenMM System to equilibrate. topology : openmm.app.Topology OpenMM Topology of the System. positions : openmm.unit.Quantity @@ -502,7 +502,7 @@ def _get_omm_objects( topology : app.Topology OpenMM Topology object describing the parameterized system. system : openmm.System - An non-alchemical OpenMM System of the simulated system. + A non-alchemical OpenMM System of the simulated system. positions : openmm.unit.Quantity Positions of the system. comp_resids : dict[Component, npt.NDArray] diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index a5203263d..51613b761 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -19,18 +19,17 @@ from openff.units.openmm import from_openmm, to_openmm import openmm.unit as omm_unit from typing import Optional -from openmm import app import pathlib from typing import Any, Iterable import uuid import time -import numpy as np import mdtraj from mdtraj.reporters import XTCReporter from openfe.utils import without_oechem_backend, log_system_probe from gufe import ( - settings, ChemicalSystem, SmallMoleculeComponent, - ProteinComponent, SolventComponent + settings, + ChemicalSystem, + SmallMoleculeComponent, ) from openfe.protocols.openmm_utils.omm_settings import ( BasePartialChargeSettings, @@ -684,8 +683,20 @@ def run(self, *, dry=False, verbose=True, 'nc': shared_basepath / output_settings.production_trajectory_filename, 'last_checkpoint': shared_basepath / output_settings.checkpoint_storage_filename, } - if output_settings.equil_nvt_structure: + # The checkpoint file can not exist if frequency > sim length + if not output['last_checkpoint'].exists(): + output['last_checkpoint'] = None + + # The NVT PDB can be ommitted if we don't run the simulation + # Note: we could also just check the file exist + if ( + output_settings.equil_nvt_structure + and sim_settings.equilibration_length_nvt is not None + ): output['nvt_equil_pdb'] = shared_basepath / output_settings.equil_nvt_structure + else: + output['nvt_equil_pdb'] = None + if output_settings.equil_npt_structure: output['npt_equil_pdb'] = shared_basepath / output_settings.equil_npt_structure diff --git a/openfe/tests/protocols/conftest.py b/openfe/tests/protocols/conftest.py index 67cc21d1c..d705af40f 100644 --- a/openfe/tests/protocols/conftest.py +++ b/openfe/tests/protocols/conftest.py @@ -5,10 +5,19 @@ from importlib import resources from rdkit import Chem from rdkit.Geometry import Point3D +from openmm import Platform import openfe from openff.units import unit +@pytest.fixture +def available_platforms() -> set[str]: + return { + Platform.getPlatform(i).getName() + for i in range(Platform.getNumPlatforms()) + } + + @pytest.fixture def benzene_vacuum_system(benzene_modifications): return openfe.ChemicalSystem( diff --git a/openfe/tests/protocols/openmm_ahfe/__init__.py b/openfe/tests/protocols/openmm_ahfe/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openfe/tests/protocols/test_openmm_afe_solvation_protocol.py b/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py similarity index 100% rename from openfe/tests/protocols/test_openmm_afe_solvation_protocol.py rename to openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py diff --git a/openfe/tests/protocols/test_openmm_afe_slow.py b/openfe/tests/protocols/openmm_ahfe/test_ahfe_slow.py similarity index 81% rename from openfe/tests/protocols/test_openmm_afe_slow.py rename to openfe/tests/protocols/openmm_ahfe/test_ahfe_slow.py index 37b5b83aa..bde8ec0af 100644 --- a/openfe/tests/protocols/test_openmm_afe_slow.py +++ b/openfe/tests/protocols/openmm_ahfe/test_ahfe_slow.py @@ -4,42 +4,21 @@ from gufe.protocols import execute_DAG import pytest from openff.units import unit -from openmm import Platform -import os import pathlib import openfe from openfe.protocols import openmm_afe -@pytest.fixture -def available_platforms() -> set[str]: - return {Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())} - - -@pytest.fixture -def set_openmm_threads_1(): - # for vacuum sims, we want to limit threads to one - # this fixture sets OPENMM_CPU_THREADS='1' for a single test, then reverts to previously held value - previous: str | None = os.environ.get('OPENMM_CPU_THREADS') - - try: - os.environ['OPENMM_CPU_THREADS'] = '1' - yield - finally: - if previous is None: - del os.environ['OPENMM_CPU_THREADS'] - else: - os.environ['OPENMM_CPU_THREADS'] = previous - - @pytest.mark.integration # takes too long to be a slow test ~ 4 mins locally @pytest.mark.flaky(reruns=3) # pytest-rerunfailures; we can get bad minimisation @pytest.mark.parametrize('platform', ['CPU', 'CUDA']) -def test_openmm_run_engine(platform, - available_platforms, - benzene_modifications, - set_openmm_threads_1, tmpdir): +def test_openmm_run_engine( + platform, + available_platforms, + benzene_modifications, + tmpdir +): if platform not in available_platforms: pytest.skip(f"OpenMM Platform: {platform} not available") diff --git a/openfe/tests/protocols/test_solvation_afe_tokenization.py b/openfe/tests/protocols/openmm_ahfe/test_ahfe_tokenization.py similarity index 100% rename from openfe/tests/protocols/test_solvation_afe_tokenization.py rename to openfe/tests/protocols/openmm_ahfe/test_ahfe_tokenization.py diff --git a/openfe/tests/protocols/openmm_md/__init__.py b/openfe/tests/protocols/openmm_md/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openfe/tests/protocols/test_openmm_plain_md_protocols.py b/openfe/tests/protocols/openmm_md/test_plain_md_protocol.py similarity index 97% rename from openfe/tests/protocols/test_openmm_plain_md_protocols.py rename to openfe/tests/protocols/openmm_md/test_plain_md_protocol.py index 44a7018b1..edd731b82 100644 --- a/openfe/tests/protocols/test_openmm_plain_md_protocols.py +++ b/openfe/tests/protocols/openmm_md/test_plain_md_protocol.py @@ -434,8 +434,13 @@ def solvent_protocol_dag(benzene_system): def test_unit_tagging(solvent_protocol_dag, tmpdir): # test that executing the Units includes correct generation and repeat info dag_units = solvent_protocol_dag.protocol_units - with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', - return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}): + with mock.patch( + 'openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', + return_value={ + 'nc': 'simulation.xtc', + 'last_checkpoint': 'checkpoint.chk' + } + ): results = [] for u in dag_units: ret = u.execute(context=gufe.Context(tmpdir, tmpdir)) @@ -452,8 +457,13 @@ def test_unit_tagging(solvent_protocol_dag, tmpdir): def test_gather(solvent_protocol_dag, tmpdir): # check .gather behaves as expected - with mock.patch('openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', - return_value={'nc': 'file.nc', 'last_checkpoint': 'chk.nc'}): + with mock.patch( + 'openfe.protocols.openmm_md.plain_md_methods.PlainMDProtocolUnit.run', + return_value={ + 'nc': 'simulation.xtc', + 'last_checkpoint': 'checkpoint.chk' + } + ): dagres = gufe.protocols.execute_DAG(solvent_protocol_dag, shared_basedir=tmpdir, scratch_basedir=tmpdir, diff --git a/openfe/tests/protocols/openmm_md/test_plain_md_slow.py b/openfe/tests/protocols/openmm_md/test_plain_md_slow.py new file mode 100644 index 000000000..ecb7bc799 --- /dev/null +++ b/openfe/tests/protocols/openmm_md/test_plain_md_slow.py @@ -0,0 +1,144 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +import pathlib +import pytest +from openff.units import unit +from gufe.protocols import execute_DAG +from openfe.protocols import openmm_md + + +@pytest.mark.integration +@pytest.mark.parametrize('platform', ['CPU', 'CUDA']) +def test_vacuum_sim( + benzene_vacuum_system, + platform, + available_platforms, + tmpdir +): + if platform not in available_platforms: + pytest.skip(f"OpenMM Platform: {platform} is not available") + + # Run a vacuum MD simulation and check what files we get. + settings = openmm_md.PlainMDProtocol.default_settings() + settings.simulation_settings.equilibration_length_nvt = None + settings.simulation_settings.equilibration_length = 10 * unit.picosecond + settings.simulation_settings.production_length = 20 * unit.picosecond + settings.output_settings.checkpoint_interval = 40 * unit.picosecond + settings.forcefield_settings.nonbonded_method = "nocutoff" + settings.engine_settings.compute_platform = platform + + prot = openmm_md.PlainMDProtocol(settings) + + dag = prot.create( + stateA=benzene_vacuum_system, + stateB=benzene_vacuum_system, + mapping=None, + ) + + workdir = pathlib.Path(str(tmpdir)) + + r = execute_DAG( + dag, + shared_basedir=workdir, + scratch_basedir=workdir, + keep_shared=True + ) + + assert r.ok() + + assert len(r.protocol_unit_results) == 1 + + pur = r.protocol_unit_results[0] + unit_shared = tmpdir / f"shared_{pur.source_key}_attempt_0" + assert unit_shared.exists() + assert pathlib.Path(unit_shared).is_dir() + + # check the files + files = [ + "equil_npt.pdb", + "minimized.pdb", + "simulation.xtc", + "simulation.log", + "system.pdb" + ] + for file in files: + assert (unit_shared / file).exists() + + # NVT PDB should not exist + assert not (unit_shared / "equil_nvt.pdb").exists() + assert not (unit_shared / "checkpoint.chk").exists() + + # check that the output file paths are correct + assert pur.outputs['system_pdb'] == unit_shared / "system.pdb" + assert pur.outputs['minimized_pdb'] == unit_shared / "minimized.pdb" + assert pur.outputs['nc'] == unit_shared / "simulation.xtc" + assert pur.outputs['last_checkpoint'] is None + assert pur.outputs['npt_equil_pdb'] == unit_shared / "equil_npt.pdb" + assert pur.outputs['nvt_equil_pdb'] is None + + +@pytest.mark.integration +@pytest.mark.parametrize('platform', ['CUDA']) +def test_complex_solvent_sim_gpu( + benzene_complex_system, + platform, + available_platforms, + tmpdir, +): + if platform not in available_platforms: + pytest.skip(f"OpenMM Platform: {platform} is not available") + + # Run an MD simulation and check what files we get. + settings = openmm_md.PlainMDProtocol.default_settings() + settings.simulation_settings.equilibration_length_nvt = 50 * unit.picosecond + 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.engine_settings.compute_platform = platform + + prot = openmm_md.PlainMDProtocol(settings) + + dag = prot.create( + stateA=benzene_complex_system, + stateB=benzene_complex_system, + mapping=None, + ) + + workdir = pathlib.Path(str(tmpdir)) + + r = execute_DAG( + dag, + shared_basedir=workdir, + scratch_basedir=workdir, + keep_shared=True + ) + + assert r.ok() + + assert len(r.protocol_unit_results) == 1 + + pur = r.protocol_unit_results[0] + unit_shared = tmpdir / f"shared_{pur.source_key}_attempt_0" + assert unit_shared.exists() + assert pathlib.Path(unit_shared).is_dir() + + # check the files + files = [ + "checkpoint.chk", + "equil_nvt.pdb", + "equil_npt.pdb", + "minimized.pdb", + "simulation.xtc", + "simulation.log", + "system.pdb" + ] + for file in files: + assert (unit_shared / file).exists() + + # check that the output file paths are correct + assert pur.outputs['system_pdb'] == unit_shared / "system.pdb" + assert pur.outputs['minimized_pdb'] == unit_shared / "minimized.pdb" + assert pur.outputs['nc'] == unit_shared / "simulation.xtc" + 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" diff --git a/openfe/tests/protocols/openmm_md/test_plain_md_tokenization.py b/openfe/tests/protocols/openmm_md/test_plain_md_tokenization.py new file mode 100644 index 000000000..1a7cf95ae --- /dev/null +++ b/openfe/tests/protocols/openmm_md/test_plain_md_tokenization.py @@ -0,0 +1,82 @@ +# This ccode is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +import json +import pytest +import gufe +from gufe.tests.test_tokenization import GufeTokenizableTestsMixin +from openfe.protocols import openmm_md + + +@pytest.fixture +def protocol(): + return openmm_md.PlainMDProtocol( + openmm_md.PlainMDProtocol.default_settings() + ) + + +@pytest.fixture +def protocol_unit(protocol, benzene_system): + pus = protocol.create( + stateA=benzene_system, + stateB=benzene_system, + mapping=None, + ) + return list(pus.protocol_units)[0] + + +@pytest.fixture +def protocol_result(md_json): + d = json.loads(md_json, cls=gufe.tokenization.JSON_HANDLER.decoder) + pr = gufe.ProtocolResult.from_dict(d['protocol_result']) + return pr + + +class TestPlainMDProtocol(GufeTokenizableTestsMixin): + cls = openmm_md.PlainMDProtocol + key = None + repr = "PlainMDProtocol-" + + @pytest.fixture() + def instance(self, protocol): + return protocol + + def test_repr(self, instance): + """ + Overwrites the base `test_repr` call to do a bit more. + """ + assert isinstance(repr(instance), str) + assert self.repr in repr(instance) + + +class TestPlainMDProtocolUnit(GufeTokenizableTestsMixin): + cls = openmm_md.PlainMDProtocolUnit + repr = "PlainMDProtocolUnit(" + key = None + + @pytest.fixture + def instance(self, protocol_unit): + return protocol_unit + + def test_repr(self, instance): + """ + Overwrites the base `test_repr` call to do a bit more. + """ + assert isinstance(repr(instance), str) + assert self.repr in repr(instance) + + +class TestPlainMDProtocolResult(GufeTokenizableTestsMixin): + cls = openmm_md.PlainMDProtocolResult + key = None + repr = "PlainMDProtocolResult-" + + @pytest.fixture() + def instance(self, protocol_result): + return protocol_result + + def test_repr(self, instance): + """ + Overwrites the base `test_repr` call to do a bit more. + """ + assert isinstance(repr(instance), str) + assert self.repr in repr(instance) \ No newline at end of file diff --git a/openfe/tests/protocols/openmm_rfe/__init__.py b/openfe/tests/protocols/openmm_rfe/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py similarity index 100% rename from openfe/tests/protocols/test_openmm_equil_rfe_protocols.py rename to openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py diff --git a/openfe/tests/protocols/test_openmm_rfe_slow.py b/openfe/tests/protocols/openmm_rfe/test_hybrid_top_slow.py similarity index 88% rename from openfe/tests/protocols/test_openmm_rfe_slow.py rename to openfe/tests/protocols/openmm_rfe/test_hybrid_top_slow.py index f0fe5a0a7..ffad3f501 100644 --- a/openfe/tests/protocols/test_openmm_rfe_slow.py +++ b/openfe/tests/protocols/openmm_rfe/test_hybrid_top_slow.py @@ -5,41 +5,22 @@ from gufe.protocols import execute_DAG import pytest from openff.units import unit -from openmm import Platform -import os import pathlib import openfe from openfe.protocols import openmm_rfe -@pytest.fixture -def available_platforms() -> set[str]: - return {Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())} - - -@pytest.fixture -def set_openmm_threads_1(): - # for vacuum sims, we want to limit threads to one - # this fixture sets OPENMM_CPU_THREADS='1' for a single test, then reverts to previously held value - previous: str | None = os.environ.get('OPENMM_CPU_THREADS') - - try: - os.environ['OPENMM_CPU_THREADS'] = '1' - yield - finally: - if previous is None: - del os.environ['OPENMM_CPU_THREADS'] - else: - os.environ['OPENMM_CPU_THREADS'] = previous - - @pytest.mark.slow -@pytest.mark.flaky(reruns=3) # pytest-rerunfailures; we can get bad minimisation +@pytest.mark.flaky(reruns=3) # pytest-rerunfailures; we can get bad minimization @pytest.mark.parametrize('platform', ['CPU', 'CUDA']) -def test_openmm_run_engine(benzene_vacuum_system, platform, - available_platforms, benzene_modifications, - set_openmm_threads_1, tmpdir): +def test_openmm_run_engine( + benzene_vacuum_system, + platform, + available_platforms, + benzene_modifications, + tmpdir +): if platform not in available_platforms: pytest.skip(f"OpenMM Platform: {platform} not available") # this test actually runs MD diff --git a/openfe/tests/protocols/test_rfe_tokenization.py b/openfe/tests/protocols/openmm_rfe/test_hybrid_top_tokenization.py similarity index 100% rename from openfe/tests/protocols/test_rfe_tokenization.py rename to openfe/tests/protocols/openmm_rfe/test_hybrid_top_tokenization.py diff --git a/openfe/tests/protocols/test_openmmutils.py b/openfe/tests/protocols/test_openmmutils.py index 4a1a7dc4f..770b13e6b 100644 --- a/openfe/tests/protocols/test_openmmutils.py +++ b/openfe/tests/protocols/test_openmmutils.py @@ -32,7 +32,7 @@ from openmmtools import multistate from pymbar.utils import ParameterError -from ..conftest import HAS_INTERNET +from openfe.tests.conftest import HAS_INTERNET @pytest.mark.parametrize('padding, number_solv, box_vectors, box_size', [