[WIP] Fixes to net charge correction in the hybrid topology scheme

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

@@ -90,7 +90,8 @@ def __init__(self,
                  softcore_alpha=0.5,
                  softcore_LJ_v2=True,
                  softcore_LJ_v2_alpha=0.85,
-                 interpolate_old_and_new_14s=False):
+                 interpolate_old_and_new_14s=False,
+                 alchemical_water_atoms=None):
         """
         Initialize the Hybrid topology factory.
 
@@ -128,6 +129,10 @@ def __init__(self,
             Whether to turn off interactions for new exceptions (not just
             1,4s) at lambda = 0 and old exceptions at lambda = 1; if False,
             they are present in the nonbonded force.
+        alchemical_water_atoms : set of int, optional
+            Old-system indices of waters being turned into counterions. Their
+            charge change is interpolated on the same electrostatics schedule
+            as the ligand net charge change, keeping every window neutral.
         """
 
         # Assign system positions and force
@@ -146,6 +151,7 @@ def __init__(self,
         # Other options
         self._use_dispersion_correction = use_dispersion_correction
         self._interpolate_14s = interpolate_old_and_new_14s
+        self._alchemical_water_atoms = set(alchemical_water_atoms or set())
 
         # Sofcore options
         self._softcore_alpha = softcore_alpha
@@ -1761,6 +1767,35 @@ def _check_indices(idx1, idx2):
         old_nonbonded_terms = self._old_nonbonded_terms
         new_nonbonded_terms = self._new_nonbonded_terms
 
+        # Counterion (alchemical water) atoms: split their charge change over
+        # the ligand electrostatics schedules so every window stays neutral.
+        alchem_water_fracs = None
+        if self._alchemical_water_atoms:
+            q_uold = sum(
+                (old_nonbonded_terms[hybrid_to_old_map[i]][0]
+                 for i in self._atom_classes['unique_old_atoms']),
+                0.0 * unit.elementary_charge,
+            )
+            q_unew = sum(
+                (new_nonbonded_terms[hybrid_to_new_map[i]][0]
+                 for i in self._atom_classes['unique_new_atoms']),
+                0.0 * unit.elementary_charge,
+            )
+            dq_core = sum(
+                (new_nonbonded_terms[hybrid_to_new_map[i]][0]
+                 - old_nonbonded_terms[hybrid_to_old_map[i]][0]
+                 for i in self._atom_classes['core_atoms']
+                 if hybrid_to_old_map[i] not in self._alchemical_water_atoms),
+                0.0 * unit.elementary_charge,
+            )
+            denom = dq_core - q_uold + q_unew
+            if abs(denom / unit.elementary_charge) > 1e-3:
+                alchem_water_fracs = {
+                    'lambda_electrostatics_core': dq_core / denom,
+                    'lambda_electrostatics_delete': -q_uold / denom,
+                    'lambda_electrostatics_insert': q_unew / denom,
+                }
+
         # Define new global parameters for NonbondedForce
         self._hybrid_system_forces['standard_nonbonded_force'].addGlobalParameter('lambda_electrostatics_core', 0.0)
         self._hybrid_system_forces['standard_nonbonded_force'].addGlobalParameter('lambda_sterics_core', 0.0)
@@ -1855,12 +1890,20 @@ def _check_indices(idx1, idx2):
                 # instead of core_sterics force so that core_sterics_force
                 # could just be softcore.
 
-                # Interpolate between old and new charge with
-                # lambda_electrostatics core make sure to keep sterics off
-                self._hybrid_system_forces['standard_nonbonded_force'].addParticleParameterOffset(
-                    'lambda_electrostatics_core', particle_index,
-                    (charge_new - charge_old), 0, 0
-                )
+                # Interpolate old -> new charge; counterion atoms track the
+                # ligand net-charge schedule, everything else is linear.
+                if (alchem_water_fracs is not None
+                        and old_index in self._alchemical_water_atoms):
+                    for param, frac in alchem_water_fracs.items():
+                        self._hybrid_system_forces['standard_nonbonded_force'].addParticleParameterOffset(
+                            param, particle_index,
+                            (charge_new - charge_old) * frac, 0, 0
+                        )
+                else:
+                    self._hybrid_system_forces['standard_nonbonded_force'].addParticleParameterOffset(
+                        'lambda_electrostatics_core', particle_index,
+                        (charge_new - charge_old), 0, 0
+                    )
 
             # Otherwise, the particle is in the environment
             else:

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

@@ -175,7 +175,7 @@ def handle_alchemical_waters(
     system_mapping: dict,
     charge_difference: int,
     forcefield: app.ForceField,
-):
+) -> set[int]:
     """
     Add alchemical waters from a pre-defined list.
 
@@ -197,6 +197,11 @@ def handle_alchemical_waters(
     forcefield : app.ForceField
       The forcefield to use for ion parameterization.
 
+    Returns
+    -------
+    set[int]
+      Old-system atom indices of the waters converted into ions.
+
     Raises
     ------
     ValueError
@@ -217,7 +222,7 @@ def handle_alchemical_waters(
         raise ValueError(errmsg)
 
     if charge_difference == 0:
-        return None
+        return set()
 
     # get the nonbonded forces
     nbfrcs = [i for i in system.getForces()
@@ -248,12 +253,14 @@ def handle_alchemical_waters(
 
     # Loop through residues, check if they match the residue index
     # mutate the atom as necessary
+    alchemical_water_atoms: set[int] = set()
     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)
                 _fix_alchemical_water_atom_mapping(system_mapping, idx)
+                alchemical_water_atoms.add(system_mapping['new_to_old_atom_map'][idx])
 
                 if charge == o_charge:
                     nbf.setParticleParameters(
@@ -267,6 +274,8 @@ def handle_alchemical_waters(
 
                     nbf.setParticleParameters(idx, 0.0, sigma, epsilon)
 
+    return alchemical_water_atoms
+
 
 def get_alchemical_waters(
     topology: app.Topology,

src/openfe/protocols/openmm_rfe/hybridtop_units.py

@@ -392,7 +392,7 @@ def _handle_net_charge(
         system_mappings: dict[str, dict[int, int]],
         distance_cutoff: Quantity,
         forcefield: openmm.app.ForceField,
-    ) -> None:
+    ) -> set[int]:
         """
         Handle system net charge by adding an alchemical water.
 
@@ -406,10 +406,15 @@ def _handle_net_charge(
         system_mappings : dict[str, dict[int, int]]
         distance_cutoff : Quantity
         forcefield: openmm.app.ForceField
+
+        Returns
+        -------
+        set[int]
+          Old-system atom indices of waters converted into ions (empty if none).
         """
         # Base case, return if no net charge
         if charge_difference == 0:
-            return
+            return set()
 
         # Get the residue ids for waters to turn alchemical
         alchem_water_resids = _rfe_utils.topologyhelpers.get_alchemical_waters(
@@ -420,7 +425,7 @@ def _handle_net_charge(
         )
 
         # In-place modify state B alchemical waters to ions
-        _rfe_utils.topologyhelpers.handle_alchemical_waters(
+        return _rfe_utils.topologyhelpers.handle_alchemical_waters(
             water_resids=alchem_water_resids,
             topology=stateB_topology,
             system=stateB_system,
@@ -545,9 +550,10 @@ 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.
+        alchemical_water_atoms: set[int] = set()
         if settings["alchemical_settings"].explicit_charge_correction:
             forcefield = states_inputs["A"]["generator"].forcefield
-            self._handle_net_charge(
+            alchemical_water_atoms = self._handle_net_charge(
                 stateA_topology=stateA_topology,
                 stateA_positions=stateA_positions,
                 stateB_topology=stateB_topology,
@@ -557,6 +563,7 @@ def _filter_small_mols(smols, state):
                 distance_cutoff=settings["alchemical_settings"].explicit_charge_correction_cutoff,
                 forcefield=forcefield,
             )
+        system_mappings["alchemical_water_atoms"] = alchemical_water_atoms
 
         # Finally get the state B positions
         stateB_positions = _rfe_utils.topologyhelpers.set_and_check_new_positions(
@@ -639,6 +646,7 @@ def _get_alchemical_system(
             softcore_LJ_v2=softcore_LJ_v2,
             softcore_LJ_v2_alpha=alchemical_settings.softcore_alpha,
             interpolate_old_and_new_14s=alchemical_settings.turn_off_core_unique_exceptions,
+            alchemical_water_atoms=system_mappings.get("alchemical_water_atoms", set()),
         )
 
         return hybrid_factory, hybrid_factory.hybrid_system

src/openfe/tests/protocols/openmm_rfe/test_hybrid_top_protocol.py

@@ -2191,11 +2191,11 @@ def test_handle_alchemical_wats(
 def _assert_total_charge(system, atom_classes, chgA, chgB):
     nonbond = [f for f in system.getForces() if isinstance(f, NonbondedForce)]
 
-    offsets = {}
+    offsets: dict[int, float] = {}
     for i in range(nonbond[0].getNumParticleParameterOffsets()):
         offset = nonbond[0].getParticleParameterOffset(i)
         assert len(offset) == 5
-        offsets[offset[1]] = ensure_quantity(offset[2], "openff")
+        offsets[offset[1]] = offsets.get(offset[1], 0.0) + offset[2]
 
     stateA_charges = np.zeros(system.getNumParticles())
     stateB_charges = np.zeros(system.getNumParticles())
@@ -2212,12 +2212,12 @@ def _assert_total_charge(system, atom_classes, chgA, chgB):
         # particle charge (c) is equal to 0
         # offset (c_offset) is equal to molB particle charge
         elif i in atom_classes["unique_new_atoms"]:
-            stateB_charges[i] = offsets[i].m
+            stateB_charges[i] = offsets[i]
         # particle charge (c) is equal to molA particle charge
         # offset (c_offset) is equal to difference between molB and molA
         elif i in atom_classes["core_atoms"]:
             stateA_charges[i] = c.m
-            stateB_charges[i] = c.m + offsets[i].m
+            stateB_charges[i] = c.m + offsets[i]
         # an environment atom
         else:
             assert i in atom_classes["environment_atoms"]
@@ -2228,6 +2228,27 @@ def _assert_total_charge(system, atom_classes, chgA, chgB):
     assert chgB == pytest.approx(np.sum(stateB_charges))
 
 
+def _assert_neutral_all_windows(system):
+    """Total charge of the hybrid system must be ~0 at every lambda window."""
+    from openfe.protocols.openmm_rfe._rfe_utils.lambdaprotocol import LambdaProtocol
+
+    nbf = [f for f in system.getForces() if isinstance(f, NonbondedForce)][0]
+    base = np.array([
+        nbf.getParticleParameters(i)[0].value_in_unit(omm_unit.elementary_charge)
+        for i in range(system.getNumParticles())
+    ])
+    offs = [nbf.getParticleParameterOffset(i)[:3]
+            for i in range(nbf.getNumParticleParameterOffsets())]
+
+    lp = LambdaProtocol(functions="default")
+    for lam in lp.lambda_schedule:
+        vals = {k: fn(lam) for k, fn in lp.functions.items()}
+        charges = base.copy()
+        for name, pidx, qscale in offs:
+            charges[pidx] += qscale * vals[name]
+        assert np.sum(charges) == pytest.approx(0.0, abs=1e-5)
+
+
 def test_dry_run_alchemwater_solvent(benzene_to_benzoic_mapping, solv_settings, tmp_path):
     stateA_system = openfe.ChemicalSystem(
         {
@@ -2258,12 +2279,80 @@ def test_dry_run_alchemwater_solvent(benzene_to_benzoic_mapping, solv_settings,
     results = dag_setup_unit.run(dry=True, scratch_basepath=tmp_path, shared_basepath=tmp_path)
     htf = results["hybrid_factory"]
     _assert_total_charge(htf.hybrid_system, htf._atom_classes, 0, 0)
+    _assert_neutral_all_windows(htf.hybrid_system)
 
     assert len(htf._atom_classes["core_atoms"]) == 14
     assert len(htf._atom_classes["unique_new_atoms"]) == 3
     assert len(htf._atom_classes["unique_old_atoms"]) == 1
 
 
+@pytest.mark.parametrize(
+    "mapping_name",
+    [
+        "benzene_to_benzoic_mapping",  # 0 -> -1: charge appears on state B (anion)
+        "benzoic_to_benzene_mapping",  # -1 -> 0: charge disappears from state A (anion)
+        "benzene_to_aniline_mapping",  # 0 -> +1: charge appears on state B (cation)
+        "aniline_to_benzene_mapping",  # +1 -> 0: charge disappears from state A (cation)
+    ],
+)
+def test_dry_run_alchemwater_solvent_directionality(mapping_name, solv_settings, tmp_path, request):
+    """
+    Regression test for the alchemical-water net-charge correction: the hybrid
+    system must stay exactly neutral at every lambda window regardless of
+    whether the net formal charge sits on state A or state B, and regardless
+    of its sign.
+    """
+    mapping = request.getfixturevalue(mapping_name)
+    assert abs(mapping.get_alchemical_charge_difference()) == 1
+
+    stateA_system = openfe.ChemicalSystem(
+        {"ligand": mapping.componentA, "solvent": openfe.SolventComponent()}
+    )
+    stateB_system = openfe.ChemicalSystem(
+        {"ligand": mapping.componentB, "solvent": openfe.SolventComponent()}
+    )
+    solv_settings.alchemical_settings.explicit_charge_correction = True
+    protocol = openmm_rfe.RelativeHybridTopologyProtocol(settings=solv_settings)
+    dag = protocol.create(stateA=stateA_system, stateB=stateB_system, mapping=mapping)
+    dag_setup_unit = _get_units(dag.protocol_units, HybridTopologySetupUnit)[0]
+    results = dag_setup_unit.run(dry=True, scratch_basepath=tmp_path, shared_basepath=tmp_path)
+    htf = results["hybrid_factory"]
+
+    _assert_neutral_all_windows(htf.hybrid_system)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize(
+    "mapping_name",
+    ["benzene_to_benzoic_mapping", "benzoic_to_benzene_mapping"],
+)
+def test_setup_complex_alchemwater_directionality(
+    mapping_name, solv_settings, tmp_path, request, T4_protein_component,
+):
+    """
+    As test_dry_run_alchemwater_solvent_directionality, but in a complex
+    (protein+solvent) system.
+    """
+    mapping = request.getfixturevalue(mapping_name)
+    solvent = openfe.SolventComponent()
+    stateA_system = openfe.ChemicalSystem(
+        {"ligand": mapping.componentA, "solvent": solvent, "protein": T4_protein_component}
+    )
+    stateB_system = openfe.ChemicalSystem(
+        {"ligand": mapping.componentB, "solvent": solvent, "protein": T4_protein_component}
+    )
+    solv_settings.solvation_settings.solvent_padding = "0.9 nm"
+    solv_settings.solvation_settings.box_shape = "dodecahedron"
+    solv_settings.alchemical_settings.explicit_charge_correction = True
+    protocol = openmm_rfe.RelativeHybridTopologyProtocol(settings=solv_settings)
+    dag = protocol.create(stateA=stateA_system, stateB=stateB_system, mapping=mapping)
+    dag_setup_unit = _get_units(dag.protocol_units, HybridTopologySetupUnit)[0]
+    results = dag_setup_unit.run(dry=True, scratch_basepath=tmp_path, shared_basepath=tmp_path)
+    htf = results["hybrid_factory"]
+
+    _assert_neutral_all_windows(htf.hybrid_system)
+
+
 @pytest.mark.slow
 @pytest.mark.parametrize(
     "mapping_name,chgA,chgB,correction,core_atoms,new_uniq,old_uniq",