Add CMAPTorsionForce support

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 <https://github.com/OpenFreeEnergy/openfe/pull/1695>`_
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

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)

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,
+    }