diff --git a/news/cmap_support.rst b/news/cmap_support.rst new file mode 100644 index 000000000..bd7220025 --- /dev/null +++ b/news/cmap_support.rst @@ -0,0 +1,24 @@ +**Added:** + +* The ``HybridTopologyFactory`` supports building hybrid OpenMM systems which contain ``CMAPTorsionForces`` on non-alchemical atoms. This should allow for simulations using + Amber ff19SB. See `PR #1695 `_ + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* diff --git a/openfe/protocols/openmm_rfe/_rfe_utils/relative.py b/openfe/protocols/openmm_rfe/_rfe_utils/relative.py index ac6e46bb9..096199a68 100644 --- a/openfe/protocols/openmm_rfe/_rfe_utils/relative.py +++ b/openfe/protocols/openmm_rfe/_rfe_utils/relative.py @@ -215,6 +215,9 @@ def __init__(self, self._handle_periodic_torsion_force() + # add cmap terms if possible + self._handle_cmap_torsion_force() + if has_nonbonded_force: self._handle_nonbonded() if not (len(self._old_system_exceptions.keys()) == 0 and @@ -229,6 +232,155 @@ def __init__(self, self._omm_hybrid_topology = self._create_hybrid_topology() logger.info("Hybrid system created") + @staticmethod + def _verify_cmap_compatibility( + cmap_old: openmm.CMAPTorsionForce, + cmap_new: openmm.CMAPTorsionForce, + ) -> tuple[ + int, + int, + int, + int, + ]: + """ + Verify CMAPTorsionForce compatibility between two systems. + + Parameters + ---------- + cmap_old : openmm.CMAPTorsionForce + CMAPTorsionForce from the old system + cmap_new : openmm.CMAPTorsionForce + CMAPTorsionForce from the new system + + Returns + ------- + tuple + (old_num_maps, new_num_maps, old_num_torsions, new_num_torsions) + four integers describing the number of maps and + torsions in each force. + + Raises + ------ + RuntimeError + If only one of the forces is present, or if the number of maps or the + number of torsions differs between the two forces. + """ + logger.info("CMAPTorsionForce found checking compatibility") + + # some quick checks on compatibility like the number of maps and total number of terms + old_num_maps = cmap_old.getNumMaps() + new_num_maps = cmap_new.getNumMaps() + if old_num_maps != new_num_maps: + raise RuntimeError( + f"Incompatible CMAPTorsionForce between end states expected to have same number of maps, " + f"found old: {old_num_maps} and new: {new_num_maps}") + + old_num_torsions = cmap_old.getNumTorsions() + new_num_torsions = cmap_new.getNumTorsions() + if old_num_torsions != new_num_torsions: + raise RuntimeError( + f"Incompatible CMAPTorsionForce between end states expected to have same number of torsions, " + f"found old: {old_num_torsions} and new: {new_num_torsions}") + + return old_num_maps, new_num_maps, old_num_torsions, new_num_torsions + + def _handle_cmap_torsion_force(self): + """ + This method does the following in order: + - Some basic checks that the CMAPTorsionForce exists in both old/new systems. + - Adds the CMAPTorsionForce from the old system + - Checks that the new system CMAPTorsionForce terms are equal to the old system's (we do not allow for alchemically changing CMAP terms). + """ + cmap_old = self._old_system_forces.get("CMAPTorsionForce", None) + cmap_new = self._new_system_forces.get("CMAPTorsionForce", None) + + # if only one has cmap raise an error + if (cmap_new is None) ^ (cmap_old is None): + raise RuntimeError(f"Inconsistent CMAPTorsionForce between end states expected to be present in both" + f"but found in old: {bool(cmap_old)} and new: {bool(cmap_new)}") + + if cmap_new == cmap_old is None: + logger.info("No CMAPTorsionForce found. Skipping adding force.") + return + + # verify compatibility and extract numbers of maps and torsions + ( + old_num_maps, + new_num_maps, + old_num_torsions, + new_num_torsions + ) = self._verify_cmap_compatibility( + cmap_old, cmap_new + ) + + logger.info("Adding CMAPTorsionForce to hybrid system") + # start looping through the old system terms and add them to the hybrid system + # track the terms we add so we can cross compare with the new system and also make sure we don't hit + # an index in the alchemical region + hybrid_cmap_force = openmm.CMAPTorsionForce() + self._hybrid_system.addForce(hybrid_cmap_force) + self._hybrid_system_forces["cmap_torsion_force"] = hybrid_cmap_force + + old_system_maps = {} + old_system_terms = {} + + logger.info("Adding CMAP forces") + # add all the old maps + for i in range(old_num_maps): + size, energy = cmap_old.getMapParameters(i) + old_system_maps[i] = (size, energy) + # also add the map to the hybrid system + hybrid_cmap_force.addMap(size, energy) + + logger.info("Adding CMAP force terms") + # now add the terms we need to map from the old to the new index + old_to_hybrid_index = self._old_to_hybrid_map + new_to_hybrid_index = self._new_to_hybrid_map + for i in range(old_num_torsions): + # get the parameters for the torsion using the same notation as OpenMM + map_index, a1, a2, a3, a4, b1, b2, b3, b4 = cmap_old.getTorsionParameters(i) + atom_ids = [a1, a2, a3, a4, b1, b2, b3, b4] + # map to hybrid indices + hybrid_atom_ids = [old_to_hybrid_index[a_id] for a_id in atom_ids] + # add to the hybrid system using the hybrid index + hybrid_cmap_force.addTorsion(map_index, *hybrid_atom_ids) + # track the atoms we add in the hybrid system to cross compare with new system + old_system_terms[tuple(hybrid_atom_ids)] = map_index + + # gather all alchemical atoms, use a copy so we don't change the groups + alchemical_atoms = self._atom_classes["core_atoms"].copy() + alchemical_atoms.update(self._atom_classes["unique_old_atoms"], self._atom_classes["unique_new_atoms"]) + # check if any of the atoms added are in the alchemical region + old_added_atoms = {atom_id for atoms in old_system_terms.keys() for atom_id in atoms} + if overlap_atoms := alchemical_atoms.intersection(old_added_atoms): + raise RuntimeError( + f"Incompatible CMAPTorsionForce term found in alchemical region for old system atoms {overlap_atoms}") + + # now loop over the new system force and check the terms are compatible + # we expect to add no new terms + for i in range(new_num_maps): + size, energy = cmap_new.getMapParameters(i) + if (size, energy) != old_system_maps[i]: + raise RuntimeError(f"Incompatible CMAPTorsionForce map parameters found between end states for map {i} " + f"expected {old_system_maps[i]} found {(size, energy)}") + + for i in range(new_num_torsions): + map_index, a1, a2, a3, a4, b1, b2, b3, b4 = cmap_new.getTorsionParameters(i) + atom_ids = [a1, a2, a3, a4, b1, b2, b3, b4] + # map to hybrid indices + hybrid_atom_ids = [new_to_hybrid_index[a_id] for a_id in atom_ids] + # check its in the old system terms + if tuple(hybrid_atom_ids) not in old_system_terms.keys(): + raise RuntimeError( + f"Incompatible CMAPTorsionForce term found between end states for atoms {hybrid_atom_ids} " + f"not found in old system terms.") + # check the map index is the same + if map_index != old_system_terms[tuple(hybrid_atom_ids)]: + raise RuntimeError( + f"Incompatible CMAPTorsionForce map index found between end states for atoms {hybrid_atom_ids} " + f"expected {old_system_terms[tuple(hybrid_atom_ids)]} found {map_index}") + logger.info("CMAPTorsionForce added to the hybrid system") + @staticmethod def _check_bounds(value, varname, minmax=(0, 1)): """ @@ -320,7 +472,7 @@ def _check_unknown_forces(forces, system_name): # TODO: double check that CMMotionRemover is ok being here known_forces = {'HarmonicBondForce', 'HarmonicAngleForce', 'PeriodicTorsionForce', 'NonbondedForce', - 'MonteCarloBarostat', 'CMMotionRemover'} + 'MonteCarloBarostat', 'CMMotionRemover', 'CMAPTorsionForce'} force_names = forces.keys() unknown_forces = set(force_names) - set(known_forces) diff --git a/openfe/tests/conftest.py b/openfe/tests/conftest.py index d0012b822..5e304b695 100644 --- a/openfe/tests/conftest.py +++ b/openfe/tests/conftest.py @@ -1,5 +1,6 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/openfe +import gzip import os import pathlib import urllib.error @@ -9,15 +10,21 @@ import gufe import mdtraj import numpy as np +import openmm import pandas as pd import pytest -from gufe import AtomMapper, LigandAtomMapping, SmallMoleculeComponent -from openff.units import unit +from gufe import AtomMapper, LigandAtomMapping, ProteinComponent, SmallMoleculeComponent +from openff.toolkit import ForceField +from openff.units import unit as offunit +from openmm import unit as ommunit from rdkit import Chem from rdkit.Chem import AllChem import openfe +from openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol +from openfe.protocols.openmm_rfe._rfe_utils.relative import HybridTopologyFactory from openfe.protocols.openmm_septop.utils import deserialize +from openfe.tests.protocols.openmm_rfe.helpers import make_htf class SlowTests: @@ -343,19 +350,19 @@ def am1bcc_ref_charges(): "ambertools":[ 0.146957, -0.918943, 0.025557, 0.025557, 0.025557, 0.347657, 0.347657 - ] * unit.elementary_charge, + ] * offunit.elementary_charge, "openeye": [ 0.14713, -0.92016, 0.02595, 0.02595, 0.02595, 0.34759, 0.34759 - ] * unit.elementary_charge, + ] * offunit.elementary_charge, "nagl": [ 0.170413, -0.930417, 0.021593, 0.021593, 0.021593, 0.347612, 0.347612 - ] * unit.elementary_charge, + ] * offunit.elementary_charge, "espaloma": [ 0.017702, -0.966793, 0.063076, 0.063076, 0.063076, 0.379931, 0.379931 - ] * unit.elementary_charge, + ] * offunit.elementary_charge, } # fmt: skip return ref_chgs @@ -373,3 +380,140 @@ def am1bcc_ref_charges(): HAS_ESPALOMA = True except ModuleNotFoundError: HAS_ESPALOMA = False + + +@pytest.fixture(scope="module") +def chlorobenzene(): + """Load chlorobenzene with partial charges from sdf file.""" + with resources.as_file(resources.files("openfe.tests.data.htf")) as f: + yield SmallMoleculeComponent.from_sdf_file(f / "t4_lysozyme_data" / "chlorobenzene.sdf") + + +@pytest.fixture(scope="module") +def fluorobenzene(): + """Load fluorobenzene with partial charges from sdf file.""" + with resources.as_file(resources.files("openfe.tests.data.htf")) as f: + yield SmallMoleculeComponent.from_sdf_file(f / "t4_lysozyme_data" / "fluorobenzene.sdf") + + +@pytest.fixture(scope="module") +def chlorobenzene_to_fluorobenzene_mapping(chlorobenzene, fluorobenzene): + """Return a mapping from chlorobenzene to fluorobenzene.""" + return LigandAtomMapping( + componentA=chlorobenzene, + componentB=fluorobenzene, + componentA_to_componentB={ + # perfect one-to-one mapping + 0: 0, + 1: 1, + 2: 2, + 3: 3, + 4: 4, + 5: 5, + 6: 6, + 7: 7, + 8: 8, + 9: 9, + 10: 10, + 11: 11, + }, + ) + + +@pytest.fixture(scope="module") +def t4_lysozyme_solvated(): + """Load the T4 lysozyme L99A structure and solvent from the pdb file.""" + with resources.as_file(resources.files("openfe.tests.data.htf")) as f: + with gzip.open(f / "t4_lysozyme_data" / "t4_lysozyme_solvated.pdb.gz", "rb") as gzf: + yield ProteinComponent.from_pdb_file(gzf) + + +def apply_box_vectors_and_fix_nb_force( + hybrid_topology_factory: HybridTopologyFactory, force_field: ForceField +): + """ + Edit the systems in the hybrid topology factory to have the correct box vectors and nonbonded force settings for the T4 lysozyme system. + """ + hybrid_system = hybrid_topology_factory.hybrid_system + # as we use a pre-solvated system, we need to correct the nonbonded methods and cutoffs and set the box vectors + box_vectors = [ + openmm.vec3.Vec3(x=6.90789161545809, y=0.0, z=0.0) * ommunit.nanometer, + openmm.vec3.Vec3(x=0.0, y=6.90789161545809, z=0.0) * ommunit.nanometer, + openmm.vec3.Vec3(x=3.453945807729045, y=3.453945807729045, z=4.88461700499211) + * ommunit.nanometer, + ] + hybrid_system.setDefaultPeriodicBoxVectors(*box_vectors) + for force in hybrid_system.getForces(): + if isinstance(force, openmm.NonbondedForce): + force.setNonbondedMethod(openmm.NonbondedForce.PME) + force.setCutoffDistance( + force_field.get_parameter_handler("Electrostatics").cutoff.m_as(offunit.nanometer) + * ommunit.nanometer + ) + force.setUseDispersionCorrection(False) + force.setUseSwitchingFunction(False) + elif isinstance(force, openmm.CustomNonbondedForce): + force.setCutoffDistance( + force_field.get_parameter_handler("Electrostatics").cutoff.m_as(offunit.nanometer) + * ommunit.nanometer + ) + force.setNonbondedMethod(force.CutoffPeriodic) + force.setUseLongRangeCorrection(False) + force.setUseSwitchingFunction(False) + + # make sure both end state systems have the same cutoff method and distance + for end_state in [hybrid_topology_factory._old_system, hybrid_topology_factory._new_system]: + end_state.setDefaultPeriodicBoxVectors(*box_vectors) + for force in end_state.getForces(): + if isinstance(force, openmm.NonbondedForce): + force.setNonbondedMethod(openmm.NonbondedForce.PME) + force.setCutoffDistance( + force_field.get_parameter_handler("Electrostatics").cutoff.m_as( + offunit.nanometer + ) + * ommunit.nanometer + ) + force.setUseDispersionCorrection(False) + force.setUseSwitchingFunction(False) + + +@pytest.fixture(scope="module") +def htf_cmap_chlorobenzene_to_fluorobenzene( + chlorobenzene_to_fluorobenzene_mapping, t4_lysozyme_solvated +): + """Generate the htf for chlorobenzene to fluorobenzene with a CMAP term.""" + settings = RelativeHybridTopologyProtocol.default_settings() + # make sure we interpolate the 1-4 exceptions involving dummy atoms if present + settings.alchemical_settings.turn_off_core_unique_exceptions = True + small_ff = settings.forcefield_settings.small_molecule_forcefield + if ".offxml" not in small_ff: + small_ff += ".offxml" + ff = ForceField(small_ff) + # update the default force fields to include a force field with CMAP terms + settings.forcefield_settings.forcefields = [ + "amber/protein.ff19SB.xml", # cmap amber ff + "amber/tip3p_standard.xml", # TIP3P and recommended monovalent ion parameters + "amber/tip3p_HFE_multivalent.xml", # for divalent ions + "amber/phosaa19SB.xml", # Handles THE TPO + ] + htf = make_htf( + mapping=chlorobenzene_to_fluorobenzene_mapping, + protein=t4_lysozyme_solvated, + settings=settings, + ) + # apply box vectors and fix nonbonded force settings so we can use PME + apply_box_vectors_and_fix_nb_force(hybrid_topology_factory=htf, force_field=ff) + hybrid_system = htf.hybrid_system + forces = {force.getName(): force for force in hybrid_system.getForces()} + + return { + "htf": htf, + "hybrid_system": hybrid_system, + "forces": forces, + "mapping": chlorobenzene_to_fluorobenzene_mapping, + "chlorobenzene": chlorobenzene_to_fluorobenzene_mapping.componentA, + "fluorobenzene": chlorobenzene_to_fluorobenzene_mapping.componentB, + "electrostatic_scale": ff.get_parameter_handler("Electrostatics").scale14, + "vdW_scale": ff.get_parameter_handler("vdW").scale14, + "force_field": ff, + } diff --git a/openfe/tests/data/htf/__init__.py b/openfe/tests/data/htf/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/openfe/tests/data/htf/t4_lysozyme_data/chlorobenzene.sdf b/openfe/tests/data/htf/t4_lysozyme_data/chlorobenzene.sdf new file mode 100644 index 000000000..5da79e4a3 --- /dev/null +++ b/openfe/tests/data/htf/t4_lysozyme_data/chlorobenzene.sdf @@ -0,0 +1,34 @@ +chlorobenzene + RDKit 3D + + 12 12 0 0 0 0 0 0 0 0999 V2000 + -6.3030 6.9214 0.5047 Cl 0 0 0 0 0 0 0 0 0 0 0 0 + -7.2858 7.1609 -0.9477 C 0 0 0 0 0 0 0 0 0 0 0 0 + -7.8665 6.0569 -1.5923 C 0 0 0 0 0 0 0 0 0 0 0 0 + -8.6603 6.2343 -2.7247 C 0 0 0 0 0 0 0 0 0 0 0 0 + -8.8944 7.5191 -3.2149 C 0 0 0 0 0 0 0 0 0 0 0 0 + -8.3035 8.6211 -2.6006 C 0 0 0 0 0 0 0 0 0 0 0 0 + -7.5003 8.4444 -1.4758 C 0 0 0 0 0 0 0 0 0 0 0 0 + -7.7098 5.0628 -1.2014 H 0 0 0 0 0 0 0 0 0 0 0 0 + -9.0935 5.3789 -3.2180 H 0 0 0 0 0 0 0 0 0 0 0 0 + -9.5294 7.6705 -4.0709 H 0 0 0 0 0 0 0 0 0 0 0 0 + -8.4705 9.6119 -2.9925 H 0 0 0 0 0 0 0 0 0 0 0 0 + -7.0579 9.3041 -0.9946 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 1 0 + 2 3 2 0 + 2 7 1 0 + 3 4 1 0 + 3 8 1 0 + 4 5 2 0 + 4 9 1 0 + 5 6 1 0 + 5 10 1 0 + 6 7 2 0 + 6 11 1 0 + 7 12 1 0 +M END +> (1) +-0.097566666666666677 0.017233333333333326 -0.12416666666666668 -0.12216666666666667 -0.12916666666666668 -0.12216666666666667 -0.12416666666666668 0.14683333333333332 0.13683333333333333 +0.13483333333333333 0.13683333333333333 0.14683333333333332 + +$$$$ diff --git a/openfe/tests/data/htf/t4_lysozyme_data/fluorobenzene.sdf b/openfe/tests/data/htf/t4_lysozyme_data/fluorobenzene.sdf new file mode 100644 index 000000000..0f3fcd273 --- /dev/null +++ b/openfe/tests/data/htf/t4_lysozyme_data/fluorobenzene.sdf @@ -0,0 +1,34 @@ +fluorobenzene + RDKit 3D + + 12 12 0 0 0 0 0 0 0 0999 V2000 + -6.5362 6.9783 0.1601 F 0 0 0 0 0 0 0 0 0 0 0 0 + -7.2858 7.1609 -0.9477 C 0 0 0 0 0 0 0 0 0 0 0 0 + -7.8665 6.0569 -1.5923 C 0 0 0 0 0 0 0 0 0 0 0 0 + -8.6603 6.2343 -2.7247 C 0 0 0 0 0 0 0 0 0 0 0 0 + -8.8944 7.5191 -3.2149 C 0 0 0 0 0 0 0 0 0 0 0 0 + -8.3035 8.6211 -2.6006 C 0 0 0 0 0 0 0 0 0 0 0 0 + -7.5003 8.4444 -1.4758 C 0 0 0 0 0 0 0 0 0 0 0 0 + -7.7098 5.0628 -1.2014 H 0 0 0 0 0 0 0 0 0 0 0 0 + -9.0935 5.3789 -3.2180 H 0 0 0 0 0 0 0 0 0 0 0 0 + -9.5294 7.6705 -4.0709 H 0 0 0 0 0 0 0 0 0 0 0 0 + -8.4705 9.6119 -2.9925 H 0 0 0 0 0 0 0 0 0 0 0 0 + -7.0579 9.3041 -0.9946 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 1 0 + 2 3 2 0 + 2 7 1 0 + 3 4 1 0 + 3 8 1 0 + 4 5 2 0 + 4 9 1 0 + 5 6 1 0 + 5 10 1 0 + 6 7 2 0 + 6 11 1 0 + 7 12 1 0 +M END +> (1) +-0.14198333333333335 0.12381666666666666 -0.16608333333333336 -0.10508333333333333 -0.14508333333333334 -0.10508333333333333 -0.16608333333333336 0.14791666666666664 0.13691666666666666 +0.13591666666666666 0.13691666666666666 0.14791666666666664 + +$$$$ diff --git a/openfe/tests/data/htf/t4_lysozyme_data/t4_lysozyme_solvated.pdb.gz b/openfe/tests/data/htf/t4_lysozyme_data/t4_lysozyme_solvated.pdb.gz new file mode 100644 index 000000000..923dc0a1b Binary files /dev/null and b/openfe/tests/data/htf/t4_lysozyme_data/t4_lysozyme_solvated.pdb.gz differ diff --git a/openfe/tests/protocols/openmm_rfe/helpers.py b/openfe/tests/protocols/openmm_rfe/helpers.py new file mode 100644 index 000000000..a0c09bdd6 --- /dev/null +++ b/openfe/tests/protocols/openmm_rfe/helpers.py @@ -0,0 +1,206 @@ +from itertools import chain + +import numpy as np +import openmm +from gufe import LigandAtomMapping, ProteinComponent, SolventComponent +from openff.units.openmm import ensure_quantity, from_openmm, to_openmm +from openmm import app, unit +from openmmforcefields.generators import SystemGenerator + +from openfe.protocols.openmm_rfe import _rfe_utils +from openfe.protocols.openmm_rfe._rfe_utils.relative import HybridTopologyFactory +from openfe.protocols.openmm_utils import system_creation + + +def make_htf( + mapping: LigandAtomMapping, + settings, + protein: ProteinComponent = None, + solvent: SolventComponent = None, +) -> HybridTopologyFactory: + """Code copied from the RBFE protocol to make an HTF.""" + + system_generator = SystemGenerator( + forcefields=settings.forcefield_settings.forcefields, + small_molecule_forcefield=settings.forcefield_settings.small_molecule_forcefield, + forcefield_kwargs={ + "constraints": app.HBonds, + "rigidWater": True, + "hydrogenMass": settings.forcefield_settings.hydrogen_mass * unit.amu, + "removeCMMotion": settings.integrator_settings.remove_com, + }, + periodic_forcefield_kwargs={ + "nonbondedMethod": app.PME, + "nonbondedCutoff": 0.9 * unit.nanometers, + }, + barostat=openmm.MonteCarloBarostat( + ensure_quantity(settings.thermo_settings.pressure, "openmm"), + ensure_quantity(settings.thermo_settings.temperature, "openmm"), + settings.integrator_settings.barostat_frequency.m, + ), + cache=None, + ) + small_mols = [mapping.componentA, mapping.componentB] + # copy a lot of code from the RHT protocol + off_small_mols = { + "stateA": [(mapping.componentA, mapping.componentA.to_openff())], + "stateB": [(mapping.componentB, mapping.componentB.to_openff())], + "both": [ + (m, m.to_openff()) + for m in small_mols + if (m != mapping.componentA and m != mapping.componentB) + ], + } + + # c. force the creation of parameters + # This is necessary because we need to have the FF templates + # registered ahead of solvating the system. + for smc, mol in chain( + off_small_mols["stateA"], off_small_mols["stateB"], off_small_mols["both"] + ): + system_generator.create_system(mol.to_topology().to_openmm(), molecules=[mol]) + + # c. get OpenMM Modeller + a dictionary of resids for each component + stateA_modeller, comp_resids = system_creation.get_omm_modeller( + # add the protein if passed + protein_comp=protein, + # add the solvent if passed + solvent_comp=solvent, + small_mols=dict(chain(off_small_mols["stateA"], off_small_mols["both"])), + omm_forcefield=system_generator.forcefield, + solvent_settings=settings.solvation_settings, + ) + # d. get topology & positions + # Note: roundtrip positions to remove vec3 issues + stateA_topology = stateA_modeller.getTopology() + stateA_positions = to_openmm(from_openmm(stateA_modeller.getPositions())) + + # e. create the stateA System + stateA_system = system_generator.create_system( + stateA_modeller.topology, + molecules=[m for _, m in chain(off_small_mols["stateA"], off_small_mols["both"])], + ) + + # 2. Get stateB system + # a. get the topology + stateB_topology, stateB_alchem_resids = _rfe_utils.topologyhelpers.combined_topology( + stateA_topology, + # zeroth item (there's only one) then get the OFF representation + off_small_mols["stateB"][0][1].to_topology().to_openmm(), + exclude_resids=comp_resids[mapping.componentA], + ) + + # b. get a list of small molecules for stateB + stateB_system = system_generator.create_system( + stateB_topology, + molecules=[m for _, m in chain(off_small_mols["stateB"], off_small_mols["both"])], + ) + + # c. Define correspondence mappings between the two systems + ligand_mappings = _rfe_utils.topologyhelpers.get_system_mappings( + mapping.componentA_to_componentB, + stateA_system, + stateA_topology, + comp_resids[mapping.componentA], + stateB_system, + stateB_topology, + stateB_alchem_resids, + # These are non-optional settings for this method + fix_constraints=True, + ) + + # e. Finally get the positions + stateB_positions = _rfe_utils.topologyhelpers.set_and_check_new_positions( + ligand_mappings, + stateA_topology, + stateB_topology, + old_positions=ensure_quantity(stateA_positions, "openmm"), + insert_positions=ensure_quantity(off_small_mols["stateB"][0][1].conformers[0], "openmm"), + ) + return HybridTopologyFactory( + old_system=stateA_system, + old_positions=stateA_positions, + old_topology=stateA_topology, + new_system=stateB_system, + new_positions=stateB_positions, + new_topology=stateB_topology, + old_to_new_atom_map=ligand_mappings["old_to_new_atom_map"], + old_to_new_core_atom_map=ligand_mappings["old_to_new_core_atom_map"], + use_dispersion_correction=settings.alchemical_settings.use_dispersion_correction, + softcore_alpha=settings.alchemical_settings.softcore_alpha, + softcore_LJ_v2=True, + softcore_LJ_v2_alpha=settings.alchemical_settings.softcore_alpha, + interpolate_old_and_new_14s=settings.alchemical_settings.turn_off_core_unique_exceptions, + ) + + +def _make_system_with_cmap( + map_sizes: list[int], + mapped_torsions: list[tuple[int, int, int, int, int, int, int, int, int]] | None = None, + num_atoms: int = 8, +) -> tuple[openmm.System, openmm.app.Topology, openmm.unit.Quantity]: + """ + Build an OpenMM System with a CMAP term based on the provided mapping data. + + Parameters + ---------- + map_sizes : list[int] + A list of integers specifying the sizes of the CMAP grids to be added. + mapped_torsions : list[tuple[int, int, int, int, int, int, int, int, int]], optional + A list of tuples specifying the atom indices for each mapped torsion. + Each tuple should contain 9 integers: the first is the map index, + followed by the 8 atom indices defining the two dihedrals. + If None, a default torsion will be added using the first 8 atoms. + num_atoms : int, optional + The total number of atoms in the system must be >= 8. Default is 8 the minimum required for a single CMAP. + + Returns + ------- + tuple[openmm.System, openmm.app.Topology, openmm.unit.Quantity] + A tuple containing the OpenMM System, Topology, and Positions. + """ + assert num_atoms >= 8, "num_atoms must be at least 8 to accommodate mapped torsions" + system = openmm.System() + # add dummy forces to avoid errors + for force in [ + openmm.NonbondedForce, + openmm.HarmonicBondForce, + openmm.HarmonicAngleForce, + openmm.PeriodicTorsionForce, + ]: + system.addForce(force()) + + for _ in range(num_atoms): + system.addParticle(12.0) # Add carbon-like particles + + # create a CMAP force if we have map sizes to add + if map_sizes: + cmap_force = openmm.CMAPTorsionForce() + + for map_size in map_sizes: + # Create a grid for the CMAP + grid = [0.0] * (map_size * map_size) + cmap_force.addMap(map_size, grid) + + if mapped_torsions is None: + # add a single cmap term for all atoms using the first map + mapped_torsions = [(0, 0, 1, 2, 3, 4, 5, 6, 7)] + + for torsion in mapped_torsions: + cmap_force.addTorsion(torsion[0], *torsion[1:]) + + system.addForce(cmap_force) + + # build a basic topology for the number of atoms bonding each atom to the next + topology = openmm.app.Topology() + chain = topology.addChain() + res = topology.addResidue("RES", chain) + atoms = [] + for i in range(num_atoms): + atom = topology.addAtom(f"C{i + 1}", openmm.app.element.carbon, res) + atoms.append(atom) + if i > 0: + topology.addBond(atoms[i - 1], atoms[i]) + # build a fake set of positions + positions = openmm.unit.Quantity(np.zeros((num_atoms, 3)), openmm.unit.nanometer) + return system, topology, positions diff --git a/openfe/tests/protocols/openmm_rfe/test_hybrid_factory.py b/openfe/tests/protocols/openmm_rfe/test_hybrid_factory.py new file mode 100644 index 000000000..e6454553c --- /dev/null +++ b/openfe/tests/protocols/openmm_rfe/test_hybrid_factory.py @@ -0,0 +1,230 @@ +import copy + +import openmm +import pytest +from openmm import app, unit + +from openfe.protocols.openmm_rfe import _rfe_utils +from openfe.protocols.openmm_rfe._rfe_utils.relative import HybridTopologyFactory +from openfe.tests.protocols.openmm_rfe.helpers import _make_system_with_cmap + + +def test_cmap_system_no_dummy_pme_energy(htf_cmap_chlorobenzene_to_fluorobenzene): + """ + Test that we can make a hybrid topology for a system with conserved CMAP terms not in the alchemical region and that + the hybrid energy matches the end state energy. + """ + htf = htf_cmap_chlorobenzene_to_fluorobenzene["htf"] + # make sure the cmap force was added to the internal store + assert "cmap_torsion_force" in htf._hybrid_system_forces + hybrid_system = htf.hybrid_system + # make sure we can find the force in the system + forces = htf_cmap_chlorobenzene_to_fluorobenzene["forces"] + assert isinstance(forces["CMAPTorsionForce"], openmm.CMAPTorsionForce) + + integrator = openmm.LangevinIntegrator( + 300 * unit.kelvin, 1.0 / unit.picosecond, 0.002 * unit.picoseconds + ) + platform = openmm.Platform.getPlatformByName("Reference") + default_lambda = _rfe_utils.lambdaprotocol.LambdaProtocol() + + hybrid_simulation = app.Simulation( + topology=htf.omm_hybrid_topology, + system=hybrid_system, + integrator=integrator, + platform=platform, + ) + for end_state, ref_system, ref_top, pos in [ + (0, htf._old_system, htf._old_topology, htf._old_positions), + (1, htf._new_system, htf._new_topology, htf._new_positions), + ]: + # set lambda + # set all lambda values to the current end state + for name, func in default_lambda.functions.items(): + val = func(end_state) + hybrid_simulation.context.setParameter(name, val) + # set positions + hybrid_simulation.context.setPositions(pos) + # get the hybrid system energy + hybrid_state = hybrid_simulation.context.getState(getEnergy=True) + hybrid_energy = hybrid_state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) + # now create a reference simulation + ref_simulation = app.Simulation( + topology=ref_top, + system=ref_system, + integrator=copy.deepcopy(integrator), + platform=platform, + ) + ref_simulation.context.setPositions(pos) + ref_state = ref_simulation.context.getState(getEnergy=True) + ref_energy = ref_state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) + # energies should be the same + assert ref_energy == pytest.approx(hybrid_energy, rel=1e-5) + # make sure the energy is non-zero to avoid false positives + assert 0.0 != pytest.approx(hybrid_energy) + + +def test_cmap_missing_cmap_error(): + """Test that an error is raised if a CMAPTorsionForce is only present in one of the end states.""" + with pytest.raises( + RuntimeError, + match="Inconsistent CMAPTorsionForce between end states expected to be present in both", + ): + old_system, old_topology, old_positions = _make_system_with_cmap([4]) + new_system, new_topology, new_positions = _make_system_with_cmap([]) + _ = HybridTopologyFactory( + old_system=old_system, + old_topology=old_topology, + old_positions=old_positions, + new_system=new_system, + new_topology=new_topology, + new_positions=new_positions, + # map all atoms so that one of the cmap atoms is in the alchemical core region + old_to_new_atom_map={0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7}, + old_to_new_core_atom_map={4: 4}, # atom 4 is part of the cmap torsion + ) + + +def test_verify_cmap_incompatible_maps_error(): + """Test that an error is raised if the number of CMAP terms differ between the end states.""" + old_cmap = openmm.CMAPTorsionForce() + new_cmap = openmm.CMAPTorsionForce() + old_cmap.addMap(2, [0.0] * 2 * 2) # add one map + new_cmap.addMap(2, [0.0] * 2 * 2) # add one map + new_cmap.addMap(2, [0.0] * 2 * 2) # add a second map to make them incompatible + with pytest.raises( + RuntimeError, + match="Incompatible CMAPTorsionForce between end states expected to have same number of maps, found old: 1 and new: 2", + ): + _ = HybridTopologyFactory._verify_cmap_compatibility(old_cmap, new_cmap) + + +def test_verify_cmap_incompatible_torsions_error(): + """Test that an error is raised if the number of CMAP torsions differ between the end states.""" + old_cmap = openmm.CMAPTorsionForce() + new_cmap = openmm.CMAPTorsionForce() + old_cmap.addMap(2, [0.0] * 2 * 2) # add one map + new_cmap.addMap(2, [0.0] * 2 * 2) # add one map + # add torsions + old_cmap.addTorsion(0, 0, 1, 2, 3, 4, 5, 6, 7) + new_cmap.addTorsion(0, 0, 1, 2, 3, 4, 5, 6, 7) + new_cmap.addTorsion(0, 1, 2, 3, 4, 5, 6, 7, 8) # add a second torsion to make them incompatible + with pytest.raises( + RuntimeError, + match="Incompatible CMAPTorsionForce between end states expected to have same number of torsions, found old: 1 and new: 2", + ): + _ = HybridTopologyFactory._verify_cmap_compatibility(old_cmap, new_cmap) + + +def test_cmap_maps_incompatible_error(): + """Test that an error is raised if the CMAP maps differ between the end states using a dummy system. + In this case the map parameters differ for map index 0 between the old and new systems + """ + old_system, old_topology, old_positions = _make_system_with_cmap([4]) + new_system, new_topology, new_positions = _make_system_with_cmap([3]) + with pytest.raises( + RuntimeError, + match="Incompatible CMAPTorsionForce map parameters found between end states for map 0 expected", + ): + _ = HybridTopologyFactory( + old_system=old_system, + old_topology=old_topology, + old_positions=old_positions, + new_system=new_system, + new_topology=new_topology, + new_positions=new_positions, + # map all atoms so they end up in the environment + old_to_new_atom_map={0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7}, + old_to_new_core_atom_map={}, + ) + + +def test_cmap_torsions_incompatible_error(): + """Test that an error is raised if the CMAP torsions differ between the end states using a dummy system. + In this case there is an extra cmap torsion in the new system not present in the old system.""" + old_system, old_topology, old_positions = _make_system_with_cmap([4], num_atoms=12) + new_system, new_topology, new_positions = _make_system_with_cmap( + [4], + num_atoms=12, + mapped_torsions=[ + # change the mapped atoms from the default + (0, 4, 5, 6, 7, 8, 9, 10, 11) + ], + ) + with pytest.raises( + RuntimeError, match="Incompatible CMAPTorsionForce term found between end states for atoms " + ): + _ = HybridTopologyFactory( + old_system=old_system, + old_topology=old_topology, + old_positions=old_positions, + new_system=new_system, + new_topology=new_topology, + new_positions=new_positions, + # map all atoms so they end up in the environment + old_to_new_atom_map={ + 0: 0, + 1: 1, + 2: 2, + 3: 3, + 4: 4, + 5: 5, + 6: 6, + 7: 7, + 8: 8, + 9: 9, + 10: 10, + 11: 11, + }, + old_to_new_core_atom_map={}, + ) + + +def test_cmap_map_index_incompatible_error(): + """Test that an error is raised if the CMAP map indices differ between the end states using a dummy system. + In this case the map index for the single cmap torsion differs between the old and new systems.""" + old_system, old_topology, old_positions = _make_system_with_cmap([4, 5]) + new_system, new_topology, new_positions = _make_system_with_cmap( + [4, 5], + mapped_torsions=[ + # change the map index from the default + (1, 0, 1, 2, 3, 4, 5, 6, 7) + ], + ) + # modify one of the torsions in the new system to make them incompatible + with pytest.raises( + RuntimeError, + match="Incompatible CMAPTorsionForce map index found between end states for atoms ", + ): + _ = HybridTopologyFactory( + old_system=old_system, + old_topology=old_topology, + old_positions=old_positions, + new_system=new_system, + new_topology=new_topology, + new_positions=new_positions, + # map all atoms so they end up in the environment + old_to_new_atom_map={0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7}, + old_to_new_core_atom_map={}, + ) + + +def test_cmap_in_alchemical_region_error(): + """Test that an error is raised if a CMAP torsion is in the alchemical region.""" + old_system, old_topology, old_positions = _make_system_with_cmap([4]) + new_system, new_topology, new_positions = _make_system_with_cmap([4]) + with pytest.raises( + RuntimeError, + match="Incompatible CMAPTorsionForce term found in alchemical region for old system atoms", + ): + _ = HybridTopologyFactory( + old_system=old_system, + old_topology=old_topology, + old_positions=old_positions, + new_system=new_system, + new_topology=new_topology, + new_positions=new_positions, + # map all atoms so that one of the cmap atoms is in the alchemical core region + old_to_new_atom_map={0: 0, 1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7}, + old_to_new_core_atom_map={4: 4}, # atom 4 is part of the cmap torsion + )