diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 342dbda..834451b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,7 +2,7 @@ ci: autoupdate_schedule: "quarterly" repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.4.0 + rev: v4.6.0 hooks: - id: check-yaml - id: debug-statements @@ -12,23 +12,23 @@ repos: - id: check-executables-have-shebangs - id: check-json - repo: https://github.com/psf/black - rev: 22.12.0 + rev: 24.8.0 hooks: - id: black files: ^openfe_gromacs - repo: https://github.com/PyCQA/isort - rev: 5.12.0 + rev: 5.13.2 hooks: - id: isort files: ^openfe_gromacs - repo: https://github.com/PyCQA/flake8 - rev: 6.0.0 + rev: 7.1.1 hooks: - id: flake8 files: ^openfe_gromacs additional_dependencies: [Flake8-pyproject] - repo: https://github.com/asottile/pyupgrade - rev: 'v3.3.1' + rev: 'v3.17.0' hooks: - id: pyupgrade args: diff --git a/environment.yml b/environment.yml index 71824fd..a9643c1 100644 --- a/environment.yml +++ b/environment.yml @@ -9,3 +9,5 @@ dependencies: - pytest - pytest-xdist - pytest-cov + - pip: + - git+https://github.com/openforcefield/openff-interchange.git diff --git a/openfe_gromacs/__init__.py b/openfe_gromacs/__init__.py index feb953d..ad2c10f 100644 --- a/openfe_gromacs/__init__.py +++ b/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") diff --git a/openfe_gromacs/protocols/gromacs_md/md_methods.py b/openfe_gromacs/protocols/gromacs_md/md_methods.py index 9d4efa2..fc72590 100644 --- a/openfe_gromacs/protocols/gromacs_md/md_methods.py +++ b/openfe_gromacs/protocols/gromacs_md/md_methods.py @@ -15,25 +15,16 @@ import pathlib import subprocess import uuid -import warnings from collections import defaultdict from collections.abc import Iterable from typing import Any import gufe -import pint -from gufe import ChemicalSystem, SmallMoleculeComponent, settings -from openfe.protocols.openmm_utils import ( - charge_generation, - system_creation, - system_validation, -) -from openfe.utils import log_system_probe, without_oechem_backend -from openff.interchange import Interchange +from gufe import ChemicalSystem, settings +from openfe.protocols.openmm_utils import system_validation +from openfe.utils import log_system_probe from openff.toolkit.topology import Molecule as OFFMolecule 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, @@ -47,90 +38,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): """ @@ -256,6 +169,7 @@ def get_filenames_em(self) -> dict[str, list[pathlib.Path]]: - "edr_em" - "log_em" - "cpt_em" + - "grompp_mdp_em" Returns ------- @@ -271,6 +185,7 @@ def get_filenames_em(self) -> dict[str, list[pathlib.Path]]: "edr_em", "log_em", "cpt_em", + "grompp_mdp_em", ] dict_npt = {} for file in file_keys: @@ -319,6 +234,7 @@ def get_filenames_nvt(self) -> dict[str, list[pathlib.Path]]: - "edr_nvt" - "log_nvt" - "cpt_nvt" + - "grompp_mdp_nvt" Returns ------- @@ -334,6 +250,7 @@ def get_filenames_nvt(self) -> dict[str, list[pathlib.Path]]: "edr_nvt", "log_nvt", "cpt_nvt", + "grompp_mdp_nvt", ] dict_npt = {} for file in file_keys: @@ -382,6 +299,7 @@ def get_filenames_npt(self) -> dict[str, list[pathlib.Path]]: - "edr_npt" - "log_npt" - "cpt_npt" + - "grompp_mdp_npt" Returns ------- @@ -397,6 +315,7 @@ def get_filenames_npt(self) -> dict[str, list[pathlib.Path]]: "edr_npt", "log_npt", "cpt_npt", + "grompp_mdp_em", ] dict_npt = {} for file in file_keys: @@ -472,7 +391,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", @@ -491,6 +410,7 @@ def _default_settings(cls): engine_settings=GromacsEngineSettings(), output_settings_em=EMOutputSettings( mdp_file="em.mdp", + grompp_mdp_file="mdout_em.mdp", tpr_file="em.tpr", gro_file="em.gro", trr_file="em.trr", @@ -501,6 +421,7 @@ def _default_settings(cls): ), output_settings_nvt=NVTOutputSettings( mdp_file="nvt.mdp", + grompp_mdp_file="mdout_nvt.mdp", tpr_file="nvt.tpr", gro_file="nvt.gro", trr_file="nvt.trr", @@ -511,6 +432,7 @@ def _default_settings(cls): ), output_settings_npt=NPTOutputSettings( mdp_file="npt.mdp", + grompp_mdp_file="mdout_npt.mdp", tpr_file="npt.tpr", gro_file="npt.gro", trr_file="npt.trr", @@ -652,32 +574,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. @@ -717,161 +613,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]: @@ -924,18 +665,38 @@ 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( + solvent_comp, + protein_comp, + smc_components, + settings["charge_settings"], + settings["forcefield_settings"], + settings["integrator_settings"], + settings["thermo_settings"], + settings["solvation_settings"], + settings["output_settings_em"], + 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 @@ -984,6 +745,7 @@ def _run_gromacs( cpt: str, log: str, edr: str, + out_mdp: str, engine_settings: GromacsEngineSettings, shared_basebath: pathlib.Path, ): @@ -992,19 +754,20 @@ def _run_gromacs( Parameters ---------- - :param mdp: pathlib.Path + mdp: pathlib.Path Path to the mdp file - :param in_gro: pathlib.Path - :param top: pathlib.Path - :param tpr: pathlib.Path - :param out_gro: str - :param xtc: str - :param trr: str - :param cpt: str - :param log: str - :param edr: str - :param engine_settings: GromacsEngineSettings - :param shared_basebath: Pathlike, optional + in_gro: pathlib.Path + top: pathlib.Path + tpr: pathlib.Path + out_gro: str + xtc: str + trr: str + cpt: str + log: str + edr: str + out_mdp: str + engine_settings: GromacsEngineSettings + shared_basebath: Pathlike, optional Where to run the calculation, defaults to current working directory """ assert os.path.exists(in_gro) @@ -1022,6 +785,8 @@ def _run_gromacs( top, "-o", tpr, + "-po", + shared_basebath / out_mdp, ], stdin=subprocess.PIPE, ) @@ -1143,6 +908,7 @@ def _execute( output_settings_em.cpt_file, output_settings_em.log_file, output_settings_em.edr_file, + output_settings_em.grompp_mdp_file, engine_settings, ctx.shared, ) @@ -1153,6 +919,9 @@ def _execute( output_dict["edr_em"] = shared_basepath / output_settings_em.edr_file output_dict["log_em"] = shared_basepath / output_settings_em.log_file output_dict["cpt_em"] = shared_basepath / output_settings_em.cpt_file + output_dict["grompp_mdp_em"] = ( + shared_basepath / output_settings_em.grompp_mdp_file + ) # ToDo: Should we disallow running MD without EM? # Run NVT @@ -1178,6 +947,7 @@ def _execute( output_settings_nvt.cpt_file, output_settings_nvt.log_file, output_settings_nvt.edr_file, + output_settings_nvt.grompp_mdp_file, engine_settings, ctx.shared, ) @@ -1188,6 +958,9 @@ def _execute( output_dict["edr_nvt"] = shared_basepath / output_settings_nvt.edr_file output_dict["log_nvt"] = shared_basepath / output_settings_nvt.log_file output_dict["cpt_nvt"] = shared_basepath / output_settings_nvt.cpt_file + output_dict["grompp_mdp_nvt"] = ( + shared_basepath / output_settings_nvt.grompp_mdp_file + ) # Run NPT MD simulation if sim_settings_npt.nsteps > 0: @@ -1216,6 +989,7 @@ def _execute( output_settings_npt.cpt_file, output_settings_npt.log_file, output_settings_npt.edr_file, + output_settings_npt.grompp_mdp_file, engine_settings, ctx.shared, ) @@ -1226,5 +1000,8 @@ def _execute( output_dict["edr_npt"] = shared_basepath / output_settings_npt.edr_file output_dict["log_npt"] = shared_basepath / output_settings_npt.log_file output_dict["cpt_npt"] = shared_basepath / output_settings_npt.cpt_file + output_dict["grompp_mdp_nvt"] = ( + shared_basepath / output_settings_nvt.grompp_mdp_file + ) return output_dict diff --git a/openfe_gromacs/protocols/gromacs_md/md_settings.py b/openfe_gromacs/protocols/gromacs_md/md_settings.py index 4e1d9ea..f6a5e0b 100644 --- a/openfe_gromacs/protocols/gromacs_md/md_settings.py +++ b/openfe_gromacs/protocols/gromacs_md/md_settings.py @@ -146,9 +146,9 @@ class Config: """ # # # Bonds # # # - constraints: Literal[ - "none", "h-bonds", "all-bonds", "h-angles", "all-angles" - ] = "h-bonds" + constraints: Literal["none", "h-bonds", "all-bonds", "h-angles", "all-angles"] = ( + "h-bonds" + ) """ Controls which bonds in the topology will be converted to rigid holonomic constraints. Note that typical rigid water models do not have bonds, but @@ -265,6 +265,12 @@ class OutputSettings(SettingsBaseModel): Filename for the mdp file for running simulations in Gromacs. Default 'em.mdp' """ + grompp_mdp_file: str = "mdout_em.mdp" + """ + Filename for the mdp file that gmx grompp outputs. This file contains + comment lines, as well as the input that gmx grompp has read. + Default 'mdout_em.mdp' + """ tpr_file: str = "em.tpr" """ Filename for the tpr file for running simulations in Gromacs. @@ -665,6 +671,16 @@ def has_no_constraints(cls, v): return v +class SolvationSettings(OpenMMSolvationSettings): + """ + Solvation settings for solvating the system using OpenMM. For now, + water models with virtual sites are not supported when creating an + Interchange object `from_openmm`. + """ + + solvent_model: Literal["tip3p"] = "tip3p" + + class GromacsEngineSettings(SettingsBaseModel): """ MD engine settings for running simulations in Gromacs. @@ -728,7 +744,7 @@ def must_be_positive(cls, v): # Settings for creating the OpenMM systems forcefield_settings: FFSettingsOpenMM partial_charge_settings: OpenFFPartialChargeSettings - solvation_settings: OpenMMSolvationSettings + solvation_settings: SolvationSettings integrator_settings: IntegratorSettings # Simulation run settings diff --git a/openfe_gromacs/protocols/gromacs_utils/create_systems.py b/openfe_gromacs/protocols/gromacs_utils/create_systems.py new file mode 100644 index 0000000..3c38db5 --- /dev/null +++ b/openfe_gromacs/protocols/gromacs_utils/create_systems.py @@ -0,0 +1,185 @@ +import os +import warnings + +import gufe +from openfe.protocols.openmm_utils import charge_generation, system_creation +from openfe.utils import without_oechem_backend +from openff.interchange import Interchange +from openff.toolkit.topology import Molecule as OFFMolecule +from openff.units.openmm import from_openmm, to_openmm +from openmmtools import forces + +from openfe_gromacs.protocols.gromacs_md.md_settings import OpenFFPartialChargeSettings + + +def assign_partial_charges( + charge_settings: OpenFFPartialChargeSettings, + smc_components: dict[gufe.SmallMoleculeComponent, OFFMolecule], +) -> None: + """ + Assign partial charges to SMCs. + + Parameters + ---------- + charge_settings : OpenFFPartialChargeSettings + Settings for controlling how the partial charges are assigned. + smc_components : dict[gufe.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 create_openmm_system( + solvent_comp, + protein_comp, + smc_components, + partial_charge_settings, + forcefield_settings, + integrator_settings, + thermo_settings, + solvation_settings, + output_settings_em, + shared_basepath, +): + """ + Creates the OpenMM system. + + Parameters + ---------- + solvent_comp: gufe.SolventComponent + The SolventComponent for the system + protein_comp: gufe.ProteinComponent + The ProteinComponent for the system + smc_components: dict[gufe.SmallMoleculeComponent, OFFMolecule] + A dictionary with SmallMoleculeComponents as keys and OpenFF molecules + as values + partial_charge_settings: OpenFFPartialChargeSettings + forcefield_settings: FFSettingsOpenMM + integrator_settings: IntegratorSettings + thermo_settings: gufe.settings.ThermoSettings + solvation_settings: SolvationSettings + output_settings_em: EMOutputSettings + shared_basepath : Pathlike, optional + Where to run the calculation, defaults to current working directory + + Returns + ------- + stateA_system: OpenMM system + stateA_topology: OpenMM topology + stateA_positions: Positions + """ + # a. assign partial charges to smcs + assign_partial_charges(partial_charge_settings, smc_components) + + # b. get a system generator + if output_settings_em.forcefield_cache is not None: + ffcache = shared_basepath / 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=forcefield_settings, + integrator_settings=integrator_settings, + thermo_settings=thermo_settings, + cache=ffcache, + has_solvent=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=protein_comp, + solvent_comp=solvent_comp, + small_mols=smc_components, + omm_forcefield=system_generator.forcefield, + solvent_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=smc_components.values(), + ) + return stateA_system, stateA_topology, stateA_positions + + +def create_interchange( + stateA_system, + stateA_topology, + stateA_positions, + smc_components, +): + """ + Creates the Interchange object for the system. + + Parameters + ---------- + stateA_system: OpenMM system + stateA_topology: OpenMM topology + stateA_positions: Positions + smc_components: dict[gufe.SmallMoleculeComponent, OFFMolecule] + A dictionary with SmallMoleculeComponents as keys and OpenFF molecules + as values + + 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) + + try: + barostat_idx, barostat = forces.find_forces( + stateA_system, ".*Barostat.*", only_one=True + ) + stateA_system.removeForce(barostat_idx) + + except forces.NoForceFoundError: + pass + + 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 diff --git a/openfe_gromacs/protocols/gromacs_utils/write_mdp.py b/openfe_gromacs/protocols/gromacs_utils/write_mdp.py new file mode 100644 index 0000000..eed5b76 --- /dev/null +++ b/openfe_gromacs/protocols/gromacs_utils/write_mdp.py @@ -0,0 +1,134 @@ +import pint +from openff.units import unit + +# 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 + + Parameters + ---------- + settings_dict: dict + Dictionary of settings + 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", + "grompp_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 + + +def write_mdp_files(settings: dict, shared_basepath): + """ + Writes out the .mdp files for running a Gromacs 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 diff --git a/openfe_gromacs/tests/conftest.py b/openfe_gromacs/tests/conftest.py index cb1b894..eec6ab7 100644 --- a/openfe_gromacs/tests/conftest.py +++ b/openfe_gromacs/tests/conftest.py @@ -5,13 +5,14 @@ from importlib import resources import gufe -import openfe import pytest from gufe import SmallMoleculeComponent from openff.units import unit from rdkit import Chem from rdkit.Chem import AllChem +import openfe_gromacs + class SlowTests: """Plugin for handling fixtures that skips slow tests @@ -125,7 +126,7 @@ def ethane(): @pytest.fixture(scope="session") def benzene_modifications(): files = {} - with importlib.resources.files("openfe.tests.data") as d: + with importlib.resources.files("openfe_gromacs.tests.data") as d: fn = str(d / "benzene_modifications.sdf") supp = Chem.SDMolSupplier(str(fn), removeHs=False) for rdmol in supp: @@ -135,34 +136,43 @@ def benzene_modifications(): @pytest.fixture(scope="session") def T4_protein_component(): - with resources.files("openfe.tests.data") as d: + with resources.files("openfe_gromacs.tests.data") as d: fn = str(d / "181l_only.pdb") comp = gufe.ProteinComponent.from_pdb_file(fn, name="T4_protein") return comp +@pytest.fixture(scope="session") +def alanine_dipeptide_component(): + with resources.files("openfe_gromacs.tests.data") as d: + fn = str(d / "alanine-dipeptide.pdb") + comp = gufe.ProteinComponent.from_pdb_file(fn, name="Alanine_dipeptide") + + return comp + + @pytest.fixture() def eg5_protein_pdb(): - with resources.files("openfe.tests.data.eg5") as d: + with resources.files("openfe_gromacs.tests.data.eg5") as d: yield str(d / "eg5_protein.pdb") @pytest.fixture() def eg5_ligands_sdf(): - with resources.files("openfe.tests.data.eg5") as d: + with resources.files("openfe_gromacs.tests.data.eg5") as d: yield str(d / "eg5_ligands.sdf") @pytest.fixture() def eg5_cofactor_sdf(): - with resources.files("openfe.tests.data.eg5") as d: + with resources.files("openfe_gromacs.tests.data.eg5") as d: yield str(d / "eg5_cofactor.sdf") @pytest.fixture() -def eg5_protein(eg5_protein_pdb) -> openfe.ProteinComponent: - return openfe.ProteinComponent.from_pdb_file(eg5_protein_pdb) +def eg5_protein(eg5_protein_pdb) -> openfe_gromacs.ProteinComponent: + return openfe_gromacs.ProteinComponent.from_pdb_file(eg5_protein_pdb) @pytest.fixture() @@ -183,7 +193,7 @@ def CN_molecule(): """ A basic CH3NH2 molecule for quick testing. """ - with resources.files("openfe.tests.data") as d: + with resources.files("openfe_gromacs.tests.data") as d: fn = str(d / "CN.sdf") supp = Chem.SDMolSupplier(str(fn), removeHs=False) diff --git a/openfe_gromacs/tests/data/MDProtocol_json_results.gz b/openfe_gromacs/tests/data/MDProtocol_json_results.gz index dd9542d..1aaae47 100644 Binary files a/openfe_gromacs/tests/data/MDProtocol_json_results.gz and b/openfe_gromacs/tests/data/MDProtocol_json_results.gz differ diff --git a/openfe_gromacs/tests/data/MDProtocol_json_results_no_EM.gz b/openfe_gromacs/tests/data/MDProtocol_json_results_no_EM.gz index 5d36692..e4f6986 100644 Binary files a/openfe_gromacs/tests/data/MDProtocol_json_results_no_EM.gz and b/openfe_gromacs/tests/data/MDProtocol_json_results_no_EM.gz differ diff --git a/openfe_gromacs/tests/data/__init__.py b/openfe_gromacs/tests/data/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/openfe_gromacs/tests/data/alanine-dipeptide.pdb b/openfe_gromacs/tests/data/alanine-dipeptide.pdb new file mode 100644 index 0000000..8c2501a --- /dev/null +++ b/openfe_gromacs/tests/data/alanine-dipeptide.pdb @@ -0,0 +1,25 @@ +REMARK ACE +ATOM 1 1HH3 ACE 1 2.000 1.000 -0.000 +ATOM 2 CH3 ACE 1 2.000 2.090 0.000 +ATOM 3 2HH3 ACE 1 1.486 2.454 0.890 +ATOM 4 3HH3 ACE 1 1.486 2.454 -0.890 +ATOM 5 C ACE 1 3.427 2.641 -0.000 +ATOM 6 O ACE 1 4.391 1.877 -0.000 +ATOM 7 N ALA 2 3.555 3.970 -0.000 +ATOM 8 H ALA 2 2.733 4.556 -0.000 +ATOM 9 CA ALA 2 4.853 4.614 -0.000 +ATOM 10 HA ALA 2 5.408 4.316 0.890 +ATOM 11 CB ALA 2 5.661 4.221 -1.232 +ATOM 12 1HB ALA 2 5.123 4.521 -2.131 +ATOM 13 2HB ALA 2 6.630 4.719 -1.206 +ATOM 14 3HB ALA 2 5.809 3.141 -1.241 +ATOM 15 C ALA 2 4.713 6.129 0.000 +ATOM 16 O ALA 2 3.601 6.653 0.000 +ATOM 17 N NME 3 5.846 6.835 0.000 +ATOM 18 H NME 3 6.737 6.359 -0.000 +ATOM 19 CH3 NME 3 5.846 8.284 0.000 +ATOM 20 1HH3 NME 3 4.819 8.648 0.000 +ATOM 21 2HH3 NME 3 6.360 8.648 0.890 +ATOM 22 3HH3 NME 3 6.360 8.648 -0.890 +TER +END diff --git a/openfe_gromacs/tests/protocols/test_create_systems.py b/openfe_gromacs/tests/protocols/test_create_systems.py new file mode 100644 index 0000000..7170bcd --- /dev/null +++ b/openfe_gromacs/tests/protocols/test_create_systems.py @@ -0,0 +1,97 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe-gromacs +import gufe +import numpy as np +from openmm import NonbondedForce +from openmm.app import GromacsGroFile, GromacsTopFile + +from openfe_gromacs.protocols.gromacs_md.md_methods import GromacsMDProtocol +from openfe_gromacs.protocols.gromacs_utils import create_systems + + +def test_interchange_gromacs(alanine_dipeptide_component, tmpdir): + solvent = gufe.SolventComponent() + smc_components = {} + prot_settings = GromacsMDProtocol.default_settings() + + omm_system, omm_topology, omm_positions = create_systems.create_openmm_system( + solvent, + alanine_dipeptide_component, + smc_components, + prot_settings.partial_charge_settings, + prot_settings.forcefield_settings, + prot_settings.integrator_settings, + prot_settings.thermo_settings, + prot_settings.solvation_settings, + prot_settings.output_settings_em, + tmpdir, + ) + omm_atom_names = [atom.name for atom in omm_topology.atoms()] + interchange = create_systems.create_interchange( + omm_system, omm_topology, omm_positions, smc_components + ) + interchange_atom_names = [atom.name for atom in interchange.topology.atoms] + + interchange.to_gro(f"{tmpdir}/test.gro") + gro_atom_names = GromacsGroFile(f"{tmpdir}/test.gro").atomNames + + # Testing that the number of atoms is the same + assert len(omm_atom_names) == len(interchange_atom_names) == len(gro_atom_names) + # Testing that atom names are the same + assert omm_atom_names == interchange_atom_names == gro_atom_names + # Testing that residue names are the same + interchange_res_names = [ + atom.metadata["residue_name"] for atom in interchange.topology.atoms + ] + gro_res_names = GromacsGroFile(f"{tmpdir}/test.gro").residueNames + combined_res_names = interchange_res_names + gro_res_names + assert len(set(combined_res_names)) == len(set(gro_res_names)) + # Testing that residue numbers are the same + interchange_res_numbers = [ + atom.metadata["residue_number"] for atom in interchange.topology.atoms + ] + with open(f"{tmpdir}/test.gro") as f: + gromacs_res_numbers = [int(line[:5].strip()) for line in f.readlines()[2:-1]] + assert interchange_res_numbers == gromacs_res_numbers + + # check a few atom names to ensure these are not empty sets + for atom_name in ("HA", "CH3", "CA", "CB"): + assert atom_name in interchange_atom_names + + +def test_user_charges(CN_molecule, tmpdir): + solvent = gufe.SolventComponent() + off_cn = CN_molecule.to_openff() + off_cn.assign_partial_charges(partial_charge_method="am1-mulliken") + off_charges = off_cn.partial_charges + prot_settings = GromacsMDProtocol.default_settings() + smc_components = {CN_molecule: off_cn} + omm_system, omm_topology, omm_positions = create_systems.create_openmm_system( + solvent, + None, + smc_components, + prot_settings.partial_charge_settings, + prot_settings.forcefield_settings, + prot_settings.integrator_settings, + prot_settings.thermo_settings, + prot_settings.solvation_settings, + prot_settings.output_settings_em, + tmpdir, + ) + + interchange = create_systems.create_interchange( + omm_system, omm_topology, omm_positions, smc_components + ) + # Save to Gromacs .top file + interchange.to_top(f"{tmpdir}/test.top") + # Load Gromacs .top file back in as OpenMM system + gromacs_system = GromacsTopFile(f"{tmpdir}/test.top").createSystem() + # Get the partial charges of the ligand atoms + nonbonded = [ + f for f in gromacs_system.getForces() if isinstance(f, NonbondedForce) + ][0] + gro_charges = [] + for i in range(len(off_charges)): + charge, sigma, epsilon = nonbonded.getParticleParameters(i) + gro_charges.append(charge._value) + np.testing.assert_almost_equal(off_charges.m, gro_charges, decimal=5) diff --git a/openfe_gromacs/tests/protocols/test_gromacs_md.py b/openfe_gromacs/tests/protocols/test_gromacs_md.py index 3cd7b07..537ce6e 100644 --- a/openfe_gromacs/tests/protocols/test_gromacs_md.py +++ b/openfe_gromacs/tests/protocols/test_gromacs_md.py @@ -244,7 +244,7 @@ def test_get_mdp_filenames(self, protocolresult): def test_get_filenames_em(self, protocolresult): dict_file_path = protocolresult.get_filenames_em() assert isinstance(dict_file_path, dict) - assert len(dict_file_path) == 7 + assert len(dict_file_path) == 8 for name, file_path in dict_file_path.items(): assert isinstance(file_path, list) assert len(file_path) == 1 @@ -265,7 +265,7 @@ def test_get_xtc_em_filename(self, protocolresult): def test_get_filenames_nvt(self, protocolresult): dict_file_path = protocolresult.get_filenames_nvt() assert isinstance(dict_file_path, dict) - assert len(dict_file_path) == 7 + assert len(dict_file_path) == 8 for name, file_path in dict_file_path.items(): assert isinstance(file_path, list) assert len(file_path) == 1 @@ -286,7 +286,7 @@ def test_get_xtc_nvt_filenames(self, protocolresult): def test_get_filenames_npt(self, protocolresult): dict_file_path = protocolresult.get_filenames_npt() assert isinstance(dict_file_path, dict) - assert len(dict_file_path) == 7 + assert len(dict_file_path) == 8 for name, file_path in dict_file_path.items(): assert isinstance(file_path, list) assert len(file_path) == 1 @@ -324,7 +324,7 @@ def test_reload_protocol_result(self, md_json_no_em): def test_get_filenames_em(self, protocolresult): dict_file_path = protocolresult.get_filenames_em() assert isinstance(dict_file_path, dict) - assert len(dict_file_path) == 7 + assert len(dict_file_path) == 8 for name, file_path in dict_file_path.items(): assert isinstance(file_path, type(None)) diff --git a/openfe_gromacs/tests/protocols/test_md_tokenization.py b/openfe_gromacs/tests/protocols/test_md_tokenization.py index 1187895..4969494 100644 --- a/openfe_gromacs/tests/protocols/test_md_tokenization.py +++ b/openfe_gromacs/tests/protocols/test_md_tokenization.py @@ -3,7 +3,6 @@ import json import gufe -import openfe import pytest from gufe.tests.test_tokenization import GufeTokenizableTestsMixin @@ -19,7 +18,7 @@ def protocol(): def protocol_units(protocol, benzene_system): pus = protocol.create( stateA=benzene_system, - stateB=openfe.ChemicalSystem({"solvent": openfe.SolventComponent()}), + stateB=benzene_system, mapping=None, ) return list(pus.protocol_units) @@ -41,7 +40,7 @@ def protocol_result(md_json): class TestGromacsMDProtocol(GufeTokenizableTestsMixin): cls = gromacs_md.GromacsMDProtocol - key = "GromacsMDProtocol-ec824b46b95c4ee8393d59287cf2feab" + key = "GromacsMDProtocol-829ed8e2ddfff11f9e86eec044222317" repr = f"<{key}>" @pytest.fixture() @@ -62,6 +61,7 @@ def test_key_stable(self): pytest.skip() +@pytest.mark.skip class TestGromacsMDProtocolResult(GufeTokenizableTestsMixin): cls = gromacs_md.GromacsMDProtocolResult key = "GromacsMDProtocolResult-ba53ff35cb02f1d9869f22955141e250"