diff --git a/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py b/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py index fa5465c88..526651488 100644 --- a/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py +++ b/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 + # + 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], diff --git a/openfe/tests/data/cdk8.zip b/openfe/tests/data/cdk8.zip new file mode 100644 index 000000000..18298f3c3 Binary files /dev/null and b/openfe/tests/data/cdk8.zip differ diff --git a/openfe/utils/ligand_utils.py b/openfe/utils/ligand_utils.py new file mode 100644 index 000000000..d20415d8e --- /dev/null +++ b/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 diff --git a/openfecli/tests/commands/test_plan_rbfe_network.py b/openfecli/tests/commands/test_plan_rbfe_network.py index 2fc4a9947..6703fc2b7 100644 --- a/openfecli/tests/commands/test_plan_rbfe_network.py +++ 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,44 @@ 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) + print(result.output) + 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 + 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 + @pytest.fixture def custom_yaml_settings():