From 403d38e02f0976f001a5a59d8e9148f5684a3682 Mon Sep 17 00:00:00 2001
From: hannahbaumann <hbaumann@uci.edu>
Date: Mon, 18 May 2026 16:50:08 +0200
Subject: [PATCH 1/9] Initial test at what we may need to change to support
 charge changes for explicitly solvated systems

---
 .../openmm_rfe/_rfe_utils/topologyhelpers.py  | 77 +++++++++++++++++--
 .../openmm_rfe/hybridtop_protocols.py         | 17 +++-
 .../protocols/openmm_rfe/hybridtop_units.py   | 15 +++-
 .../openmm_rfe/test_hybrid_top_protocol.py    | 50 ++++++++++--
 4 files changed, 144 insertions(+), 15 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index 7181f4ce5..2a4e6069b 100644
--- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/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
 
@@ -123,10 +124,14 @@ 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,
+    positive_ion_resname: str,
+    negative_ion_resname: str,
+    water_resname: str,
 ):
     """
     Add alchemical waters from a pre-defined list.
@@ -150,7 +155,7 @@ def handle_alchemical_waters(
       The name of a negative ion to replace the water with if the absolute
       charge difference is negative.
     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
     ------
@@ -172,16 +177,16 @@ def handle_alchemical_waters(
         raise ValueError(errmsg)
 
     if charge_difference > 0:
-        ion_resname = solvent_component.positive_ion.strip('-+').upper()
+        ion_resname = positive_ion_resname
     elif charge_difference < 0:
-        ion_resname = solvent_component.negative_ion.strip('-+').upper()
+        ion_resname = negative_ion_resname
     # if there's no charge difference then just skip altogether
     else:
         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
+        water_resname,
     )
 
     # get the nonbonded forces
@@ -433,7 +438,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)
 
@@ -720,3 +724,60 @@ def set_and_check_new_positions(mapping, old_topology, new_topology,
             logging.warning(wmsg)
 
     return new_pos_array * omm_unit.angstrom
+
+
+def _get_ion_resnames_from_topology(topology: app.Topology) -> tuple[str, str]:
+    """
+    Infer positive and negative ion residue names from a topology by
+    finding the most common monovalent ion of each charge type.
+    Falls back to NA/CL if none are found.
+
+    Parameters
+    ----------
+    topology : app.Topology
+        The topology to search for ions.
+
+    Returns
+    -------
+    pos_ion : str
+        The residue name of the most abundant positive monovalent ion (Na, K).
+    neg_ion : str
+        The residue name of the most abundant negative monovalent ion (Cl).
+    """
+    known_positive = {'NA', 'K'}
+    # This doesn't make much sense yet to check it, since it's only Cl, but
+    # leaving it here for now so we can add other neg ions.
+    known_negative = {'CL'}
+
+    pos_counts = Counter(
+        r.name for r in topology.residues() if r.name in known_positive
+    )
+    neg_counts = Counter(
+        r.name for r in topology.residues() if r.name in known_negative
+    )
+
+    if not pos_counts:
+        wmsg = (
+            "Could not find any known positive monovalent ions "
+            f"(searched for {known_positive}) in the topology. "
+            "Defaulting to NA for explicit charge correction."
+        )
+        warnings.warn(wmsg)
+        logger.warning(wmsg)
+        pos_ion = 'NA'
+    else:
+        pos_ion = max(pos_counts, key=pos_counts.get)
+
+    if not neg_counts:
+        wmsg = (
+            "Could not find any known negative monovalent ions "
+            f"(searched for {known_negative}) in the topology. "
+            "Defaulting to CL for explicit charge correction."
+        )
+        warnings.warn(wmsg)
+        logger.warning(wmsg)
+        neg_ion = 'CL'
+    else:
+        neg_ion = max(neg_counts, key=neg_counts.get)
+
+    return pos_ion, neg_ion
diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
index ed83e0459..c25b333b9 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
+++ b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
@@ -28,6 +28,7 @@
     ProteinComponent,
     ProteinMembraneComponent,
     SmallMoleculeComponent,
+    SolvatedPDBComponent,
     SolventComponent,
     settings,
 )
@@ -50,6 +51,7 @@
     OpenMMSolvationSettings,
     RelativeHybridTopologyProtocolSettings,
 )
+from . import _rfe_utils
 from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult
 from .hybridtop_units import (
     HybridTopologyMultiStateAnalysisUnit,
@@ -435,7 +437,20 @@ def _validate_charge_difference(
             )
             raise ValueError(errmsg)
 
-        ion = {-1: solvent_component.positive_ion, 1: solvent_component.negative_ion}[difference]
+        # resolve ion names from SolventComponent or topology
+        if isinstance(solvent_component, SolventComponent):
+            positive_ion = solvent_component.positive_ion.strip(
+                '-+').upper()
+            negative_ion = solvent_component.negative_ion.strip(
+                '-+').upper()
+        elif isinstance(solvent_component, SolvatedPDBComponent):
+            positive_ion, negative_ion = (
+                _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
+                    solvent_component.to_openmm_topology()
+                )
+            )
+
+        ion = {-1: positive_ion, 1: negative_ion}[difference]
 
         wmsg = (
             f"A charge difference of {difference} is observed "
diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_units.py b/src/openfe/protocols/openmm_rfe/hybridtop_units.py
index cd634f391..e6d6d0626 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_units.py
+++ b/src/openfe/protocols/openmm_rfe/hybridtop_units.py
@@ -410,6 +410,17 @@ def _handle_net_charge(
         if charge_difference == 0:
             return
 
+        # Resolve ion names from the solvent component if available,
+        # otherwise infer from the topology (for explicitly solvated systems).
+        if isinstance(solvent_component, SolventComponent):
+            positive_ion = solvent_component.positive_ion.strip('-+').upper()
+            negative_ion = solvent_component.negative_ion.strip('-+').upper()
+        else:
+            positive_ion, negative_ion = (
+                _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
+                    stateA_topology)
+            )
+
         # Get the residue ids for waters to turn alchemical
         alchem_water_resids = _rfe_utils.topologyhelpers.get_alchemical_waters(
             topology=stateA_topology,
@@ -425,7 +436,9 @@ def _handle_net_charge(
             system=stateB_system,
             system_mapping=system_mappings,
             charge_difference=charge_difference,
-            solvent_component=solvent_component,
+            positive_ion_resname=positive_ion,
+            negative_ion_resname=negative_ion,
+            water_resname='HOH'   # Need to check if this is also true for systems not prepped with addSolvent
         )
 
     def _get_omm_objects(
diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
index fb297e248..0a15971ea 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
@@ -1215,6 +1215,36 @@ def test_dry_run_membrane_complex(
         )
 
 
+def test_validate_charge_difference_membrane_system(
+    a2a_protein_membrane_component,
+    a2a_ligands,
+):
+    ligA = next(c for c in a2a_ligands if c.name == "4g")
+    ligB = next(c for c in a2a_ligands if c.name == "4h")
+
+    mapping = openfe.LigandAtomMapping(
+        componentA=ligA,
+        componentB=ligB,
+        componentA_to_componentB={i: i for i in range(36)},
+    )
+
+    settings = openmm_rfe.RelativeHybridTopologyProtocol.default_settings()
+
+    # Mock get_alchemical_charge_difference to return non-zero
+    with mock.patch.object(
+        mapping,
+        "get_alchemical_charge_difference",
+        return_value=1,
+    ):
+        protocol = openmm_rfe.RelativeHybridTopologyProtocol(settings=settings)
+        protocol._validate_charge_difference(
+            mapping=mapping,
+            nonbonded_method=settings.forcefield_settings.nonbonded_method,
+            explicit_charge_correction=True,
+            solvent_component=a2a_protein_membrane_component,
+        )
+
+
 def test_lambda_schedule_default():
     lambdas = openmm_rfe._rfe_utils.lambdaprotocol.LambdaProtocol(functions="default")
     assert len(lambdas.lambda_schedule) == 10
@@ -1980,7 +2010,9 @@ def test_handle_alchemwats_incorrect_count(
             system=system,
             system_mapping={},
             charge_difference=1,
-            solvent_component=openfe.SolventComponent(),
+            positive_ion_resname='NA',
+            negative_ion_resname='CL',
+            water_resname='HOH',
         )
 
 
@@ -2004,7 +2036,9 @@ def test_handle_alchemwats_too_many_nbf(
             system=new_system,
             system_mapping={},
             charge_difference=1,
-            solvent_component=openfe.SolventComponent(),
+            positive_ion_resname='NA',
+            negative_ion_resname='CL',
+            water_resname='HOH',
         )
 
 
@@ -2026,7 +2060,9 @@ def test_handle_alchemwats_vsite_water(
             system=system,
             system_mapping={},
             charge_difference=1,
-            solvent_component=openfe.SolventComponent(),
+            positive_ion_resname='NA',
+            negative_ion_resname='CL',
+            water_resname='HOH',
         )
 
 
@@ -2054,7 +2090,9 @@ def test_handle_alchemwats_incorrect_atom(
             system=new_system,
             system_mapping=benzene_self_system_mapping,
             charge_difference=1,
-            solvent_component=openfe.SolventComponent(),
+            positive_ion_resname='NA',
+            negative_ion_resname='CL',
+            water_resname='HOH',
         )
 
 
@@ -2073,7 +2111,9 @@ def test_handle_alchemical_wats(
         system=system,
         system_mapping=benzene_self_system_mapping,
         charge_difference=1,
-        solvent_component=openfe.SolventComponent(),
+        positive_ion_resname='NA',
+        negative_ion_resname='CL',
+        water_resname='HOH',
     )
 
     # check the mappings

From 1c218a29ee343026cbcec22779ebfe78621df7bc Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 18 May 2026 14:55:42 +0000
Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../openmm_rfe/hybridtop_protocols.py         | 14 ++++-----
 .../protocols/openmm_rfe/hybridtop_units.py   | 11 ++++---
 .../openmm_rfe/test_hybrid_top_protocol.py    | 30 +++++++++----------
 3 files changed, 25 insertions(+), 30 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
index c25b333b9..4be7eff5a 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
+++ b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
@@ -40,6 +40,7 @@
     settings_validation,
     system_validation,
 )
+from . import _rfe_utils
 from .equil_rfe_settings import (
     AlchemicalSettings,
     IntegratorSettings,
@@ -51,7 +52,6 @@
     OpenMMSolvationSettings,
     RelativeHybridTopologyProtocolSettings,
 )
-from . import _rfe_utils
 from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult
 from .hybridtop_units import (
     HybridTopologyMultiStateAnalysisUnit,
@@ -439,15 +439,11 @@ def _validate_charge_difference(
 
         # resolve ion names from SolventComponent or topology
         if isinstance(solvent_component, SolventComponent):
-            positive_ion = solvent_component.positive_ion.strip(
-                '-+').upper()
-            negative_ion = solvent_component.negative_ion.strip(
-                '-+').upper()
+            positive_ion = solvent_component.positive_ion.strip("-+").upper()
+            negative_ion = solvent_component.negative_ion.strip("-+").upper()
         elif isinstance(solvent_component, SolvatedPDBComponent):
-            positive_ion, negative_ion = (
-                _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
-                    solvent_component.to_openmm_topology()
-                )
+            positive_ion, negative_ion = _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
+                solvent_component.to_openmm_topology()
             )
 
         ion = {-1: positive_ion, 1: negative_ion}[difference]
diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_units.py b/src/openfe/protocols/openmm_rfe/hybridtop_units.py
index e6d6d0626..988e0a990 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_units.py
+++ b/src/openfe/protocols/openmm_rfe/hybridtop_units.py
@@ -413,12 +413,11 @@ def _handle_net_charge(
         # Resolve ion names from the solvent component if available,
         # otherwise infer from the topology (for explicitly solvated systems).
         if isinstance(solvent_component, SolventComponent):
-            positive_ion = solvent_component.positive_ion.strip('-+').upper()
-            negative_ion = solvent_component.negative_ion.strip('-+').upper()
+            positive_ion = solvent_component.positive_ion.strip("-+").upper()
+            negative_ion = solvent_component.negative_ion.strip("-+").upper()
         else:
-            positive_ion, negative_ion = (
-                _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
-                    stateA_topology)
+            positive_ion, negative_ion = _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
+                stateA_topology
             )
 
         # Get the residue ids for waters to turn alchemical
@@ -438,7 +437,7 @@ def _handle_net_charge(
             charge_difference=charge_difference,
             positive_ion_resname=positive_ion,
             negative_ion_resname=negative_ion,
-            water_resname='HOH'   # Need to check if this is also true for systems not prepped with addSolvent
+            water_resname="HOH",  # Need to check if this is also true for systems not prepped with addSolvent
         )
 
     def _get_omm_objects(
diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
index 0a15971ea..00c44ac14 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
@@ -2010,9 +2010,9 @@ def test_handle_alchemwats_incorrect_count(
             system=system,
             system_mapping={},
             charge_difference=1,
-            positive_ion_resname='NA',
-            negative_ion_resname='CL',
-            water_resname='HOH',
+            positive_ion_resname="NA",
+            negative_ion_resname="CL",
+            water_resname="HOH",
         )
 
 
@@ -2036,9 +2036,9 @@ def test_handle_alchemwats_too_many_nbf(
             system=new_system,
             system_mapping={},
             charge_difference=1,
-            positive_ion_resname='NA',
-            negative_ion_resname='CL',
-            water_resname='HOH',
+            positive_ion_resname="NA",
+            negative_ion_resname="CL",
+            water_resname="HOH",
         )
 
 
@@ -2060,9 +2060,9 @@ def test_handle_alchemwats_vsite_water(
             system=system,
             system_mapping={},
             charge_difference=1,
-            positive_ion_resname='NA',
-            negative_ion_resname='CL',
-            water_resname='HOH',
+            positive_ion_resname="NA",
+            negative_ion_resname="CL",
+            water_resname="HOH",
         )
 
 
@@ -2090,9 +2090,9 @@ def test_handle_alchemwats_incorrect_atom(
             system=new_system,
             system_mapping=benzene_self_system_mapping,
             charge_difference=1,
-            positive_ion_resname='NA',
-            negative_ion_resname='CL',
-            water_resname='HOH',
+            positive_ion_resname="NA",
+            negative_ion_resname="CL",
+            water_resname="HOH",
         )
 
 
@@ -2111,9 +2111,9 @@ def test_handle_alchemical_wats(
         system=system,
         system_mapping=benzene_self_system_mapping,
         charge_difference=1,
-        positive_ion_resname='NA',
-        negative_ion_resname='CL',
-        water_resname='HOH',
+        positive_ion_resname="NA",
+        negative_ion_resname="CL",
+        water_resname="HOH",
     )
 
     # check the mappings

From f1ca989e9224ab5239f6788c82553f5894d341a1 Mon Sep 17 00:00:00 2001
From: hannahbaumann <hbaumann@uci.edu>
Date: Thu, 21 May 2026 15:47:39 +0200
Subject: [PATCH 3/9] Get ion parameters from ff if no ions present in system

---
 .../openmm_rfe/_rfe_utils/topologyhelpers.py  | 263 ++++++++++--------
 .../openmm_rfe/hybridtop_protocols.py         |  24 +-
 .../protocols/openmm_rfe/hybridtop_units.py   |  20 +-
 .../openmm_rfe/test_hybrid_top_protocol.py    |  95 +++++--
 4 files changed, 226 insertions(+), 176 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index 2a4e6069b..06a7a6212 100644
--- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -22,74 +22,165 @@
 from openmm import NonbondedForce, System, app
 from openmm import unit as omm_unit
 
-from openfe import SolventComponent
-
 logger = logging.getLogger(__name__)
 
 
+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_ion_and_water_parameters(
     topology: app.Topology,
     system: System,
-    ion_resname: str,
+    forcefield: app.ForceField,
+    charge_difference: int,
     water_resname: str = 'HOH',
-):
+) -> tuple:
     """
-    Get ion, and water (oxygen and hydrogen) atoms parameters.
+    Get ion NonbondedForce parameters and water charges.
+
+    Scans the topology for the most abundant monovalent ion of the appropriate
+    charge sign, identified by single-atom residues with a partial charge
+    close to +/- 1.
+    If no matching ion is found, falls back to NA or CL.
 
     Parameters
     ----------
     topology : app.Topology
-      The topology to search for the ion and water
-    system : app.System
-      The system associated with the input topology object.
-    ion_resname : str
-      The residue name of the ion to get parameters for
+        The topology to search.
+    system : openmm.System
+        The system associated with the topology.
+    forcefield : app.ForceField
+        Used as a fallback to obtain ion parameters when no matching ion is
+        present in the topology.
+    charge_difference : int
+        Charge difference between state A and state B, used to determine the
+        required ion charge sign.
     water_resname : str
-      The residue name of the water to get parameters for. Default 'HOH'.
+        Residue name of 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
-      The partial charge of the water oxygen.
-    h_charge : float
-      The partial charge of the water hydrogen.
+    ion_charge, ion_sigma, ion_epsilon : openmm.unit.Quantity
+        NonbondedForce parameters for the ion.
+    o_charge, h_charge : openmm.unit.Quantity
+        Partial charges of the water oxygen and hydrogen.
+    """
+    nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
 
-    Raises
-    ------
-    ValueError
-      If there are no ``ion_resname`` or ``water_resname`` named residues in
-      the input ``topology``.
+    desired_sign = np.sign(charge_difference)
 
-    Attribution
-    -----------
-    Based on `perses.utils.charge_changing.get_ion_and_water_parameters`.
-    """
-    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)
+    ion_counts = Counter()
+    ion_atom_indices = {}
+
+    o_charge = None
+    h_charge = None
+
+    for residue in topology.residues():
+        atoms = list(residue.atoms())
+
+        # water parameters
+        if o_charge is None and residue.name == water_resname:
+            for atom in atoms:
+                charge, _, _ = nbf.getParticleParameters(atom.index)
+                if atom.element.symbol == 'O':
+                    o_charge = charge
+                elif atom.element.symbol == 'H':
+                    h_charge = charge
+            continue
+
+        # ion candidates
+        if len(atoms) != 1:
+            continue
 
-    ion_index = _find_atom(topology, ion_resname, None)
-    oxygen_index = _find_atom(topology, water_resname, 'O')
-    hydrogen_index = _find_atom(topology, water_resname, 'H')
+        atom = atoms[0]
+        if atom.element is None:
+            continue
 
-    nbf = [i for i in system.getForces()
-           if isinstance(i, NonbondedForce)][0]
+        charge, _, _ = nbf.getParticleParameters(atom.index)
+        charge_val = charge.value_in_unit(omm_unit.elementary_charge)
 
-    ion_charge, ion_sigma, ion_epsilon = nbf.getParticleParameters(ion_index)
-    o_charge, _, _ = nbf.getParticleParameters(oxygen_index)
-    h_charge, _, _ = nbf.getParticleParameters(hydrogen_index)
+        # Only monovalent ions for now
+        if not np.isclose(abs(charge_val), 1.0, atol=0.2):
+            continue
+
+        if np.sign(charge_val) != desired_sign:
+            continue
+
+        ion_counts[residue.name] += 1
+        ion_atom_indices.setdefault(
+            residue.name,
+            atom.index,
+        )
+
+    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."
+        )
+
+    if ion_counts:
+        # Use the most abundant matching ion found in the topology
+        best_resname = ion_counts.most_common(1)[0][0]
+        ion_atom_index = ion_atom_indices[best_resname]
+        ion_charge, ion_sigma, ion_epsilon = nbf.getParticleParameters(
+            ion_atom_index
+        )
+    else:
+        # 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
+
+        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)
+
+        ion_charge, ion_sigma, ion_epsilon = _get_ion_parameters_from_forcefield(
+            forcefield, fallback_resname, fallback_element
+        )
 
     return ion_charge, ion_sigma, ion_epsilon, o_charge, h_charge
 
@@ -129,8 +220,7 @@ def handle_alchemical_waters(
     system: System,
     system_mapping: dict,
     charge_difference: int,
-    positive_ion_resname: str,
-    negative_ion_resname: str,
+    forcefield: app.ForceField,
     water_resname: str,
 ):
     """
@@ -148,12 +238,8 @@ 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.
 
@@ -176,17 +262,11 @@ def handle_alchemical_waters(
                   f"difference: {abs(charge_difference)}")
         raise ValueError(errmsg)
 
-    if charge_difference > 0:
-        ion_resname = positive_ion_resname
-    elif charge_difference < 0:
-        ion_resname = negative_ion_resname
-    # 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,
-        water_resname,
+        topology, system, forcefield, charge_difference, water_resname='HOH',
     )
 
     # get the nonbonded forces
@@ -724,60 +804,3 @@ def set_and_check_new_positions(mapping, old_topology, new_topology,
             logging.warning(wmsg)
 
     return new_pos_array * omm_unit.angstrom
-
-
-def _get_ion_resnames_from_topology(topology: app.Topology) -> tuple[str, str]:
-    """
-    Infer positive and negative ion residue names from a topology by
-    finding the most common monovalent ion of each charge type.
-    Falls back to NA/CL if none are found.
-
-    Parameters
-    ----------
-    topology : app.Topology
-        The topology to search for ions.
-
-    Returns
-    -------
-    pos_ion : str
-        The residue name of the most abundant positive monovalent ion (Na, K).
-    neg_ion : str
-        The residue name of the most abundant negative monovalent ion (Cl).
-    """
-    known_positive = {'NA', 'K'}
-    # This doesn't make much sense yet to check it, since it's only Cl, but
-    # leaving it here for now so we can add other neg ions.
-    known_negative = {'CL'}
-
-    pos_counts = Counter(
-        r.name for r in topology.residues() if r.name in known_positive
-    )
-    neg_counts = Counter(
-        r.name for r in topology.residues() if r.name in known_negative
-    )
-
-    if not pos_counts:
-        wmsg = (
-            "Could not find any known positive monovalent ions "
-            f"(searched for {known_positive}) in the topology. "
-            "Defaulting to NA for explicit charge correction."
-        )
-        warnings.warn(wmsg)
-        logger.warning(wmsg)
-        pos_ion = 'NA'
-    else:
-        pos_ion = max(pos_counts, key=pos_counts.get)
-
-    if not neg_counts:
-        wmsg = (
-            "Could not find any known negative monovalent ions "
-            f"(searched for {known_negative}) in the topology. "
-            "Defaulting to CL for explicit charge correction."
-        )
-        warnings.warn(wmsg)
-        logger.warning(wmsg)
-        neg_ion = 'CL'
-    else:
-        neg_ion = max(neg_counts, key=neg_counts.get)
-
-    return pos_ion, neg_ion
diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
index 4be7eff5a..84c051095 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
+++ b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
@@ -374,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.
@@ -441,18 +441,18 @@ def _validate_charge_difference(
         if isinstance(solvent_component, SolventComponent):
             positive_ion = solvent_component.positive_ion.strip("-+").upper()
             negative_ion = solvent_component.negative_ion.strip("-+").upper()
+            ion = {-1: positive_ion, 1: 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"
+            )
         elif isinstance(solvent_component, SolvatedPDBComponent):
-            positive_ion, negative_ion = _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
-                solvent_component.to_openmm_topology()
+            wmsg = (
+                f"A charge difference of {difference} is observed "
+                "between the end states. This will be addressed by "
+                "transforming a water into an ion."
             )
-
-        ion = {-1: positive_ion, 1: 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.info(wmsg)
 
     @staticmethod
@@ -571,6 +571,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
 
diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_units.py b/src/openfe/protocols/openmm_rfe/hybridtop_units.py
index 988e0a990..779cf8cee 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_units.py
+++ b/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,22 +404,12 @@ 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:
             return
 
-        # Resolve ion names from the solvent component if available,
-        # otherwise infer from the topology (for explicitly solvated systems).
-        if isinstance(solvent_component, SolventComponent):
-            positive_ion = solvent_component.positive_ion.strip("-+").upper()
-            negative_ion = solvent_component.negative_ion.strip("-+").upper()
-        else:
-            positive_ion, negative_ion = _rfe_utils.topologyhelpers._get_ion_resnames_from_topology(
-                stateA_topology
-            )
-
         # Get the residue ids for waters to turn alchemical
         alchem_water_resids = _rfe_utils.topologyhelpers.get_alchemical_waters(
             topology=stateA_topology,
@@ -435,8 +425,7 @@ def _handle_net_charge(
             system=stateB_system,
             system_mapping=system_mappings,
             charge_difference=charge_difference,
-            positive_ion_resname=positive_ion,
-            negative_ion_resname=negative_ion,
+            forcefield=forcefield,
             water_resname="HOH",  # Need to check if this is also true for systems not prepped with addSolvent
         )
 
@@ -557,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,
@@ -565,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
diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
index 00c44ac14..7a81c876a 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
@@ -1890,7 +1890,7 @@ def benzene_solvent_openmm_system(benzene_modifications):
     positions = to_openmm(from_openmm(modeller.getPositions()))
     system = system_generator.create_system(topology, molecules=[offmol])
 
-    return system, topology, positions
+    return system, topology, positions, system_generator.forcefield
 
 
 @pytest.fixture(scope="session")
@@ -1940,7 +1940,7 @@ def benzene_self_system_mapping(benzene_solvent_openmm_system):
     alchemical transformation (this technically doesn't work in practice
     because the RFE protocol expects an alchemical component).
     """
-    system, topology, positions = benzene_solvent_openmm_system
+    system, topology, positions, forcefield = benzene_solvent_openmm_system
 
     res = [r for r in topology.residues()]
     benzene_res = [r for r in res if r.name == "UNK"][0]
@@ -1962,28 +1962,67 @@ def benzene_self_system_mapping(benzene_solvent_openmm_system):
     return system_mapping
 
 
-@pytest.mark.parametrize(
-    "ion, water",
-    [
-        ["NA", "SOL"],
-        ["NX", "WAT"],
-    ],
-)
-def test_get_ion_water_parameters_unknownresname(ion, water, benzene_solvent_openmm_system):
-    system, topology, positions = benzene_solvent_openmm_system
+def test_get_ion_water_parameters_fallback_to_forcefield(benzene_modifications):
+    """
+    Check that when no matching ion is found in the topology,
+    parameters are obtained from the forcefield (NA/CL fallback).
+    """
+    smc = benzene_modifications["benzene"]
+    offmol = smc.to_openff()
+    settings = openmm_rfe.RelativeHybridTopologyProtocol.default_settings()
+
+    system_generator = system_creation.get_system_generator(
+        forcefield_settings=settings.forcefield_settings,
+        integrator_settings=settings.integrator_settings,
+        thermo_settings=settings.thermo_settings,
+        cache=None,
+        has_solvent=True,
+    )
+    system_generator.create_system(
+        offmol.to_topology().to_openmm(),
+        molecules=[offmol],
+    )
+
+    modeller, _ = system_creation.get_omm_modeller(
+        protein_comp=None,
+        solvent_comp=openfe.SolventComponent(ion_concentration=0 * unit.molar),
+        small_mols={smc: offmol},
+        omm_forcefield=system_generator.forcefield,
+        solvent_settings=settings.solvation_settings,
+    )
+
+    topology = modeller.getTopology()
+    system = system_generator.create_system(topology, molecules=[offmol])
+
+    with pytest.warns(UserWarning, match="No monovalent ion"):
+        ion_charge, ion_sigma, ion_epsilon, o_charge, h_charge = (
+            topologyhelpers._get_ion_and_water_parameters(
+                topology, system, system_generator.forcefield,
+                charge_difference=1,
+            )
+        )
 
-    errmsg = "Error encountered when attempting to explicitly handle"
+    assert ion_charge.value_in_unit(omm_unit.elementary_charge) == pytest.approx(1.0)
 
+
+def test_get_ion_water_parameters_bad_water(benzene_solvent_openmm_system):
+    """
+    Check that a ValueError is raised when the water residue name
+    is not found in the topology.
+    """
+    system, topology, positions, forcefield = benzene_solvent_openmm_system
+
+    errmsg = "Could not find water residue"
     with pytest.raises(ValueError, match=errmsg):
         topologyhelpers._get_ion_and_water_parameters(
-            topology, system, ion_resname=ion, water_resname=water
+            topology, system, forcefield, charge_difference=1, water_resname="SOL"
         )
 
 
 def test_get_alchemical_waters_no_waters(
     benzene_solvent_openmm_system,
 ):
-    system, topology, positions = benzene_solvent_openmm_system
+    system, topology, positions, forcefield = benzene_solvent_openmm_system
 
     errmsg = "There are no waters"
 
@@ -1999,7 +2038,7 @@ def test_handle_alchemwats_incorrect_count(
     """
     Check that an error is thrown when charge_difference != len(water_resids)
     """
-    system, topology, positions = benzene_solvent_openmm_system
+    system, topology, positions, forcefield = benzene_solvent_openmm_system
 
     errmsg = "There should be as many alchemical water residues:"
 
@@ -2010,8 +2049,7 @@ def test_handle_alchemwats_incorrect_count(
             system=system,
             system_mapping={},
             charge_difference=1,
-            positive_ion_resname="NA",
-            negative_ion_resname="CL",
+            forcefield=forcefield,
             water_resname="HOH",
         )
 
@@ -2022,7 +2060,7 @@ def test_handle_alchemwats_too_many_nbf(
     """
     Check that an error is thrown when there are multiple NonbondedForces
     """
-    system, topology, positions = benzene_solvent_openmm_system
+    system, topology, positions, forcefield = benzene_solvent_openmm_system
 
     new_system = copy.deepcopy(system)
     new_system.addForce(NonbondedForce())
@@ -2036,8 +2074,7 @@ def test_handle_alchemwats_too_many_nbf(
             system=new_system,
             system_mapping={},
             charge_difference=1,
-            positive_ion_resname="NA",
-            negative_ion_resname="CL",
+            forcefield=forcefield,
             water_resname="HOH",
         )
 
@@ -2060,8 +2097,7 @@ def test_handle_alchemwats_vsite_water(
             system=system,
             system_mapping={},
             charge_difference=1,
-            positive_ion_resname="NA",
-            negative_ion_resname="CL",
+            forcefield=app.ForceField(),
             water_resname="HOH",
         )
 
@@ -2073,7 +2109,7 @@ def test_handle_alchemwats_incorrect_atom(
     """
     Check that an error is thrown when charge_difference != len(water_resids)
     """
-    system, topology, positions = benzene_solvent_openmm_system
+    system, topology, positions, forcefield = benzene_solvent_openmm_system
 
     # modify the charge of hydrogen atom 25
     new_system = copy.deepcopy(system)  # protect the session scoped object
@@ -2090,8 +2126,7 @@ def test_handle_alchemwats_incorrect_atom(
             system=new_system,
             system_mapping=benzene_self_system_mapping,
             charge_difference=1,
-            positive_ion_resname="NA",
-            negative_ion_resname="CL",
+            forcefield=forcefield,
             water_resname="HOH",
         )
 
@@ -2100,7 +2135,7 @@ def test_handle_alchemical_wats(
     benzene_solvent_openmm_system,
     benzene_self_system_mapping,
 ):
-    system, topology, positions = benzene_solvent_openmm_system
+    system, topology, positions, forcefield = benzene_solvent_openmm_system
 
     n_env = len(benzene_self_system_mapping["old_to_new_env_atom_map"])
     n_core = len(benzene_self_system_mapping["old_to_new_core_atom_map"])
@@ -2111,8 +2146,7 @@ def test_handle_alchemical_wats(
         system=system,
         system_mapping=benzene_self_system_mapping,
         charge_difference=1,
-        positive_ion_resname="NA",
-        negative_ion_resname="CL",
+        forcefield=forcefield,
         water_resname="HOH",
     )
 
@@ -2129,10 +2163,11 @@ def test_handle_alchemical_wats(
     # system parameters checks
     nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
     # check the oxygen parameters
-    i_chg, i_sig, i_eps, o_chg, h_chg = topologyhelpers._get_ion_and_water_parameters(
-        topology, system, "NA", "HOH"
+    i_chg, i_sig, i_eps, o_chg, h_chg = (topologyhelpers._get_ion_and_water_parameters(
+        topology, system, forcefield, charge_difference=1)
     )
 
+
     charge, sigma, epsilon = nbf.getParticleParameters(24)
     assert charge == 1.0 * omm_unit.elementary_charge == i_chg
     assert sigma == i_sig

From 7f06c1cc54e444f53cf4cffb6b01de3a5564704c Mon Sep 17 00:00:00 2001
From: hannahbaumann <hbaumann@uci.edu>
Date: Thu, 21 May 2026 16:24:56 +0200
Subject: [PATCH 4/9] Split of water and ion into separate functions

---
 .../openmm_rfe/_rfe_utils/topologyhelpers.py  | 188 ++++++++++--------
 .../openmm_rfe/test_hybrid_top_protocol.py    |  22 +-
 .../openmm_rfe/test_hybrid_top_validation.py  |   2 +-
 3 files changed, 111 insertions(+), 101 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index 06a7a6212..3fb376faf 100644
--- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -62,128 +62,139 @@ def _get_ion_parameters_from_forcefield(
 
     return nbf.getParticleParameters(0)
 
-
-def _get_ion_and_water_parameters(
+def _get_water_parameters(
     topology: app.Topology,
     system: System,
-    forcefield: app.ForceField,
-    charge_difference: int,
     water_resname: str = 'HOH',
 ) -> tuple:
     """
-    Get ion NonbondedForce parameters and water charges.
-
-    Scans the topology for the most abundant monovalent ion of the appropriate
-    charge sign, identified by single-atom residues with a partial charge
-    close to +/- 1.
-    If no matching ion is found, falls back to NA or CL.
+    Get water oxygen and hydrogen partial charges from the system.
 
     Parameters
     ----------
     topology : app.Topology
-        The topology to search.
-    system : openmm.System
-        The system associated with the topology.
-    forcefield : app.ForceField
-        Used as a fallback to obtain ion parameters when no matching ion is
-        present in the topology.
-    charge_difference : int
-        Charge difference between state A and state B, used to determine the
-        required ion charge sign.
+      The topology to search for water atoms.
+    system : app.System
+      The system associated with the input topology object.
     water_resname : str
-        Residue name of water. Default 'HOH'.
+      The residue name of the water. Default 'HOH'.
 
     Returns
     -------
-    ion_charge, ion_sigma, ion_epsilon : openmm.unit.Quantity
-        NonbondedForce parameters for the ion.
-    o_charge, h_charge : openmm.unit.Quantity
-        Partial charges of the water oxygen and hydrogen.
+    o_charge : openmm.unit.Quantity
+      The partial charge of the water oxygen.
+    h_charge : openmm.unit.Quantity
+      The partial charge of the water hydrogen.
+
+    Raises
+    ------
+    ValueError
+      If no water residue named ``water_resname`` with O and H atoms
+      is found in the topology.
     """
     nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
 
-    desired_sign = np.sign(charge_difference)
-
-    ion_counts = Counter()
-    ion_atom_indices = {}
-
     o_charge = None
     h_charge = None
 
     for residue in topology.residues():
-        atoms = list(residue.atoms())
-
-        # water parameters
-        if o_charge is None and residue.name == water_resname:
-            for atom in atoms:
-                charge, _, _ = nbf.getParticleParameters(atom.index)
-                if atom.element.symbol == 'O':
-                    o_charge = charge
-                elif atom.element.symbol == 'H':
-                    h_charge = charge
+        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.
+
+    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.
+
+    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]
 
-        # ion candidates
+    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)
-
-        # Only monovalent ions for now
-        if not np.isclose(abs(charge_val), 1.0, atol=0.2):
-            continue
-
-        if np.sign(charge_val) != desired_sign:
-            continue
-
-        ion_counts[residue.name] += 1
-        ion_atom_indices.setdefault(
-            residue.name,
-            atom.index,
-        )
-
-    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."
-        )
+        if (abs(abs(charge_val) - 1.0) < 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]
-        ion_atom_index = ion_atom_indices[best_resname]
-        ion_charge, ion_sigma, ion_epsilon = nbf.getParticleParameters(
-            ion_atom_index
-        )
-    else:
-        # 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
-
-        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 nbf.getParticleParameters(ion_atom_indices[best_resname])
 
-        ion_charge, ion_sigma, ion_epsilon = _get_ion_parameters_from_forcefield(
-            forcefield, fallback_resname, fallback_element
-        )
+    # 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]]],
@@ -265,9 +276,10 @@ def handle_alchemical_waters(
     if charge_difference == 0:
         return None
 
-    ion_charge, ion_sigma, ion_epsilon, o_charge, h_charge = _get_ion_and_water_parameters(
-        topology, system, forcefield, charge_difference, water_resname='HOH',
+    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()
diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
index 7a81c876a..b77a098ce 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
@@ -1962,10 +1962,10 @@ def benzene_self_system_mapping(benzene_solvent_openmm_system):
     return system_mapping
 
 
-def test_get_ion_water_parameters_fallback_to_forcefield(benzene_modifications):
+def test_get_ion_parameters_fallback_to_forcefield(benzene_modifications):
     """
     Check that when no matching ion is found in the topology,
-    parameters are obtained from the forcefield (NA/CL fallback).
+    parameters are obtained from the forcefield (NA fallback).
     """
     smc = benzene_modifications["benzene"]
     offmol = smc.to_openff()
@@ -1995,17 +1995,15 @@ def test_get_ion_water_parameters_fallback_to_forcefield(benzene_modifications):
     system = system_generator.create_system(topology, molecules=[offmol])
 
     with pytest.warns(UserWarning, match="No monovalent ion"):
-        ion_charge, ion_sigma, ion_epsilon, o_charge, h_charge = (
-            topologyhelpers._get_ion_and_water_parameters(
-                topology, system, system_generator.forcefield,
-                charge_difference=1,
-            )
+        ion_charge, ion_sigma, ion_epsilon = topologyhelpers._get_ion_parameters(
+            topology, system, charge_difference=1,
+            forcefield=system_generator.forcefield,
         )
 
     assert ion_charge.value_in_unit(omm_unit.elementary_charge) == pytest.approx(1.0)
 
 
-def test_get_ion_water_parameters_bad_water(benzene_solvent_openmm_system):
+def test_get_water_parameters_bad_water(benzene_solvent_openmm_system):
     """
     Check that a ValueError is raised when the water residue name
     is not found in the topology.
@@ -2014,8 +2012,8 @@ def test_get_ion_water_parameters_bad_water(benzene_solvent_openmm_system):
 
     errmsg = "Could not find water residue"
     with pytest.raises(ValueError, match=errmsg):
-        topologyhelpers._get_ion_and_water_parameters(
-            topology, system, forcefield, charge_difference=1, water_resname="SOL"
+        topologyhelpers._get_water_parameters(
+            topology, system, water_resname="SOL"
         )
 
 
@@ -2163,8 +2161,8 @@ def test_handle_alchemical_wats(
     # system parameters checks
     nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
     # check the oxygen parameters
-    i_chg, i_sig, i_eps, o_chg, h_chg = (topologyhelpers._get_ion_and_water_parameters(
-        topology, system, forcefield, charge_difference=1)
+    i_chg, i_sig, i_eps = topologyhelpers._get_ion_parameters(
+        topology, system, charge_difference=1, forcefield=forcefield
     )
 
 
diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py
index 6d5db1249..c29b1b37f 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py
@@ -441,7 +441,7 @@ def test_get_charge_difference(mapping_name, result, request, caplog):
     mapping = request.getfixturevalue(mapping_name)
     caplog.set_level(logging.INFO)
 
-    ion = r"Na+" if result == -1 else r"Cl-"
+    ion = r"NA" if result == -1 else r"CL"
     msg = (
         f"A charge difference of {result} is observed "
         "between the end states. This will be addressed by "

From f508d3442e903744fa27fdf6246332c215600bba Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 21 May 2026 14:25:50 +0000
Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../protocols/openmm_rfe/test_hybrid_top_protocol.py     | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
index b77a098ce..58d5cf00c 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
@@ -1996,7 +1996,9 @@ def test_get_ion_parameters_fallback_to_forcefield(benzene_modifications):
 
     with pytest.warns(UserWarning, match="No monovalent ion"):
         ion_charge, ion_sigma, ion_epsilon = topologyhelpers._get_ion_parameters(
-            topology, system, charge_difference=1,
+            topology,
+            system,
+            charge_difference=1,
             forcefield=system_generator.forcefield,
         )
 
@@ -2012,9 +2014,7 @@ def test_get_water_parameters_bad_water(benzene_solvent_openmm_system):
 
     errmsg = "Could not find water residue"
     with pytest.raises(ValueError, match=errmsg):
-        topologyhelpers._get_water_parameters(
-            topology, system, water_resname="SOL"
-        )
+        topologyhelpers._get_water_parameters(topology, system, water_resname="SOL")
 
 
 def test_get_alchemical_waters_no_waters(
@@ -2165,7 +2165,6 @@ def test_handle_alchemical_wats(
         topology, system, charge_difference=1, forcefield=forcefield
     )
 
-
     charge, sigma, epsilon = nbf.getParticleParameters(24)
     assert charge == 1.0 * omm_unit.elementary_charge == i_chg
     assert sigma == i_sig

From 2cb287af63f1b8a212cf5f0c47b5e631d3b3fa64 Mon Sep 17 00:00:00 2001
From: hannahbaumann <hbaumann@uci.edu>
Date: Mon, 1 Jun 2026 10:18:19 +0200
Subject: [PATCH 6/9] Address review comments

---
 .../openmm_rfe/_rfe_utils/topologyhelpers.py  |  2 +-
 .../openmm_rfe/hybridtop_protocols.py         | 21 +++++--------------
 .../openmm_rfe/test_hybrid_top_validation.py  |  3 +--
 3 files changed, 7 insertions(+), 19 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index 3fb376faf..a4e7cc693 100644
--- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -166,7 +166,7 @@ def _get_ion_parameters(
             continue
         charge, _, _ = nbf.getParticleParameters(atom.index)
         charge_val = charge.value_in_unit(omm_unit.elementary_charge)
-        if (abs(abs(charge_val) - 1.0) < 0.01 and np.sign(charge_val) == desired_sign):
+        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
diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
index 84c051095..35ef51375 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
+++ b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py
@@ -437,22 +437,11 @@ def _validate_charge_difference(
             )
             raise ValueError(errmsg)
 
-        # resolve ion names from SolventComponent or topology
-        if isinstance(solvent_component, SolventComponent):
-            positive_ion = solvent_component.positive_ion.strip("-+").upper()
-            negative_ion = solvent_component.negative_ion.strip("-+").upper()
-            ion = {-1: positive_ion, 1: 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"
-            )
-        elif isinstance(solvent_component, SolvatedPDBComponent):
-            wmsg = (
-                f"A charge difference of {difference} is observed "
-                "between the end states. This will be addressed by "
-                "transforming a water into an ion."
-            )
+        wmsg = (
+            f"A charge difference of {difference} is observed "
+            "between the end states. This will be addressed by "
+            "transforming a water into an ion of opposite charge."
+        )
         logger.info(wmsg)
 
     @staticmethod
diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py
index c29b1b37f..f9d219e34 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_validation.py
@@ -441,11 +441,10 @@ def test_get_charge_difference(mapping_name, result, request, caplog):
     mapping = request.getfixturevalue(mapping_name)
     caplog.set_level(logging.INFO)
 
-    ion = r"NA" if result == -1 else r"CL"
     msg = (
         f"A charge difference of {result} is observed "
         "between the end states. This will be addressed by "
-        f"transforming a water into a {ion} ion"
+        f"transforming a water into an ion of opposite charge"
     )
 
     openmm_rfe.RelativeHybridTopologyProtocol._validate_charge_difference(

From 3d96013fb842390ec7285fccf525ad4c461b5f14 Mon Sep 17 00:00:00 2001
From: hannahbaumann <hbaumann@uci.edu>
Date: Thu, 4 Jun 2026 11:48:18 +0200
Subject: [PATCH 7/9] Also get water parameters from ff

---
 .../openmm_rfe/_rfe_utils/topologyhelpers.py  | 136 +++++-------------
 .../protocols/openmm_rfe/hybridtop_units.py   |   1 -
 .../openmm_rfe/test_hybrid_top_protocol.py    |  17 ---
 3 files changed, 35 insertions(+), 119 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index a4e7cc693..5c689eccb 100644
--- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -25,101 +25,36 @@
 logger = logging.getLogger(__name__)
 
 
-def _get_ion_parameters_from_forcefield(
+def _get_ff_parameters(
     forcefield: app.ForceField,
-    ion_resname: str,
-    ion_element: app.Element,
-) -> tuple:
+    smiles: str,
+) -> list[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.
+    Get NonbondedForce parameters for all atoms in a molecule by creating
+    a minimal dummy system from a SMILES string with the given forcefield.
 
     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.
+      The force field to use for parameterisation.
+    smiles : str
+      The SMILES string of the molecule to parameterise
+      (e.g. '[Na+]', '[Cl-]', 'O').
 
     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,
-    water_resname: str = 'HOH',
-) -> tuple:
-    """
-    Get water oxygen and hydrogen partial charges from the system.
-
-    Parameters
-    ----------
-    topology : app.Topology
-      The topology to search for water atoms.
-    system : app.System
-      The system associated with the input topology object.
-    water_resname : str
-      The residue name of the water. Default 'HOH'.
-
-    Returns
-    -------
-    o_charge : openmm.unit.Quantity
-      The partial charge of the water oxygen.
-    h_charge : openmm.unit.Quantity
-      The partial charge of the water hydrogen.
-
-    Raises
-    ------
-    ValueError
-      If no water residue named ``water_resname`` with O and H atoms
-      is found in the topology.
+    list[tuple]
+      A list of (charge, sigma, epsilon) tuples for each atom in the
+      molecule, in the same order as the SMILES.
     """
-    nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
+    from openff.toolkit import Molecule
 
-    o_charge = None
-    h_charge = None
+    offmol = Molecule.from_smiles(smiles)
+    dummy_top = offmol.to_topology().to_openmm()
+    dummy_system = forcefield.createSystem(dummy_top)
+    nbf = [f for f in dummy_system.getForces() if isinstance(f, NonbondedForce)][0]
 
-    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
+    return [nbf.getParticleParameters(i) for i in range(dummy_top.getNumAtoms())]
 
 def _get_ion_parameters(
     topology: app.Topology,
@@ -178,23 +113,19 @@ def _get_ion_parameters(
 
     # 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
+        fallback_smiles = '[Na+]'
     else:
-        fallback_resname = 'CL'
-        fallback_element = app.Element.getByAtomicNumber(17)  # Cl
+        fallback_smiles = '[Cl-]'
 
     wmsg = (
         f"No monovalent ion with the appropriate charge sign found in "
-        f"the topology. Defaulting to '{fallback_resname}' and obtaining "
+        f"the topology. Defaulting to '{fallback_smiles}' 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,
-    )
+    return _get_ff_parameters(forcefield, fallback_smiles)[0]
 
 def _fix_alchemical_water_atom_mapping(
     system_mapping: dict[str, Union[dict[int, int], list[int]]],
@@ -232,7 +163,6 @@ def handle_alchemical_waters(
     system_mapping: dict,
     charge_difference: int,
     forcefield: app.ForceField,
-    water_resname: str,
 ):
     """
     Add alchemical waters from a pre-defined list.
@@ -251,8 +181,6 @@ def handle_alchemical_waters(
       The charge difference between state A and state B.
     forcefield : app.ForceField
       The forcefield to use for ion parameterization.
-    water_resname : str
-      The residue name of the water to get parameters for.
 
     Raises
     ------
@@ -276,11 +204,6 @@ def handle_alchemical_waters(
     if charge_difference == 0:
         return None
 
-    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()
               if isinstance(i, NonbondedForce)]
@@ -290,8 +213,7 @@ def handle_alchemical_waters(
     # for convenience just grab the first & only entry
     nbf = nbfrcs[0]
 
-    # Loop through residues, check if they match the residue index
-    # mutate the atom as necessary
+    # Check for virtual site waters
     for res in topology.residues():
         if res.index in water_resids:
             # if the number of atoms > 3, then we have virtual sites which are
@@ -301,6 +223,18 @@ def handle_alchemical_waters(
                           "are not currently supported as alchemical waters")
                 raise ValueError(errmsg)
 
+    ion_charge, ion_sigma, ion_epsilon = _get_ion_parameters(
+        topology, system, charge_difference, forcefield,
+    )
+    # Get water parameters (O is index 0, H is index 1 and 2)
+    water_params = _get_ff_parameters(forcefield, '[H]O[H]')
+    o_charge, _, _ = water_params[0]
+    h_charge, _, _ = water_params[1]
+
+    # Loop through residues, check if they match the residue index
+    # mutate the atom as necessary
+    for res in topology.residues():
+        if res.index in water_resids:
             for at in res.atoms():
                 idx = at.index
                 charge, sigma, epsilon = nbf.getParticleParameters(idx)
diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_units.py b/src/openfe/protocols/openmm_rfe/hybridtop_units.py
index 779cf8cee..5f4c1c9f3 100644
--- a/src/openfe/protocols/openmm_rfe/hybridtop_units.py
+++ b/src/openfe/protocols/openmm_rfe/hybridtop_units.py
@@ -426,7 +426,6 @@ def _handle_net_charge(
             system_mapping=system_mappings,
             charge_difference=charge_difference,
             forcefield=forcefield,
-            water_resname="HOH",  # Need to check if this is also true for systems not prepped with addSolvent
         )
 
     def _get_omm_objects(
diff --git a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
index 58d5cf00c..62f3072c8 100644
--- a/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
+++ b/src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py
@@ -2005,18 +2005,6 @@ def test_get_ion_parameters_fallback_to_forcefield(benzene_modifications):
     assert ion_charge.value_in_unit(omm_unit.elementary_charge) == pytest.approx(1.0)
 
 
-def test_get_water_parameters_bad_water(benzene_solvent_openmm_system):
-    """
-    Check that a ValueError is raised when the water residue name
-    is not found in the topology.
-    """
-    system, topology, positions, forcefield = benzene_solvent_openmm_system
-
-    errmsg = "Could not find water residue"
-    with pytest.raises(ValueError, match=errmsg):
-        topologyhelpers._get_water_parameters(topology, system, water_resname="SOL")
-
-
 def test_get_alchemical_waters_no_waters(
     benzene_solvent_openmm_system,
 ):
@@ -2048,7 +2036,6 @@ def test_handle_alchemwats_incorrect_count(
             system_mapping={},
             charge_difference=1,
             forcefield=forcefield,
-            water_resname="HOH",
         )
 
 
@@ -2073,7 +2060,6 @@ def test_handle_alchemwats_too_many_nbf(
             system_mapping={},
             charge_difference=1,
             forcefield=forcefield,
-            water_resname="HOH",
         )
 
 
@@ -2096,7 +2082,6 @@ def test_handle_alchemwats_vsite_water(
             system_mapping={},
             charge_difference=1,
             forcefield=app.ForceField(),
-            water_resname="HOH",
         )
 
 
@@ -2125,7 +2110,6 @@ def test_handle_alchemwats_incorrect_atom(
             system_mapping=benzene_self_system_mapping,
             charge_difference=1,
             forcefield=forcefield,
-            water_resname="HOH",
         )
 
 
@@ -2145,7 +2129,6 @@ def test_handle_alchemical_wats(
         system_mapping=benzene_self_system_mapping,
         charge_difference=1,
         forcefield=forcefield,
-        water_resname="HOH",
     )
 
     # check the mappings

From 656b41e1cc436cb0efddad027782c4ad6c51e270 Mon Sep 17 00:00:00 2001
From: hannahbaumann <hbaumann@uci.edu>
Date: Thu, 4 Jun 2026 13:36:16 +0200
Subject: [PATCH 8/9] Some fixes

---
 .../openmm_rfe/_rfe_utils/topologyhelpers.py  | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index 5c689eccb..b3ff708b2 100644
--- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -28,7 +28,7 @@
 def _get_ff_parameters(
     forcefield: app.ForceField,
     smiles: str,
-) -> list[tuple]:
+) -> dict[str, tuple]:
     """
     Get NonbondedForce parameters for all atoms in a molecule by creating
     a minimal dummy system from a SMILES string with the given forcefield.
@@ -43,9 +43,9 @@ def _get_ff_parameters(
 
     Returns
     -------
-    list[tuple]
-      A list of (charge, sigma, epsilon) tuples for each atom in the
-      molecule, in the same order as the SMILES.
+    dict[str, tuple]
+      A dictionary keyed by element symbol, with (charge, sigma, epsilon)
+      tuples as values.
     """
     from openff.toolkit import Molecule
 
@@ -54,7 +54,10 @@ def _get_ff_parameters(
     dummy_system = forcefield.createSystem(dummy_top)
     nbf = [f for f in dummy_system.getForces() if isinstance(f, NonbondedForce)][0]
 
-    return [nbf.getParticleParameters(i) for i in range(dummy_top.getNumAtoms())]
+    return {
+        atom.element.symbol: nbf.getParticleParameters(atom.index)
+        for atom in dummy_top.atoms()
+    }
 
 def _get_ion_parameters(
     topology: app.Topology,
@@ -125,7 +128,8 @@ def _get_ion_parameters(
     warnings.warn(wmsg)
     logger.warning(wmsg)
 
-    return _get_ff_parameters(forcefield, fallback_smiles)[0]
+    ion_params = _get_ff_parameters(forcefield, fallback_smiles)
+    return next(iter(ion_params.values()))
 
 def _fix_alchemical_water_atom_mapping(
     system_mapping: dict[str, Union[dict[int, int], list[int]]],
@@ -226,10 +230,10 @@ def handle_alchemical_waters(
     ion_charge, ion_sigma, ion_epsilon = _get_ion_parameters(
         topology, system, charge_difference, forcefield,
     )
-    # Get water parameters (O is index 0, H is index 1 and 2)
+    # Get water parameters
     water_params = _get_ff_parameters(forcefield, '[H]O[H]')
-    o_charge, _, _ = water_params[0]
-    h_charge, _, _ = water_params[1]
+    o_charge = water_params['O'][0]
+    h_charge = water_params['H'][0]
 
     # Loop through residues, check if they match the residue index
     # mutate the atom as necessary

From 7bb40cca5cdb4a996011985699aa14cd2b949bd4 Mon Sep 17 00:00:00 2001
From: hannahbaumann <hbaumann@uci.edu>
Date: Thu, 4 Jun 2026 14:10:11 +0200
Subject: [PATCH 9/9] Small fix

---
 .../openmm_rfe/_rfe_utils/topologyhelpers.py  | 20 +++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index b3ff708b2..b410d583b 100644
--- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -77,17 +77,22 @@ def _get_ion_parameters(
     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.
+      The charge difference between state A and state B,
+      calculated as stateA_formal_charge - stateB_formal_charge.
+      Positive values indicate a positive ion is needed,
+      negative values a negative ion.
     forcefield : app.ForceField
       The force field to use for parameterization if no ion is found
       in the topology.
 
     Returns
     -------
-    ion_charge : openmm.unit.Quantity
-    ion_sigma : openmm.unit.Quantity
-    ion_epsilon : openmm.unit.Quantity
+    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
     """
     nbf = [i for i in system.getForces() if isinstance(i, NonbondedForce)][0]
 
@@ -182,7 +187,10 @@ def handle_alchemical_waters(
     system_mapping : dictionary
       A dictionary of system mappings between the stateA and stateB systems
     charge_difference : int
-      The charge difference between state A and state B.
+      The charge difference between state A and state B,
+      calculated as stateA_formal_charge - stateB_formal_charge.
+      Positive values indicate a positive ion is needed,
+      negative values a negative ion.
     forcefield : app.ForceField
       The forcefield to use for ion parameterization.