Merge pull request #43 from OpenFreeEnergy/protocol_updates

Some protocol updates: MDPs dict, vacuum check

openfe_gromacs/protocols/gromacs_md/md_methods.py

@@ -15,6 +15,7 @@
 import pathlib
 import subprocess
 import uuid
+import warnings
 from collections import defaultdict
 from collections.abc import Iterable
 from typing import Any
@@ -54,7 +55,7 @@
 
 logger = logging.getLogger(__name__)
 
-
+# Settings that are not exposed to the user
 PRE_DEFINED_SETTINGS = {
     "tinit": 0 * unit.picosecond,
     "init_step": 0,
@@ -198,16 +199,17 @@ def get_top_filename(self) -> list[pathlib.Path]:
 
         return top
 
-    def get_mdp_filenames(self) -> list[list[pathlib.Path]]:
+    def get_mdp_filenames(self) -> list[dict[str, pathlib.Path]]:
         """
-        Get a list of paths to the .mdp files
+        Get a dictionary of paths to the .mdp files
 
         Returns
         -------
-        mdps : list[list[pathlib.Path]]
-          list of paths (pathlib.Path) to the mdp files for energy minimization,
+        mdps : list[dict[str, pathlib.Path]]
+          dictionary of paths (pathlib.Path) to the mdp files for energy minimization,
           NVT and NPT MD runs
         """
+
         mdps = [
             pus[0].outputs["mdp_files"]
             for pus in self.data.values()
@@ -556,7 +558,16 @@ def _create(
             stateA
         )
 
-        system_name = "Solvent MD" if solvent_comp is not None else "Vacuum MD"
+        # Raise an error when no SolventComponent is provided as this Protocol
+        # currently does not support vacuum simulations
+        if solvent_comp is None:
+            errmsg = (
+                "No SolventComponent provided. This protocol currently does"
+                " not support vacuum simulations."
+            )
+            raise ValueError(errmsg)
+
+        system_name = "Solvent MD"
 
         for comp in [protein_comp] + small_mols:
             if comp is not None:
@@ -729,11 +740,11 @@ def _write_mdp_files(self, settings: dict, shared_basepath):
 
         Returns
         -------
-        mdps: list
-          List of file paths to mdp files.
+        mdps: dict
+          Dictionary of file paths to mdp files.
         """
 
-        mdps = []
+        mdps = {}
         if settings["sim_settings_em"].nsteps > 0:
             settings_dict = (
                 settings["sim_settings_em"].dict()
@@ -742,7 +753,7 @@ def _write_mdp_files(self, settings: dict, shared_basepath):
                 | PRE_DEFINED_SETTINGS_EM
             )
             mdp = _dict2mdp(settings_dict, shared_basepath)
-            mdps.append(mdp)
+            mdps["em"] = mdp
         if settings["sim_settings_nvt"].nsteps > 0:
             settings_dict = (
                 settings["sim_settings_nvt"].dict()
@@ -751,7 +762,7 @@ def _write_mdp_files(self, settings: dict, shared_basepath):
                 | PRE_DEFINED_SETTINGS_MD
             )
             mdp = _dict2mdp(settings_dict, shared_basepath)
-            mdps.append(mdp)
+            mdps["nvt"] = mdp
         if settings["sim_settings_npt"].nsteps > 0:
             settings_dict = (
                 settings["sim_settings_npt"].dict()
@@ -760,7 +771,7 @@ def _write_mdp_files(self, settings: dict, shared_basepath):
                 | PRE_DEFINED_SETTINGS_MD
             )
             mdp = _dict2mdp(settings_dict, shared_basepath)
-            mdps.append(mdp)
+            mdps["npt"] = mdp
 
         return mdps
 
@@ -783,6 +794,15 @@ def _create_interchange(self, settings, components, shared_basepath):
         stateA_interchange: Interchange object
           The interchange object of the system.
         """
+        # Set the environment variable for using the experimental interchange
+        # functionality `from_openmm` and raise a warning
+        os.environ["INTERCHANGE_EXPERIMENTAL"] = "1"
+        war = (
+            "Environment variable INTERCHANGE_EXPERIMENTAL=1 is set for using "
+            "the interchange functionality 'from_openmm' which is not well "
+            "tested yet."
+        )
+        warnings.warn(war)
         # Create the stateA system
         # Create a dictionary of OFFMol for each SMC for bookkeeping
         smc_components: dict[SmallMoleculeComponent, OFFMolecule]
@@ -916,15 +936,6 @@ def run(
             "small_mols": small_mols,
         }
 
-        # Raise an error when no SolventComponent is provided as this Protocol
-        # currently does not support vacuum simulations
-        if not solvent_comp:
-            errmsg = (
-                "No SolventComponent provided. This protocol currently does"
-                " not support vacuum simulations."
-            )
-            raise ValueError(errmsg)
-
         # 1. Write out .mdp files
         mdps = self._write_mdp_files(settings, shared_basepath)
 
@@ -1106,15 +1117,11 @@ def _execute(
         if sim_settings_em.nsteps > 0:
             if verbose:
                 self.logger.info("Running energy minimization")
-            mdp = [
-                x
-                for x in mdp_files
-                if str(x).split("/")[-1] == output_settings_em.mdp_file
-            ]
+            mdp = mdp_files["em"]
             tpr = pathlib.Path(ctx.shared / output_settings_em.tpr_file)
-            assert len(mdp) == 1
+            assert mdp.exists()
             self._run_gromacs(
-                mdp[0],
+                mdp,
                 input_gro,
                 input_top,
                 tpr,
@@ -1133,20 +1140,16 @@ def _execute(
         if sim_settings_nvt.nsteps > 0:
             if verbose:
                 self.logger.info("Running an NVT MD simulation")
-            mdp = [
-                x
-                for x in mdp_files
-                if str(x).split("/")[-1] == output_settings_nvt.mdp_file
-            ]
+            mdp = mdp_files["nvt"]
             tpr = pathlib.Path(ctx.shared / output_settings_nvt.tpr_file)
-            assert len(mdp) == 1
+            assert mdp.exists()
             # If EM was run, use the output from that to run NVT MD
             if sim_settings_em.nsteps > 0:
                 gro = pathlib.Path(ctx.shared / output_settings_em.gro_file)
             else:
                 gro = input_gro
             self._run_gromacs(
-                mdp[0],
+                mdp,
                 gro,
                 input_top,
                 tpr,
@@ -1164,13 +1167,9 @@ def _execute(
         if sim_settings_npt.nsteps > 0:
             if verbose:
                 self.logger.info("Running an NPT MD simulation")
-            mdp = [
-                x
-                for x in mdp_files
-                if str(x).split("/")[-1] == output_settings_npt.mdp_file
-            ]
+            mdp = mdp_files["npt"]
             tpr = pathlib.Path(ctx.shared / output_settings_npt.tpr_file)
-            assert len(mdp) == 1
+            assert mdp.exists()
             # If EM and/or NVT MD was run, use the output coordinate file
             # from that to run NPT MD
             if sim_settings_em.nsteps > 0:
@@ -1181,7 +1180,7 @@ def _execute(
             else:
                 gro = input_gro
             self._run_gromacs(
-                mdp[0],
+                mdp,
                 gro,
                 input_top,
                 tpr,

openfe_gromacs/tests/protocols/test_gromacs_md.py

@@ -80,17 +80,14 @@ def test_no_SolventComponent(benzene_vacuum_system, tmpdir):
         settings=settings,
     )
 
-    dag = p.create(
-        stateA=benzene_vacuum_system,
-        stateB=benzene_vacuum_system,
-        mapping=None,
-    )
-    dag_unit = list(dag.protocol_units)[0]
-
     errmsg = "No SolventComponent provided. This protocol currently"
     with tmpdir.as_cwd():
         with pytest.raises(ValueError, match=errmsg):
-            dag_unit.run(dry=True)
+            p.create(
+                stateA=benzene_vacuum_system,
+                stateB=benzene_vacuum_system,
+                mapping=None,
+            )
 
 
 @pytest.fixture
@@ -135,7 +132,6 @@ def test_unit_tagging_setup_unit(solvent_protocol_dag, tmpdir):
 
 
 def test_dry_run_ffcache_none(benzene_system, monkeypatch, tmp_path_factory):
-    monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1")
     settings = GromacsMDProtocol.default_settings()
     settings.output_settings_em.forcefield_cache = None
     settings.simulation_settings_em.nsteps = 10
@@ -173,33 +169,33 @@ def test_dry_many_molecules_solvent(
     """
     A basic test flushing "will it work if you pass multiple molecules"
     """
-    monkeypatch.setenv("INTERCHANGE_EXPERIMENTAL", "1")
-    settings = GromacsMDProtocol.default_settings()
-    # Only run an EM, no MD, to make the test faster
-    settings.simulation_settings_em.nsteps = 1
-    settings.simulation_settings_nvt.nsteps = 0
-    settings.simulation_settings_npt.nsteps = 0
-    settings.simulation_settings_em.rcoulomb = 1.0 * off_unit.nanometer
-    protocol = GromacsMDProtocol(
-        settings=settings,
-    )
-
-    # create DAG from protocol and take first (and only) work unit from within
-    dag = protocol.create(
-        stateA=benzene_many_solv_system,
-        stateB=benzene_many_solv_system,
-        mapping=None,
-    )
-
-    shared_temp = tmp_path_factory.mktemp("shared")
-    scratch_temp = tmp_path_factory.mktemp("scratch")
-    gufe.protocols.execute_DAG(
-        dag,
-        shared_basedir=shared_temp,
-        scratch_basedir=scratch_temp,
-        keep_shared=False,
-        n_retries=3,
-    )
+    with pytest.warns(UserWarning, match="Environment variabl"):
+        settings = GromacsMDProtocol.default_settings()
+        # Only run an EM, no MD, to make the test faster
+        settings.simulation_settings_em.nsteps = 1
+        settings.simulation_settings_nvt.nsteps = 0
+        settings.simulation_settings_npt.nsteps = 0
+        settings.simulation_settings_em.rcoulomb = 1.0 * off_unit.nanometer
+        protocol = GromacsMDProtocol(
+            settings=settings,
+        )
+
+        # create DAG from protocol and take first (and only) work unit from within
+        dag = protocol.create(
+            stateA=benzene_many_solv_system,
+            stateB=benzene_many_solv_system,
+            mapping=None,
+        )
+
+        shared_temp = tmp_path_factory.mktemp("shared")
+        scratch_temp = tmp_path_factory.mktemp("scratch")
+        gufe.protocols.execute_DAG(
+            dag,
+            shared_basedir=shared_temp,
+            scratch_basedir=scratch_temp,
+            keep_shared=False,
+            n_retries=3,
+        )
 
 
 class TestProtocolResult:
@@ -244,8 +240,8 @@ def test_get_mdp_filenames(self, protocolresult):
         mdps = protocolresult.get_mdp_filenames()
 
         assert isinstance(mdps, list)
-        assert isinstance(mdps[0], list)
-        assert isinstance(mdps[0][0], pathlib.Path)
+        assert isinstance(mdps[0], dict)
+        assert all(isinstance(mdp, pathlib.Path) for mdp in mdps[0].values())
 
     def test_get_filenames_em(self, protocolresult):
         dict_file_path = protocolresult.get_filenames_em()

openfe_gromacs/tests/protocols/test_md_tokenization.py

@@ -64,7 +64,7 @@ def test_key_stable(self):
 
 class TestGromacsMDProtocolResult(GufeTokenizableTestsMixin):
     cls = gromacs_md.GromacsMDProtocolResult
-    key = "GromacsMDProtocolResult-f5a750033d57f5fd2177491cff9046c0"
+    key = "GromacsMDProtocolResult-4739ae9a56f36a1e7c95c503601769f4"
     repr = f"<{key}>"
 
     @pytest.fixture()