Move things to utils, test interchange creation

openfe_gromacs/__init__.py

@@ -1,3 +1,5 @@
+from importlib.metadata import version
+
 from gufe import (
     AlchemicalNetwork,
     ChemicalSystem,
@@ -19,3 +21,5 @@
     ProtocolUnitResult,
     execute_DAG,
 )
+
+__version__ = version("openfe_gromacs")

openfe_gromacs/protocols/gromacs_md/md_methods.py

@@ -15,28 +15,15 @@
 import pathlib
 import subprocess
 import uuid
-import warnings
 from collections import defaultdict
 from collections.abc import Iterable
 from typing import Any
 
-import gmxapi as gmx
 import gufe
-import pint
-from gufe import ChemicalSystem, SmallMoleculeComponent, settings
-from openfe.protocols.openmm_utils import (
-    charge_generation,
-    settings_validation,
-    system_creation,
-    system_validation,
-)
-from openfe.protocols.openmm_utils.omm_settings import BasePartialChargeSettings
-from openfe.utils import log_system_probe, without_oechem_backend
-from openff.interchange import Interchange
-from openff.toolkit.topology import Molecule as OFFMolecule
+from gufe import ChemicalSystem, settings
+from openfe.protocols.openmm_utils import system_validation
+from openfe.utils import log_system_probe
 from openff.units import unit
-from openff.units.openmm import from_openmm, to_openmm
-from openmmtools import forces
 
 from openfe_gromacs.protocols.gromacs_md.md_settings import (
     EMOutputSettings,
@@ -50,90 +37,12 @@
     NVTOutputSettings,
     NVTSimulationSettings,
     OpenFFPartialChargeSettings,
-    OpenMMSolvationSettings,
+    SolvationSettings,
 )
+from openfe_gromacs.protocols.gromacs_utils import create_systems, write_mdp
 
 logger = logging.getLogger(__name__)
 
-# Settings that are not exposed to the user
-PRE_DEFINED_SETTINGS = {
-    "tinit": 0 * unit.picosecond,
-    "init_step": 0,
-    "simulation_part": 0,
-    "comm_mode": "Linear",
-    "nstcomm": 100,
-    "comm_grps": "system",
-    "pbc": "xyz",
-    "verlet_buffer_tolerance": 0.005 * unit.kilojoule / (unit.mole * unit.picosecond),
-    "verlet_buffer_pressure_tolerance": 0.5 * unit.bar,
-    "coulomb_modifier": "Potential-shift",
-    "epsilon_r": 1,
-    "epsilon_rf": 0,
-    "fourierspacing": 0.12 * unit.nanometer,
-    "lj_pme_comb_rule": "Geometric",
-    "ewald_geometry": "3d",
-    "epsilon_surface": 0,
-    "lincs_warnangle": 30 * unit.degree,
-    "continuation": "no",
-    "morse": "no",
-}
-
-PRE_DEFINED_SETTINGS_EM = {
-    "emtol": 10.0 * unit.kilojoule / (unit.mole * unit.nanometer),
-    "emstep": 0.01 * unit.nanometer,
-}
-
-PRE_DEFINED_SETTINGS_MD = {
-    "nsttcouple": -1,
-    "tc_grps": "system",
-    "tau_t": 2.0 * unit.picosecond,
-    "pcoupltype": "isotropic",
-    "nstpcouple": -1,
-    "tau_p": 5 * unit.picosecond,
-    "compressibility": 4.5e-05 / unit.bar,
-}
-
-
-def _dict2mdp(settings_dict: dict, shared_basepath):
-    """
-    Write out a Gromacs .mdp file given a settings dictionary
-    :param settings_dict: dict
-          Dictionary of settings
-    :param shared_basepath: Pathlike
-          Where to save the .mdp files
-    """
-    filename = shared_basepath / settings_dict["mdp_file"]
-    # Remove non-mdp settings from the dictionary
-    non_mdps = [
-        "forcefield_cache",
-        "mdp_file",
-        "tpr_file",
-        "trr_file",
-        "xtc_file",
-        "gro_file",
-        "edr_file",
-        "log_file",
-        "cpt_file",
-    ]
-    for setting in non_mdps:
-        settings_dict.pop(setting)
-    with open(filename, "w") as f:
-        for key, value in settings_dict.items():
-            # First convert units to units in the mdp file, then remove units
-            if isinstance(value, pint.Quantity):
-                if value.is_compatible_with(unit.nanometer):
-                    value = value.to(unit.nanometer)
-                if value.is_compatible_with(unit.picosecond):
-                    value = value.to(unit.picosecond)
-                if value.is_compatible_with(unit.kelvin):
-                    value = value.to(unit.kelvin)
-                if value.is_compatible_with(unit.bar):
-                    value = value.to(unit.bar)
-                value = value.magnitude
-            # Write out all the setting, value pairs
-            f.write(f"{key} = {value}\n")
-    return filename
-
 
 class GromacsMDProtocolResult(gufe.ProtocolResult):
     """
@@ -475,7 +384,7 @@ def _default_settings(cls):
                 pressure=1 * unit.bar,
             ),
             partial_charge_settings=OpenFFPartialChargeSettings(),
-            solvation_settings=OpenMMSolvationSettings(),
+            solvation_settings=SolvationSettings(),
             integrator_settings=IntegratorSettings(),
             simulation_settings_em=EMSimulationSettings(
                 integrator="steep",
@@ -655,32 +564,6 @@ def __init__(
             generation=generation,
         )
 
-    @staticmethod
-    def _assign_partial_charges(
-        charge_settings: OpenFFPartialChargeSettings,
-        smc_components: dict[SmallMoleculeComponent, OFFMolecule],
-    ) -> None:
-        """
-        Assign partial charges to SMCs.
-
-        Parameters
-        ----------
-        charge_settings : OpenFFPartialChargeSettings
-          Settings for controlling how the partial charges are assigned.
-        smc_components : dict[SmallMoleculeComponent, openff.toolkit.Molecule]
-          Dictionary of OpenFF Molecules to add, keyed by
-          SmallMoleculeComponent.
-        """
-        for mol in smc_components.values():
-            charge_generation.assign_offmol_partial_charges(
-                offmol=mol,
-                overwrite=False,
-                method=charge_settings.partial_charge_method,
-                toolkit_backend=charge_settings.off_toolkit_backend,
-                generate_n_conformers=charge_settings.number_of_conformers,
-                nagl_model=charge_settings.nagl_model,
-            )
-
     def _handle_settings(self) -> dict[str, gufe.settings.SettingsBaseModel]:
         """
         Extract the relevant settings for an MD simulation.
@@ -720,161 +603,6 @@ def _handle_settings(self) -> dict[str, gufe.settings.SettingsBaseModel]:
 
         return settings
 
-    def _write_mdp_files(self, settings: dict, shared_basepath):
-        """
-        Writes out the .mdp files for running a Gromacs MD simulation.
-
-        Parameters
-        ----------
-        settings: dict
-          Dictionary of all the settings
-        shared_basepath : Pathlike, optional
-          Where to run the calculation, defaults to current working directory
-
-        Returns
-        -------
-        mdps: dict
-          Dictionary of file paths to mdp files.
-        """
-
-        mdps = {}
-        if settings["sim_settings_em"].nsteps > 0:
-            settings_dict = (
-                settings["sim_settings_em"].dict()
-                | settings["output_settings_em"].dict()
-                | PRE_DEFINED_SETTINGS
-                | PRE_DEFINED_SETTINGS_EM
-            )
-            mdp = _dict2mdp(settings_dict, shared_basepath)
-            mdps["em"] = mdp
-        if settings["sim_settings_nvt"].nsteps > 0:
-            settings_dict = (
-                settings["sim_settings_nvt"].dict()
-                | settings["output_settings_nvt"].dict()
-                | PRE_DEFINED_SETTINGS
-                | PRE_DEFINED_SETTINGS_MD
-            )
-            mdp = _dict2mdp(settings_dict, shared_basepath)
-            mdps["nvt"] = mdp
-        if settings["sim_settings_npt"].nsteps > 0:
-            settings_dict = (
-                settings["sim_settings_npt"].dict()
-                | settings["output_settings_npt"].dict()
-                | PRE_DEFINED_SETTINGS
-                | PRE_DEFINED_SETTINGS_MD
-            )
-            mdp = _dict2mdp(settings_dict, shared_basepath)
-            mdps["npt"] = mdp
-
-        return mdps
-
-    def _create_interchange(self, settings, components, shared_basepath):
-        """
-        Creates the Interchange object for the system.
-
-        Parameters
-        ----------
-        settings: dict
-          Dictionary of all the settings
-        components: dict
-          Dictionary of the components of the settings, solvent, protein, and
-          small molecules
-        shared_basepath : Pathlike, optional
-          Where to run the calculation, defaults to current working directory
-
-        Returns
-        -------
-        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]
-
-        # ToDo: Make this more general, check if there is a smc_component,
-        #       allow ProteinComponents, ...
-        smc_components = {i: i.to_openff() for i in components["small_mols"]}
-
-        # a. assign partial charges to smcs
-        self._assign_partial_charges(settings["charge_settings"], smc_components)
-
-        # b. get a system generator
-        if settings["output_settings_em"].forcefield_cache is not None:
-            ffcache = shared_basepath / settings["output_settings_em"].forcefield_cache
-        else:
-            ffcache = None
-
-        # Note: we block out the oechem backend for all systemgenerator
-        # linked operations to avoid any smiles operations that can
-        # go wrong when doing rdkit->OEchem roundtripping
-        with without_oechem_backend():
-            system_generator = system_creation.get_system_generator(
-                forcefield_settings=settings["forcefield_settings"],
-                integrator_settings=settings["integrator_settings"],
-                thermo_settings=settings["thermo_settings"],
-                cache=ffcache,
-                has_solvent=components["solvent_comp"] is not None,
-            )
-
-            # Force creation of smc templates so we can solvate later
-            for mol in smc_components.values():
-                system_generator.create_system(
-                    mol.to_topology().to_openmm(), molecules=[mol]
-                )
-
-            # c. get OpenMM Modeller + a resids dictionary for each component
-            stateA_modeller, comp_resids = system_creation.get_omm_modeller(
-                protein_comp=components["protein_comp"],
-                solvent_comp=components["solvent_comp"],
-                small_mols=smc_components,
-                omm_forcefield=system_generator.forcefield,
-                solvent_settings=settings["solvation_settings"],
-            )
-
-            # d. get topology & positions
-            # Note: roundtrip positions to remove vec3 issues
-            stateA_topology = stateA_modeller.getTopology()
-            stateA_positions = to_openmm(from_openmm(stateA_modeller.getPositions()))
-
-            # e. create the stateA System
-            stateA_system = system_generator.create_system(
-                stateA_topology,
-                molecules=[s.to_openff() for s in components["small_mols"]],
-            )
-
-            # 3. Create the Interchange object
-            barostat_idx, barostat = forces.find_forces(
-                stateA_system, ".*Barostat.*", only_one=True
-            )
-            stateA_system.removeForce(barostat_idx)
-
-            stateA_interchange = Interchange.from_openmm(
-                topology=stateA_topology,
-                system=stateA_system,
-                positions=stateA_positions,
-            )
-
-            for molecule_index, molecule in enumerate(
-                stateA_interchange.topology.molecules
-            ):
-                for atom in molecule.atoms:
-                    atom.metadata["residue_number"] = molecule_index + 1
-                if len([*smc_components.values()]) > 0:
-                    if molecule.n_atoms == [*smc_components.values()][0].n_atoms:
-                        for atom in molecule.atoms:
-                            atom.metadata["residue_name"] = "UNK"
-
-        return stateA_interchange
-
     def run(
         self, *, dry=False, verbose=True, scratch_basepath=None, shared_basepath=None
     ) -> dict[str, Any]:
@@ -923,18 +651,33 @@ def run(
         solvent_comp, protein_comp, small_mols = system_validation.get_components(
             stateA
         )
-        components = {
-            "solvent_comp": solvent_comp,
-            "protein_comp": protein_comp,
-            "small_mols": small_mols,
-        }
 
         # 1. Write out .mdp files
-        mdps = self._write_mdp_files(settings, shared_basepath)
+        mdps = write_mdp.write_mdp_files(settings, shared_basepath)
 
-        # 2. Create the Interchange object
-        stateA_interchange = self._create_interchange(
-            settings, components, shared_basepath
+        # 2. Create the OpenMM system
+
+        # Create a dictionary of OFFMol for each SMC for bookkeeping
+        smc_components: dict[gufe.SmallMoleculeComponent, OFFMolecule]
+        smc_components = {i: i.to_openff() for i in small_mols}
+
+        (
+            stateA_system,
+            stateA_topology,
+            stateA_positions,
+        ) = create_systems.create_openmm_system(
+            settings,
+            solvent_comp,
+            protein_comp,
+            smc_components,
+            shared_basepath,
+        )
+        # 3. Create the Interchange object
+        stateA_interchange = create_systems.create_interchange(
+            stateA_system,
+            stateA_topology,
+            stateA_positions,
+            smc_components,
         )
 
         # 4. Save .gro and .top file of the entire system