diff --git a/news/absolute_settings.rst b/news/absolute_settings.rst new file mode 100644 index 000000000..dc4d715c1 --- /dev/null +++ b/news/absolute_settings.rst @@ -0,0 +1,27 @@ +**Added:** + +* New options have been added to the ``AlchemicalSettings`` + of the ``SepTopProtocol``, ``AbsoluteSolvationProtocol`` + and ``AbsoluteBindingProtocol``. Notably, these options allow users to + control the softcore parameters as well as the use of + long range dispersion corrections. + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py index 9cafdf47b..d28a630ca 100644 --- a/openfe/protocols/openmm_afe/base.py +++ b/openfe/protocols/openmm_afe/base.py @@ -36,7 +36,7 @@ from openff.units import Quantity, unit from openff.units.openmm import ensure_quantity, from_openmm, to_openmm from openmm import app -from openmm import unit as omm_unit +from openmm import unit as ommunit from openmmforcefields.generators import SystemGenerator from openmmtools import multistate from openmmtools.alchemy import ( @@ -52,6 +52,7 @@ ) from openfe.protocols.openmm_afe.equil_afe_settings import ( + AlchemicalSettings, BaseSolvationSettings, IntegratorSettings, MultiStateOutputSettings, @@ -168,10 +169,10 @@ def _pre_equilibrate( self, system: openmm.System, topology: openmm.app.Topology, - positions: omm_unit.Quantity, + positions: ommunit.Quantity, settings: dict[str, SettingsBaseModel], dry: bool, - ) -> tuple[omm_unit.Quantity, omm_unit.Quantity]: + ) -> tuple[ommunit.Quantity, ommunit.Quantity]: """ Run a non-alchemical equilibration to get a stable system. @@ -619,6 +620,7 @@ def _get_alchemical_system( system: openmm.System, comp_resids: dict[Component, npt.NDArray], alchem_comps: dict[str, list[Component]], + alchemical_settings: AlchemicalSettings, ) -> tuple[AbsoluteAlchemicalFactory, openmm.System, list[int]]: """ Get an alchemically modified system and its associated factory @@ -633,6 +635,8 @@ def _get_alchemical_system( A dictionary of residues for each component in the System. alchem_comps : dict[str, list[Component]] A dictionary of alchemical components for each end state. + alchemical_settings : AlchemicalSettings + Settings controlling how the alchemical system is built. Returns ------- @@ -652,9 +656,26 @@ def _get_alchemical_system( alchemical_region = AlchemicalRegion( alchemical_atoms=alchemical_indices, + softcore_alpha=alchemical_settings.softcore_alpha, + annihilate_electrostatics=True, + annihilate_sterics=alchemical_settings.annihilate_sterics, + softcore_a=alchemical_settings.softcore_a, + softcore_b=alchemical_settings.softcore_b, + softcore_c=alchemical_settings.softcore_c, + softcore_beta=0.0, + softcore_d=1.0, + softcore_e=1.0, + softcore_f=2.0, ) - alchemical_factory = AbsoluteAlchemicalFactory() + alchemical_factory = AbsoluteAlchemicalFactory( + consistent_exceptions=False, + switch_width=1.0 * ommunit.angstroms, + alchemical_pme_treatment="exact", + alchemical_rf_treatment="switched", + disable_alchemical_dispersion_correction=alchemical_settings.disable_alchemical_dispersion_correction, + split_alchemical_forces=True, + ) alchemical_system = alchemical_factory.create_alchemical_system(system, alchemical_region) return alchemical_factory, alchemical_system, alchemical_indices @@ -1141,7 +1162,11 @@ def run( # 7. Get alchemical system alchem_factory, alchem_system, alchem_indices = self._get_alchemical_system( - omm_topology, restrained_omm_system, comp_resids, alchem_comps + topology=omm_topology, + system=restrained_omm_system, + comp_resids=comp_resids, + alchem_comps=alchem_comps, + alchemical_settings=settings["alchemical_settings"], ) # 8. Get compound and sampler states diff --git a/openfe/protocols/openmm_afe/equil_afe_settings.py b/openfe/protocols/openmm_afe/equil_afe_settings.py index e652cfa1b..afb74528c 100644 --- a/openfe/protocols/openmm_afe/equil_afe_settings.py +++ b/openfe/protocols/openmm_afe/equil_afe_settings.py @@ -42,9 +42,48 @@ class AlchemicalSettings(SettingsBaseModel): - """Settings for the alchemical protocol + """ + Alchemical settings for Protocols which use the + AbsoluteAlchemicalFactory. + """ + + disable_alchemical_dispersion_correction: bool = False + """ + If True, the long-range dispersion correction will not + be included for the alchemical region, avoiding the need + to recompute the correction. This can improve performance, + at the cost of accuracy. Default is False. + """ + annihilate_sterics: bool = False + """ + If True, sterics (Lennard-Jones) will be annhilated instead + of decoupled. Default is False. + """ + softcore_alpha: float = 0.5 + """ + Alchemical softcore parameter for the Lennard-Jones interactions + (default is 0.5). - Empty place holder for right now. + The generalized softcore potential formalism introduced by + Pham and Shirts, J. Chem. Phys. 135, 034114 (2011), equation 13, + is used here. The ``softcore_a``, ``softcore_b``, and + ``softcore_c`` parameters are used alongside ``softcore_alpha`` + to control how the potential is scaled. + """ + softcore_a: float = 1.0 + """ + Scaling constant ``a`` in + Eq. 13 from Pham and Shirts, J. Chem. Phys. 135, 034114 (2011). + """ + softcore_b: float = 1.0 + """ + Scaling constant ``b`` in + Eq. 13 from Pham and Shirts, J. Chem. Phys. 135, 034114 (2011). + """ + softcore_c: float = 6.0 + """ + Scaling constant ``c`` in + Eq. 13 from Pham and Shirts, J. Chem. Phys. 135, 034114 (2011). """ diff --git a/openfe/protocols/openmm_septop/base.py b/openfe/protocols/openmm_septop/base.py index 18d0209b8..000f0b14e 100644 --- a/openfe/protocols/openmm_septop/base.py +++ b/openfe/protocols/openmm_septop/base.py @@ -45,6 +45,7 @@ ) from openfe.protocols.openmm_afe.equil_afe_settings import ( + AlchemicalSettings, BaseSolvationSettings, IntegratorSettings, MultiStateOutputSettings, @@ -285,6 +286,7 @@ def _get_alchemical_system( system: openmm.System, alchem_indices_A: list[int], alchem_indices_B: list[int], + alchemical_settings: AlchemicalSettings, ) -> tuple[AbsoluteAlchemicalFactory, openmm.System]: """ Get an alchemically modified system and its associated factory @@ -293,12 +295,14 @@ def _get_alchemical_system( ---------- system : openmm.System System to alchemically modify. - alchem_indices_A: list[int] + alchem_indices_A : list[int] A list of atom indices for the alchemically modified ligand A in the system. - alchem_indices_B: list[int] + alchem_indices_B : list[int] A list of atom indices for the alchemically modified ligand B in the system. + alchemical_settings : AlchemicalSettings + Settings controlling how the alchemical system will be built. Returns ------- @@ -313,13 +317,42 @@ def _get_alchemical_system( switch_width=1.0 * unit.angstroms, alchemical_pme_treatment="exact", alchemical_rf_treatment="switched", - disable_alchemical_dispersion_correction=False, + disable_alchemical_dispersion_correction=alchemical_settings.disable_alchemical_dispersion_correction, split_alchemical_forces=True, ) + # Alchemical Region for ligand A - alchemical_region_A = AlchemicalRegion(alchemical_atoms=alchem_indices_A, name="A") + alchemical_region_A = AlchemicalRegion( + alchemical_atoms=alchem_indices_A, + name="A", + softcore_alpha=alchemical_settings.softcore_alpha, + annihilate_electrostatics=True, + annihilate_sterics=alchemical_settings.annihilate_sterics, + softcore_a=alchemical_settings.softcore_a, + softcore_b=alchemical_settings.softcore_b, + softcore_c=alchemical_settings.softcore_c, + softcore_beta=0.0, + softcore_d=1.0, + softcore_e=1.0, + softcore_f=2.0, + ) + # Alchemical Region for ligand B - alchemical_region_B = AlchemicalRegion(alchemical_atoms=alchem_indices_B, name="B") + alchemical_region_B = AlchemicalRegion( + alchemical_atoms=alchem_indices_B, + name="B", + softcore_alpha=alchemical_settings.softcore_alpha, + annihilate_electrostatics=True, + annihilate_sterics=alchemical_settings.annihilate_sterics, + softcore_a=alchemical_settings.softcore_a, + softcore_b=alchemical_settings.softcore_b, + softcore_c=alchemical_settings.softcore_c, + softcore_beta=0.0, + softcore_d=1.0, + softcore_e=1.0, + softcore_f=2.0, + ) + alchemical_system = alchemical_factory.create_alchemical_system( system, [alchemical_region_A, alchemical_region_B] ) diff --git a/openfe/protocols/openmm_septop/equil_septop_method.py b/openfe/protocols/openmm_septop/equil_septop_method.py index cefabf359..87f7aa7a0 100644 --- a/openfe/protocols/openmm_septop/equil_septop_method.py +++ b/openfe/protocols/openmm_septop/equil_septop_method.py @@ -1971,6 +1971,7 @@ def run( omm_system_AB, atom_indices_AB_A, atom_indices_AB_B, + settings["alchemical_settings"], ) # 10. Apply Restraints @@ -2254,6 +2255,7 @@ def run( omm_system_AB, atom_indices_AB_A, atom_indices_AB_B, + settings["alchemical_settings"], ) # 8. Apply Restraints diff --git a/openfe/protocols/openmm_septop/equil_septop_settings.py b/openfe/protocols/openmm_septop/equil_septop_settings.py index 6e6522acb..e1898c5fa 100644 --- a/openfe/protocols/openmm_septop/equil_septop_settings.py +++ b/openfe/protocols/openmm_septop/equil_septop_settings.py @@ -23,6 +23,7 @@ from openff.units import unit as offunit from pydantic import field_validator +from openfe.protocols.openmm_afe.equil_afe_settings import AlchemicalSettings from openfe.protocols.openmm_utils.omm_settings import ( IntegratorSettings, MDOutputSettings, @@ -36,13 +37,6 @@ from openfe.protocols.restraint_utils.settings import BaseRestraintSettings -class AlchemicalSettings(SettingsBaseModel): - """Settings for the alchemical protocol - - Empty place holder for right now. - """ - - class LambdaSettings(SettingsBaseModel): """Lambda schedule settings. diff --git a/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py b/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py index 33cd36223..206ff98c7 100644 --- a/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py +++ b/openfe/tests/protocols/openmm_ahfe/test_ahfe_protocol.py @@ -13,7 +13,15 @@ from numpy.testing import assert_allclose from openff.units import unit as offunit from openff.units.openmm import ensure_quantity, from_openmm -from openmm import CustomNonbondedForce, NonbondedForce +from openmm import ( + CustomBondForce, + CustomNonbondedForce, + HarmonicAngleForce, + HarmonicBondForce, + MonteCarloBarostat, + NonbondedForce, + PeriodicTorsionForce, +) from openmmtools.multistate.multistatesampler import MultiStateSampler import openfe @@ -86,6 +94,54 @@ def test_create_independent_repeat_ids(benzene_system): assert len(repeat_ids) == 12 +def _assert_num_forces(system, forcetype, number): + """ + Helper method to check the number of forces of a given + type in a system. + """ + forces = [f for f in system.getForces() if isinstance(f, forcetype)] + assert len(forces) == number + + +def _verify_alchemical_sterics_force_parameters( + force, + long_range=True, + alpha=0.5, + beta=0, + a=1.0, + b=1.0, + c=6.0, + d=1.0, + e=1.0, + f=2.0, +): + assert force.getUseLongRangeCorrection() is long_range + + if force.getNumGlobalParameters() == 8: + shift = 0 + else: + shift = 1 + + # Check the softcore parameters for the sterics forces + assert force.getGlobalParameterName(0 + shift) == "softcore_alpha" + assert force.getGlobalParameterName(1 + shift) == "softcore_beta" + assert force.getGlobalParameterName(2 + shift) == "softcore_a" + assert force.getGlobalParameterName(3 + shift) == "softcore_b" + assert force.getGlobalParameterName(4 + shift) == "softcore_c" + assert force.getGlobalParameterName(5 + shift) == "softcore_d" + assert force.getGlobalParameterName(6 + shift) == "softcore_e" + assert force.getGlobalParameterName(7 + shift) == "softcore_f" + + assert force.getGlobalParameterDefaultValue(0 + shift) == alpha + assert force.getGlobalParameterDefaultValue(1 + shift) == beta + assert force.getGlobalParameterDefaultValue(2 + shift) == a + assert force.getGlobalParameterDefaultValue(3 + shift) == b + assert force.getGlobalParameterDefaultValue(4 + shift) == c + assert force.getGlobalParameterDefaultValue(5 + shift) == d + assert force.getGlobalParameterDefaultValue(6 + shift) == e + assert force.getGlobalParameterDefaultValue(7 + shift) == f + + @pytest.mark.parametrize("method", ["repex", "sams", "independent", "InDePeNdENT"]) def test_dry_run_vac_benzene(benzene_system, method, protocol_dry_settings, tmpdir): protocol_dry_settings.vacuum_simulation_settings.sampler_method = method @@ -114,9 +170,110 @@ def test_dry_run_vac_benzene(benzene_system, method, protocol_dry_settings, tmpd assert len(sol_unit) == 1 with tmpdir.as_cwd(): - vac_sampler = vac_unit[0].run(dry=True)["debug"]["sampler"] + debug = vac_unit[0].run(dry=True)["debug"] + vac_sampler = debug["sampler"] assert not vac_sampler.is_periodic + # standard system + system = debug["system"] + assert system.getNumParticles() == 12 + assert len(system.getForces()) == 4 + _assert_num_forces(system, NonbondedForce, 1) + _assert_num_forces(system, HarmonicBondForce, 1) + _assert_num_forces(system, HarmonicAngleForce, 1) + _assert_num_forces(system, PeriodicTorsionForce, 1) + + # alchemical system + alchem_system = debug["alchem_system"] + assert alchem_system.getNumParticles() == 12 + assert len(alchem_system.getForces()) == 12 + _assert_num_forces(alchem_system, NonbondedForce, 1) + _assert_num_forces(alchem_system, CustomNonbondedForce, 4) + _assert_num_forces(alchem_system, CustomBondForce, 4) + _assert_num_forces(alchem_system, HarmonicBondForce, 1) + _assert_num_forces(alchem_system, HarmonicAngleForce, 1) + _assert_num_forces(alchem_system, PeriodicTorsionForce, 1) + + # Check some force contents + stericsf = [ + f + for f in alchem_system.getForces() + if isinstance(f, CustomNonbondedForce) and "U_sterics" in f.getEnergyFunction() + ] + + for force in stericsf: + _verify_alchemical_sterics_force_parameters(force) + + +@pytest.mark.parametrize( + "alpha, a, b, c, correction", + [ + [0.2, 2, 2, 1, True], + [0.35, 2.2, 1.5, 0, False], + ], +) +def test_alchemical_settings_dry_run_vacuum( + alpha, a, b, c, correction, benzene_system, protocol_dry_settings, tmpdir +): + """ + Test non default alchemical settings + """ + protocol_dry_settings.alchemical_settings.softcore_alpha = alpha + protocol_dry_settings.alchemical_settings.softcore_a = a + protocol_dry_settings.alchemical_settings.softcore_b = b + protocol_dry_settings.alchemical_settings.softcore_c = c + protocol_dry_settings.alchemical_settings.disable_alchemical_dispersion_correction = correction + protocol = openmm_afe.AbsoluteSolvationProtocol(settings=protocol_dry_settings) + + stateA = benzene_system + + stateB = ChemicalSystem({"solvent": SolventComponent()}) + + # Create DAG from protocol, get the vacuum and solvent units + # and eventually dry run the first vacuum unit + dag = protocol.create( + stateA=stateA, + stateB=stateB, + mapping=None, + ) + prot_units = list(dag.protocol_units) + + assert len(prot_units) == 2 + + vac_unit = [u for u in prot_units if isinstance(u, AbsoluteSolvationVacuumUnit)] + sol_unit = [u for u in prot_units if isinstance(u, AbsoluteSolvationSolventUnit)] + + assert len(vac_unit) == 1 + assert len(sol_unit) == 1 + + with tmpdir.as_cwd(): + debug = vac_unit[0].run(dry=True)["debug"] + + alchem_system = debug["alchem_system"] + _assert_num_forces(alchem_system, NonbondedForce, 1) + _assert_num_forces(alchem_system, CustomNonbondedForce, 4) + _assert_num_forces(alchem_system, CustomBondForce, 4) + _assert_num_forces(alchem_system, HarmonicBondForce, 1) + _assert_num_forces(alchem_system, HarmonicAngleForce, 1) + _assert_num_forces(alchem_system, PeriodicTorsionForce, 1) + + # Check some force contents + stericsf = [ + f + for f in alchem_system.getForces() + if isinstance(f, CustomNonbondedForce) and "U_sterics" in f.getEnergyFunction() + ] + + for force in stericsf: + _verify_alchemical_sterics_force_parameters( + force, + long_range=not correction, + alpha=alpha, + a=a, + b=b, + c=c, + ) + def test_confgen_fail_AFE(benzene_system, protocol_dry_settings, tmpdir): # check system parametrisation works even if confgen fails diff --git a/openfe/tests/protocols/openmm_septop/test_septop_protocol.py b/openfe/tests/protocols/openmm_septop/test_septop_protocol.py index 25b4f2e85..83090368f 100644 --- a/openfe/tests/protocols/openmm_septop/test_septop_protocol.py +++ b/openfe/tests/protocols/openmm_septop/test_septop_protocol.py @@ -17,7 +17,16 @@ from numpy.testing import assert_allclose from openff.units import unit as offunit from openff.units.openmm import ensure_quantity, from_openmm -from openmm import CustomNonbondedForce, MonteCarloBarostat, NonbondedForce +from openmm import ( + CustomBondForce, + CustomCompoundBondForce, + CustomNonbondedForce, + HarmonicAngleForce, + HarmonicBondForce, + MonteCarloBarostat, + NonbondedForce, + PeriodicTorsionForce, +) from openmmtools.alchemy import AbsoluteAlchemicalFactory, AlchemicalRegion from openmmtools.multistate.multistatesampler import MultiStateSampler @@ -38,6 +47,10 @@ from openfe.protocols.openmm_utils import system_validation from openfe.protocols.restraint_utils.geometry.boresch import BoreschRestraintGeometry from openfe.tests.protocols.conftest import compute_energy +from openfe.tests.protocols.openmm_ahfe.test_ahfe_protocol import ( + _assert_num_forces, + _verify_alchemical_sterics_force_parameters, +) E_CHARGE = 1.602176634e-19 * openmm.unit.coulomb EPSILON0 = ( @@ -717,6 +730,22 @@ def test_dry_run_benzene_toluene(benzene_toluene_dag, tmpdir): pdb = md.load_pdb("alchemical_system.pdb") assert pdb.n_atoms == 31 + # Test the solvent system + alchem_system = deserialize(serialized_system) + assert len(alchem_system.getForces()) == 14 + _assert_num_forces(alchem_system, NonbondedForce, 1) + _assert_num_forces(alchem_system, CustomNonbondedForce, 4) + _assert_num_forces(alchem_system, CustomBondForce, 4) + _assert_num_forces(alchem_system, HarmonicBondForce, 2) + _assert_num_forces(alchem_system, HarmonicAngleForce, 1) + _assert_num_forces(alchem_system, PeriodicTorsionForce, 1) + _assert_num_forces(alchem_system, MonteCarloBarostat, 1) + + # Check steric forces + for f in alchem_system.getForces(): + if isinstance(f, CustomNonbondedForce) and "U_sterics" in f.getEnergyFunction(): + _verify_alchemical_sterics_force_parameters(f) + complex_setup_output = complex_setup_unit[0].run(dry=True) serialized_topology = complex_setup_output["topology"] serialized_system = complex_setup_output["system"] @@ -732,6 +761,23 @@ def test_dry_run_benzene_toluene(benzene_toluene_dag, tmpdir): pdb = md.load_pdb("alchemical_system.pdb") assert pdb.n_atoms == 2687 + # Test the complex system + alchem_system = deserialize(serialized_system) + assert len(alchem_system.getForces()) == 15 + _assert_num_forces(alchem_system, NonbondedForce, 1) + _assert_num_forces(alchem_system, CustomNonbondedForce, 4) + _assert_num_forces(alchem_system, CustomBondForce, 4) + _assert_num_forces(alchem_system, HarmonicBondForce, 1) + _assert_num_forces(alchem_system, HarmonicAngleForce, 1) + _assert_num_forces(alchem_system, PeriodicTorsionForce, 1) + _assert_num_forces(alchem_system, CustomCompoundBondForce, 2) + _assert_num_forces(alchem_system, MonteCarloBarostat, 1) + + # Check steric forces + for f in alchem_system.getForces(): + if isinstance(f, CustomNonbondedForce) and "U_sterics" in f.getEnergyFunction(): + _verify_alchemical_sterics_force_parameters(f) + @pytest.mark.parametrize("method", ["repex", "sams", "independent"]) def test_dry_run_methods(