Utility function changes for protein mutations

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,

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,

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)

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,

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
     """

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

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': [],

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,