diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py index 88caf4242..472793956 100644 --- a/openfe/protocols/openmm_afe/base.py +++ b/openfe/protocols/openmm_afe/base.py @@ -458,8 +458,8 @@ def _get_modeller( # get OpenMM modeller + dictionary of resids for each component system_modeller, comp_resids = system_creation.get_omm_modeller( - protein_comp=protein_component, - solvent_comp=solvent_component, + protein_comps=protein_component, + solvent_comps=solvent_component, small_mols=smc_components, omm_forcefield=system_generator.forcefield, solvent_settings=solvation_settings, diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py index 24c3f1fb0..24cb35423 100644 --- a/openfe/protocols/openmm_md/plain_md_methods.py +++ b/openfe/protocols/openmm_md/plain_md_methods.py @@ -603,8 +603,8 @@ def run(self, *, dry=False, verbose=True, # 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, + protein_comps=protein_comp, + solvent_comps=solvent_comp, small_mols=smc_components, omm_forcefield=system_generator.forcefield, solvent_settings=solvation_settings, diff --git a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py index 41e83238c..cdf207705 100644 --- a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py +++ b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py @@ -390,10 +390,6 @@ def _get_indices(topology, resids): residue_name : str Name of the residue to get the indices for. """ - # TODO: remove, this shouldn't be necessary anymore - if len(resids) > 1: - raise ValueError("multiple residues were found") - # create list of openmm residues top_res = [r for r in topology.residues() if r.index in resids] @@ -502,6 +498,7 @@ def pick_H(i, j, x, y) -> int: to_del.append(pick_H(i, j, x, y)) + # to_del = set(to_del) # Making a set to avoid repeated indices (TODO: TEST!!) for idx in to_del: del no_const_old_to_new_atom_map[idx] @@ -677,7 +674,7 @@ def set_and_check_new_positions(mapping, old_topology, new_topology, state. tolerance : float Warning threshold for deviations along any dimension (x,y,z) in mapped - atoms between the "old" and "new" positions. Default 1.0. + atoms between the "old" and "new" positions, in Angstroms. Default 1.0. """ # Get the positions in Angstrom as raw numpy arrays old_pos_array = old_positions.value_in_unit(omm_unit.angstrom) diff --git a/openfe/protocols/openmm_rfe/equil_rfe_methods.py b/openfe/protocols/openmm_rfe/equil_rfe_methods.py index d2f95f2c3..1b2e7de7c 100644 --- a/openfe/protocols/openmm_rfe/equil_rfe_methods.py +++ b/openfe/protocols/openmm_rfe/equil_rfe_methods.py @@ -154,14 +154,21 @@ def _get_alchemical_charge_difference( "correction has been requested. Unfortunately " "only absolute differences of 1 are supported.") raise ValueError(errmsg) - - ion = {-1: solvent_component.positive_ion, - 1: solvent_component.negative_ion}[difference] - wmsg = (f"A charge difference of {difference} is observed " - "between the end states. This will be addressed by " - f"transforming a water into a {ion} ion") - logger.warning(wmsg) - warnings.warn(wmsg) + if solvent_component: + ion = {-1: solvent_component.positive_ion, + 1: solvent_component.negative_ion}[difference] + wmsg = (f"A charge difference of {difference} is observed " + "between the end states. This will be addressed by " + f"transforming a water into a {ion} ion") + logger.warning(wmsg) + warnings.warn(wmsg) + else: + errmsg = (f"A charge difference of {difference} is observed " + "between the end states. This would be addressed by " + f"transforming a water into an ion. However, " + "No solvent component was specified when trying to" + "apply the charge correction.") + raise ValueError(errmsg) else: wmsg = (f"A charge difference of {difference} is observed " "between the end states. No charge correction has " @@ -796,8 +803,8 @@ def run(self, *, dry=False, verbose=True, # c. get OpenMM Modeller + a dictionary of resids for each component stateA_modeller, comp_resids = system_creation.get_omm_modeller( - protein_comp=protein_comp, - solvent_comp=solvent_comp, + protein_comps=protein_comp, + solvent_comps=solvent_comp, small_mols=dict(chain(off_small_mols['stateA'], off_small_mols['both'])), omm_forcefield=system_generator.forcefield, diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 8e2e5c28b..5f51d78b3 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -342,17 +342,17 @@ class Config: """Size of the simulation timestep. Default 4.0 * unit.femtosecond.""" langevin_collision_rate: FloatQuantity['1/picosecond'] = 1.0 / unit.picosecond """Collision frequency. Default 1.0 / unit.pisecond.""" - reassign_velocities = False + reassign_velocities: bool = False """ If ``True``, velocities are reassigned from the Maxwell-Boltzmann distribution at the beginning of each MC move. Default ``False``. """ - n_restart_attempts = 20 + n_restart_attempts: int = 20 """ Number of attempts to restart from Context if there are NaNs in the energies after integration. Default 20. """ - constraint_tolerance = 1e-06 + constraint_tolerance: float = 1e-06 """Tolerance for the constraint solver. Default 1e-6.""" barostat_frequency: FloatQuantity['timestep'] = 25.0 * unit.timestep # todo: IntQuantity """ diff --git a/openfe/protocols/openmm_utils/system_creation.py b/openfe/protocols/openmm_utils/system_creation.py index 15635619b..080aab745 100644 --- a/openfe/protocols/openmm_utils/system_creation.py +++ b/openfe/protocols/openmm_utils/system_creation.py @@ -8,10 +8,9 @@ import numpy.typing as npt from openmm import app, MonteCarloBarostat from openmm import unit as omm_unit -from openff.toolkit import Molecule as OFFMol from openff.units.openmm import to_openmm, ensure_quantity from openmmforcefields.generators import SystemGenerator -from typing import Optional +from typing import Optional, Iterable from pathlib import Path from gufe.settings import OpenMMSystemGeneratorFFSettings, ThermoSettings from gufe import ( @@ -134,9 +133,9 @@ def get_system_generator( def get_omm_modeller( - protein_comp: Optional[ProteinComponent], - solvent_comp: Optional[SolventComponent], - small_mols: dict[SmallMoleculeComponent, OFFMol], + protein_comps: Optional[Iterable[ProteinComponent] | ProteinComponent], + solvent_comps: Optional[Iterable[SolventComponent] | SolventComponent], + small_mols: Optional[Iterable[SmallMoleculeComponent] | SmallMoleculeComponent], omm_forcefield : app.ForceField, solvent_settings : OpenMMSolvationSettings ) -> ModellerReturn: @@ -146,11 +145,11 @@ def get_omm_modeller( Parameters ---------- - protein_comp : Optional[ProteinComponent] + protein_comps : Optional[Iterable[ProteinComponent] or ProteinComponent] Protein Component, if it exists. - solvent_comp : Optional[ProteinCompoinent] + solvent_comps : Optional[Iterable[SolventComponent] or SolventComponent] Solvent Component, if it exists. - small_mols : dict + small_mols : Optional[Iterable[SmallMoleculeComponent] or SmallMoleculeComponent] Small molecules to add. omm_forcefield : app.ForceField ForceField object for system. @@ -188,28 +187,49 @@ def _add_small_mol(comp, # Create empty modeller system_modeller = app.Modeller(app.Topology(), []) - # If there's a protein in the system, we add it first to the Modeller - if protein_comp is not None: - system_modeller.add(protein_comp.to_openmm_topology(), - protein_comp.to_openmm_positions()) - # add missing virtual particles (from crystal waters) - system_modeller.addExtraParticles(omm_forcefield) - component_resids[protein_comp] = np.array( - [r.index for r in system_modeller.topology.residues()] - ) - # if we solvate temporarily rename water molecules to 'WAT' - # see openmm issue #4103 - if solvent_comp is not None: - for r in system_modeller.topology.residues(): - if r.name == 'HOH': - r.name = 'WAT' + # We first add all the protein components, if any + if protein_comps: + try: + protein_comps = iter(protein_comps) + except TypeError: + protein_comps = {protein_comps} # make it a set/iterable with the comp + for protein_comp in protein_comps: + system_modeller.add(protein_comp.to_openmm_topology(), + protein_comp.to_openmm_positions()) + # add missing virtual particles (from crystal waters) + system_modeller.addExtraParticles(omm_forcefield) + component_resids[protein_comp] = np.array( + [r.index for r in system_modeller.topology.residues()] + ) + # if we solvate temporarily rename water molecules to 'WAT' + # see openmm issue #4103 + if solvent_comps is not None: + for r in system_modeller.topology.residues(): + if r.name == 'HOH': + r.name = 'WAT' # Now loop through small mols - for comp, mol in small_mols.items(): - _add_small_mol(comp, mol, system_modeller, component_resids) + if small_mols: + try: + small_mols = iter(small_mols) + except TypeError: + small_mols = {small_mols} # make it a set/iterable with the comp + for small_mol_comp in small_mols: + _add_small_mol(small_mol_comp, small_mol_comp.to_openff(), system_modeller, + component_resids) # Add solvent if neeeded - if solvent_comp is not None: + if solvent_comps: + # Making it a list to make our life easier -- TODO: Maybe there's a better type for this + try: + solvent_comps = list(set(solvent_comps)) # if given iterable + except TypeError: + solvent_comps = [solvent_comps] # if not iterable, given single obj + # TODO: Support multiple solvent components? Is there a use case for it? + # Error out when we iter(have more than one solvent component in the states/systems + if len(solvent_comps) > 1: + raise ValueError("More than one solvent component found in systems. Only one supported.") + solvent_comp = solvent_comps[0] # Get the first (and only?) solvent component # Do unit conversions if necessary solvent_padding = None box_size = None diff --git a/openfe/protocols/openmm_utils/system_validation.py b/openfe/protocols/openmm_utils/system_validation.py index f1710d93e..9eefc9d92 100644 --- a/openfe/protocols/openmm_utils/system_validation.py +++ b/openfe/protocols/openmm_utils/system_validation.py @@ -5,7 +5,6 @@ Protocols. """ from typing import Optional, Tuple -from openff.toolkit import Molecule as OFFMol from gufe import ( Component, ChemicalSystem, SolventComponent, ProteinComponent, SmallMoleculeComponent @@ -37,6 +36,7 @@ def get_alchemical_components( ValueError If there are any duplicate components in states A or B. """ + # TODO: We probably want to rewrite this using sets. As in _check_states_compatibility in feflow.protocols.ProteinMutationProtocol matched_components: dict[Component, Component] = {} alchemical_components: dict[str, list[Component]] = { 'stateA': [], 'stateB': [], diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py index 513962355..a9b7ef353 100644 --- a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py +++ b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py @@ -1717,8 +1717,8 @@ def benzene_solvent_openmm_system(benzene_modifications): ) modeller, _ = system_creation.get_omm_modeller( - protein_comp=None, - solvent_comp=openfe.SolventComponent(), + protein_comps=None, + solvent_comps=openfe.SolventComponent(), small_mols={smc: offmol}, omm_forcefield=system_generator.forcefield, solvent_settings=settings.solvation_settings, @@ -1758,8 +1758,8 @@ def benzene_tip4p_solvent_openmm_system(benzene_modifications): ) modeller, _ = system_creation.get_omm_modeller( - protein_comp=None, - solvent_comp=openfe.SolventComponent(), + protein_comps=None, + solvent_comps=openfe.SolventComponent(), small_mols={smc: offmol}, omm_forcefield=system_generator.forcefield, solvent_settings=settings.solvation_settings,