diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py index 844aa8777..9e4aa3eb9 100644 --- a/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py +++ b/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: diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py index 85d0657a9..a7b9f8b68 100644 --- a/src/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py +++ b/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, diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_units.py b/src/openfe/protocols/openmm_rfe/hybridtop_units.py index 8de7845f2..a039618d8 100644 --- a/src/openfe/protocols/openmm_rfe/hybridtop_units.py +++ b/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 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 9b6b3872f..c37476268 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 @@ -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",