Merge pull request #1053 from OpenFreeEnergy/cli_charge_changes

CLI adaptive settings for charge change transformations

news/cli_adaptive_charge_settings.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* Networks planned using the CLI will now automatically use an extended protocol for transformations involving a net charge change `PR#1053 <https://github.com/OpenFreeEnergy/openfe/pull/1053>`_
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py

@@ -3,6 +3,7 @@
 import abc
 import copy
 from typing import Iterable, Callable, Type, Optional
+import warnings
 
 from gufe import (
     Protocol,
@@ -15,6 +16,7 @@
     SmallMoleculeComponent, ProteinComponent, SolventComponent,
     LigandNetwork,
 )
+from openff.units import unit
 
 
 from .abstract_alchemical_network_planner import (
@@ -33,7 +35,7 @@
     RFEComponentLabels,
 )
 from ...protocols.openmm_rfe.equil_rfe_methods import RelativeHybridTopologyProtocol
-
+from ...utils.ligand_utils import get_alchemical_charge_difference
 
 # TODO: move/or find better structure for protocol_generator combintations!
 PROTOCOL_GENERATOR = {
@@ -315,6 +317,47 @@ def __init__(
             protocol=protocol,
         )
 
+    def _build_transformation(
+        self,
+        ligand_mapping_edge: LigandAtomMapping,
+        stateA: ChemicalSystem,
+        stateB: ChemicalSystem,
+        transformation_protocol: Protocol,
+    ) -> Transformation:
+        """
+        Overwrite the default method to handle net charge change transformations with our default protocol.
+        """
+        transformation_name = self.name + "_" + stateA.name + "_" + stateB.name
+
+        protocol_settings = transformation_protocol.settings.unfrozen_copy()
+        if "vacuum" in transformation_name:
+            protocol_settings.forcefield_settings.nonbonded_method = "nocutoff"
+        elif get_alchemical_charge_difference(ligand_mapping_edge) != 0:
+            wmsg = ("Charge changing transformation between ligands "
+                    f"{ligand_mapping_edge.componentA.name} and {ligand_mapping_edge.componentB.name}. "
+                    "A more expensive protocol with 22 lambda windows, sampled "
+                    "for 20 ns each, will be used here.")
+            warnings.warn(wmsg)
+            # apply the recommended charge change settings taken from the industry benchmarking
+            # <https://github.com/OpenFreeEnergy/IndustryBenchmarks2024/blob/2df362306e2727321d55d16e06919559338c4250/industry_benchmarks/utils/plan_rbfe_network.py#L128-L146>
+            protocol_settings.alchemical_settings.explicit_charge_correction = True
+            protocol_settings.simulation_settings.production_length = 20 * unit.nanosecond
+            protocol_settings.simulation_settings.n_replicas = 22
+            protocol_settings.lambda_settings.lambda_windows = 22
+
+        transformation_protocol = transformation_protocol.__class__(
+            settings=protocol_settings
+        )
+
+        return Transformation(
+            stateA=stateA,
+            stateB=stateB,
+            mapping=ligand_mapping_edge,
+            name=transformation_name,
+            protocol=transformation_protocol,
+        )
+
+
     def __call__(
         self,
         ligands: Iterable[SmallMoleculeComponent],

openfe/utils/ligand_utils.py

@@ -0,0 +1,20 @@
+from gufe import LigandAtomMapping
+
+def get_alchemical_charge_difference(mapping: LigandAtomMapping) -> int:
+    """
+    Return the difference in formal charge between stateA and stateB defined as (formal charge A - formal charge B)
+
+    Parameters
+    ----------
+    mapping: LigandAtomMapping
+        The mapping between the end states A and B.
+
+    Returns
+    -------
+    int:
+        The difference in formal charge between the end states.
+    """
+    from rdkit import Chem
+    charge_a = Chem.rdmolops.GetFormalCharge(mapping.componentA.to_rdkit())
+    charge_b = Chem.rdmolops.GetFormalCharge(mapping.componentB.to_rdkit())
+    return charge_a - charge_b

openfecli/tests/commands/test_plan_rbfe_network.py

@@ -10,6 +10,9 @@
     plan_rbfe_network,
     plan_rbfe_network_main,
 )
+from gufe import AlchemicalNetwork
+from gufe.tokenization import JSON_HANDLER
+import json
 
 
 @pytest.fixture(scope='session')
@@ -149,6 +152,51 @@ def test_plan_rbfe_network_cofactors(eg5_files):
 
         assert result.exit_code == 0
 
+@pytest.fixture
+def cdk8_files():
+    with resources.files("openfe.tests.data") as p:
+        if not (cdk8_dir := p.joinpath("cdk8")).exists():
+            shutil.unpack_archive(cdk8_dir.with_suffix(".zip"), p)
+        pdb_path = str(cdk8_dir.joinpath("cdk8_protein.pdb"))
+        lig_path = str(cdk8_dir.joinpath("cdk8_ligands.sdf"))
+
+        yield pdb_path, lig_path
+
+def test_plan_rbfe_network_charge_changes(cdk8_files):
+    """
+    Make sure the protocol settings are changed and a warning is printed when we plan a network
+    with a net charge change.
+    """
+    runner = CliRunner()
+
+    args = [
+        '-p', cdk8_files[0],
+        '-M', cdk8_files[1],
+    ]
+
+    with runner.isolated_filesystem():
+        with pytest.warns(UserWarning, match="Charge changing transformation between ligands lig_40 and lig_41"):
+            result = runner.invoke(plan_rbfe_network, args)
+
+            assert result.exit_code == 0
+            # load the transformations and check the settings
+            network = AlchemicalNetwork.from_dict(
+                json.load(open("alchemicalNetwork/alchemicalNetwork.json"), cls=JSON_HANDLER.decoder)
+            )
+            for edge in network.edges:
+                settings = edge.protocol.settings
+                # check the charged transform
+                if edge.stateA.components["ligand"].name == "lig_40" and edge.stateB.components["ligand"].name == "lig_41":
+                    assert settings.alchemical_settings.explicit_charge_correction is True
+                    assert settings.simulation_settings.production_length.m == 20.0
+                    assert settings.simulation_settings.n_replicas == 22
+                    assert settings.lambda_settings.lambda_windows == 22
+                else:
+                    assert settings.alchemical_settings.explicit_charge_correction is False
+                    assert settings.simulation_settings.production_length.m == 5.0
+                    assert settings.simulation_settings.n_replicas == 11
+                    assert settings.lambda_settings.lambda_windows == 11
+
 
 @pytest.fixture
 def custom_yaml_settings():