[ WIP ] Support for charge changes in the HybTop protocol for explicitly solvated systems

src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py

@@ -10,6 +10,7 @@
 import itertools
 import logging
 import warnings
+from collections import Counter
 from copy import deepcopy
 from typing import Optional, Union
 
@@ -21,77 +22,179 @@
 from openmm import NonbondedForce, System, app
 from openmm import unit as omm_unit
 
-from openfe import SolventComponent
-
 logger = logging.getLogger(__name__)
 
 
-def _get_ion_and_water_parameters(
+def _get_ion_parameters_from_forcefield(
+    forcefield: app.ForceField,
+    ion_resname: str,
+    ion_element: app.Element,
+) -> tuple:
+    """
+    Get NonbondedForce parameters for a bare monovalent ion by creating a
+    minimal single-ion system via the ForceField. Used as a fallback when
+    no ion of the required type is present in the topology.
+
+    Parameters
+    ----------
+    forcefield : app.ForceField
+        The ForceField to use to parameterize the ion.
+    ion_resname : str
+        The residue name of the ion to parameterize (e.g. 'NA', 'CL').
+    ion_element : app.Element
+        The element of the ion.
+
+    Returns
+    -------
+    ion_charge : openmm.unit.Quantity
+        The partial charge of the ion.
+    ion_sigma : openmm.unit.Quantity
+        The NonbondedForce sigma parameter of the ion.
+    ion_epsilon : openmm.unit.Quantity
+        The NonbondedForce epsilon parameter of the ion.
+    """
+    dummy_top = app.Topology()
+    res = dummy_top.addResidue(ion_resname, dummy_top.addChain())
+    dummy_top.addAtom(ion_resname, ion_element, res)
+    dummy_system = forcefield.createSystem(dummy_top)
+
+    nbf = [i for i in dummy_system.getForces() if isinstance(i, NonbondedForce)][0]
+
+    return nbf.getParticleParameters(0)
+
+def _get_water_parameters(
     topology: app.Topology,
     system: System,
-    ion_resname: str,
     water_resname: str = 'HOH',
-):
+) -> tuple:
     """
-    Get ion, and water (oxygen and hydrogen) atoms parameters.
+    Get water oxygen and hydrogen partial charges from the system.
 
     Parameters
     ----------
     topology : app.Topology
-      The topology to search for the ion and water
+      The topology to search for water atoms.
     system : app.System
       The system associated with the input topology object.
-    ion_resname : str
-      The residue name of the ion to get parameters for
     water_resname : str
-      The residue name of the water to get parameters for. Default 'HOH'.
+      The residue name of the water. Default 'HOH'.
 
     Returns
     -------
-    ion_charge : float
-      The partial charge of the ion atom
-    ion_sigma : float
-      The NonbondedForce sigma parameter of the ion atom
-    ion_epsilon : float
-      The NonbondedForce epsilon parameter of the ion atom
-    o_charge : float
+    o_charge : openmm.unit.Quantity
       The partial charge of the water oxygen.
-    h_charge : float
+    h_charge : openmm.unit.Quantity
       The partial charge of the water hydrogen.
 
     Raises
     ------
     ValueError
-      If there are no ``ion_resname`` or ``water_resname`` named residues in
-      the input ``topology``.
-
-    Attribution
-    -----------
-    Based on `perses.utils.charge_changing.get_ion_and_water_parameters`.
+      If no water residue named ``water_resname`` with O and H atoms
+      is found in the topology.
     """
-    def _find_atom(topology, resname, elementname):
-        for atom in topology.atoms():
-            if atom.residue.name == resname:
-                if (elementname is None or atom.element.symbol == elementname):
-                    return atom.index
-        errmsg = ("Error encountered when attempting to explicitly handle "
-                  "charge changes using an alchemical water. No residue "
-                  f"named: {resname} found, with element {elementname}")
-        raise ValueError(errmsg)
+    nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
+
+    o_charge = None
+    h_charge = None
 
-    ion_index = _find_atom(topology, ion_resname, None)
-    oxygen_index = _find_atom(topology, water_resname, 'O')
-    hydrogen_index = _find_atom(topology, water_resname, 'H')
+    for residue in topology.residues():
+        if residue.name != water_resname:
+            continue
+        for atom in residue.atoms():
+            if atom.element is None:
+                continue
+            charge, _, _ = nbf.getParticleParameters(atom.index)
+            if atom.element.symbol == 'O':
+                o_charge = charge
+            elif atom.element.symbol == 'H':
+                h_charge = charge
+        # Stop as soon as we have both charges
+        if o_charge is not None and h_charge is not None:
+            break
+
+    if o_charge is None or h_charge is None:
+        raise ValueError(
+            f"Could not find water residue '{water_resname}' with O and H "
+            "atoms in the topology. Water parameters are required for "
+            "alchemical charge correction."
+        )
+
+    return o_charge, h_charge
+
+def _get_ion_parameters(
+    topology: app.Topology,
+    system: System,
+    charge_difference: int,
+    forcefield: app.ForceField,
+) -> tuple:
+    """
+    Get NonbondedForce parameters for the most abundant monovalent ion of
+    the appropriate charge sign found in the topology. Falls back to
+    'NA' or 'CL' parameters if no ion is found in the topology.
 
-    nbf = [i for i in system.getForces()
-           if isinstance(i, NonbondedForce)][0]
+    Parameters
+    ----------
+    topology : app.Topology
+      The topology to search for the most abundant ion type.
+    system : app.System
+      The system associated with the input topology object.
+    charge_difference : int
+      The charge difference between state A and state B. Determines
+      whether to look for a positive or negative ion.
+    forcefield : app.ForceField
+      The force field to use for parameterization if no ion is found
+      in the topology.
 
-    ion_charge, ion_sigma, ion_epsilon = nbf.getParticleParameters(ion_index)
-    o_charge, _, _ = nbf.getParticleParameters(oxygen_index)
-    h_charge, _, _ = nbf.getParticleParameters(hydrogen_index)
+    Returns
+    -------
+    ion_charge : openmm.unit.Quantity
+    ion_sigma : openmm.unit.Quantity
+    ion_epsilon : openmm.unit.Quantity
+    """
+    nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
+
+    desired_sign = np.sign(charge_difference)
+    ion_counts: Counter = Counter()
+    ion_atom_indices: dict[str, int] = {}
+
+    for residue in topology.residues():
+        atoms = list(residue.atoms())
+        if len(atoms) != 1:
+            continue
+        atom = atoms[0]
+        if atom.element is None:
+            continue
+        charge, _, _ = nbf.getParticleParameters(atom.index)
+        charge_val = charge.value_in_unit(omm_unit.elementary_charge)
+        if np.isclose(abs(charge_val), 1.0, atol=0.01) and np.sign(charge_val) == desired_sign:
+            ion_counts[residue.name] += 1
+            if residue.name not in ion_atom_indices:
+                ion_atom_indices[residue.name] = atom.index
+
+    if ion_counts:
+        # Use the most abundant matching ion found in the topology
+        best_resname = ion_counts.most_common(1)[0][0]
+        return nbf.getParticleParameters(ion_atom_indices[best_resname])
+
+    # No matching ion in topology: fall back to NA/CL from forcefield
+    if charge_difference > 0:
+        fallback_resname = 'NA'
+        fallback_element = app.Element.getByAtomicNumber(11)  # Na
+    else:
+        fallback_resname = 'CL'
+        fallback_element = app.Element.getByAtomicNumber(17)  # Cl
 
-    return ion_charge, ion_sigma, ion_epsilon, o_charge, h_charge
+    wmsg = (
+        f"No monovalent ion with the appropriate charge sign found in "
+        f"the topology. Defaulting to '{fallback_resname}' and obtaining "
+        f"parameters from the forcefield directly."
+    )
+    warnings.warn(wmsg)
+    logger.warning(wmsg)
 
+    return _get_ion_parameters_from_forcefield(
+        forcefield, fallback_resname, fallback_element,
+    )
 
 def _fix_alchemical_water_atom_mapping(
     system_mapping: dict[str, Union[dict[int, int], list[int]]],
@@ -123,10 +226,13 @@ def _fix_alchemical_water_atom_mapping(
 
 
 def handle_alchemical_waters(
-    water_resids: list[int], topology: app.Topology,
-    system: System, system_mapping: dict,
+    water_resids: list[int],
+    topology: app.Topology,
+    system: System,
+    system_mapping: dict,
     charge_difference: int,
-    solvent_component: SolventComponent,
+    forcefield: app.ForceField,
+    water_resname: str,
 ):
     """
     Add alchemical waters from a pre-defined list.
@@ -143,14 +249,10 @@ def handle_alchemical_waters(
       A dictionary of system mappings between the stateA and stateB systems
     charge_difference : int
       The charge difference between state A and state B.
-    positive_ion_resname : str
-      The name of a positive ion to replace the water with if the absolute
-      charge difference is positive.
-    negative_ion_resname : str
-      The name of a negative ion to replace the water with if the absolute
-      charge difference is negative.
+    forcefield : app.ForceField
+      The forcefield to use for ion parameterization.
     water_resname : str
-      The residue name of the water to get parameters for. Default 'HOH'.
+      The residue name of the water to get parameters for.
 
     Raises
     ------
@@ -171,18 +273,13 @@ def handle_alchemical_waters(
                   f"difference: {abs(charge_difference)}")
         raise ValueError(errmsg)
 
-    if charge_difference > 0:
-        ion_resname = solvent_component.positive_ion.strip('-+').upper()
-    elif charge_difference < 0:
-        ion_resname = solvent_component.negative_ion.strip('-+').upper()
-    # if there's no charge difference then just skip altogether
-    else:
+    if charge_difference == 0:
         return None
 
-    ion_charge, ion_sigma, ion_epsilon, o_charge, h_charge = _get_ion_and_water_parameters(
-        topology, system, ion_resname,
-        'HOH',  # Modeller always adds HOH waters
+    ion_charge, ion_sigma, ion_epsilon = _get_ion_parameters(
+        topology, system, charge_difference, forcefield,
     )
+    o_charge, h_charge = _get_water_parameters(topology, system, water_resname)
 
     # get the nonbonded forces
     nbfrcs = [i for i in system.getForces()
@@ -433,7 +530,6 @@ def _remove_constraints(old_to_new_atom_map, old_system, old_topology,
     * Very slow, needs refactoring
     * Can we drop having topologies as inputs here?
     """
-    from collections import Counter
 
     no_const_old_to_new_atom_map = deepcopy(old_to_new_atom_map)
 

src/openfe/protocols/openmm_rfe/hybridtop_protocols.py

@@ -28,6 +28,7 @@
     ProteinComponent,
     ProteinMembraneComponent,
     SmallMoleculeComponent,
+    SolvatedPDBComponent,
     SolventComponent,
     settings,
 )
@@ -39,6 +40,7 @@
     settings_validation,
     system_validation,
 )
+from . import _rfe_utils
 from .equil_rfe_settings import (
     AlchemicalSettings,
     IntegratorSettings,
@@ -372,7 +374,7 @@ def _validate_charge_difference(
         mapping: LigandAtomMapping,
         nonbonded_method: str,
         explicit_charge_correction: bool,
-        solvent_component: SolventComponent | None,
+        solvent_component: SolventComponent | SolvatedPDBComponent | None,
     ):
         """
         Validates the net charge difference between the two states.
@@ -435,12 +437,10 @@ def _validate_charge_difference(
             )
             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"
+            "transforming a water into an ion of opposite charge."
         )
         logger.info(wmsg)
 
@@ -560,6 +560,8 @@ def _validate(
         # Note: validation depends on the mapping & solvent component checks
         if stateA.contains(SolventComponent):
             solv_comp = stateA.get_components_of_type(SolventComponent)[0]
+        elif stateA.contains(SolvatedPDBComponent):
+            solv_comp = stateA.get_components_of_type(SolvatedPDBComponent)[0]
         else:
             solv_comp = None
 

src/openfe/protocols/openmm_rfe/hybridtop_units.py

@@ -390,7 +390,7 @@ def _handle_net_charge(
         charge_difference: int,
         system_mappings: dict[str, dict[int, int]],
         distance_cutoff: Quantity,
-        solvent_component: SolventComponent | None,
+        forcefield: openmm.app.ForceField,
     ) -> None:
         """
         Handle system net charge by adding an alchemical water.
@@ -404,7 +404,7 @@ def _handle_net_charge(
         charge_difference : int
         system_mappings : dict[str, dict[int, int]]
         distance_cutoff : Quantity
-        solvent_component : SolventComponent | None
+        forcefield: openmm.app.ForceField
         """
         # Base case, return if no net charge
         if charge_difference == 0:
@@ -425,7 +425,8 @@ def _handle_net_charge(
             system=stateB_system,
             system_mapping=system_mappings,
             charge_difference=charge_difference,
-            solvent_component=solvent_component,
+            forcefield=forcefield,
+            water_resname="HOH",  # Need to check if this is also true for systems not prepped with addSolvent
         )
 
     def _get_omm_objects(
@@ -545,6 +546,7 @@ def _filter_small_mols(smols, state):
         # Net charge: add alchemical water if needed
         # Must be done here as we in-place modify the particles of state B.
         if settings["alchemical_settings"].explicit_charge_correction:
+            forcefield = states_inputs["A"]["generator"].forcefield
             self._handle_net_charge(
                 stateA_topology=stateA_topology,
                 stateA_positions=stateA_positions,
@@ -553,7 +555,7 @@ def _filter_small_mols(smols, state):
                 charge_difference=mapping.get_alchemical_charge_difference(),
                 system_mappings=system_mappings,
                 distance_cutoff=settings["alchemical_settings"].explicit_charge_correction_cutoff,
-                solvent_component=solvent_component,
+                forcefield=forcefield,
             )
 
         # Finally get the state B positions