pub_soft/openfe
3
0
Fork 0
mirror of https://github.com/OpenFreeEnergy/openfe.git synced 2026-06-04 14:14:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add CMAPTorsionForce support (#1695)

Browse Source 
Adds support for CMAPTorsionForce in relative hybrid topology simulations.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Irfan Alibay <IAlibay@users.noreply.github.com>
Co-authored-by: Alyssa Travitz <31974495+atravitz@users.noreply.github.com>
This commit is contained in:
Josh Horton
2025-12-01 14:19:54 +00:00
committed by GitHub
parent 013b4b3213
commit 35a999907e
9 changed files with 831 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
news/cmap_support.rst Normal file
View File

@@ -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>

154
openfe/protocols/openmm_rfe/_rfe_utils/relative.py
View File

@@ -215,6 +215,9 @@ class HybridTopologyFactory:
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 @@ class HybridTopologyFactory:
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 @@ class HybridTopologyFactory:
# 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)

156
openfe/tests/conftest.py
View File

@@ -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 @@ from importlib import resources
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 @@ try:
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,
}

0
openfe/tests/data/htf/__init__.py Normal file
View File

34
openfe/tests/data/htf/t4_lysozyme_data/chlorobenzene.sdf Normal file
View File

@@ -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
> <atom.dprop.PartialCharge> (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
$$$$

34
openfe/tests/data/htf/t4_lysozyme_data/fluorobenzene.sdf Normal file
View File

@@ -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
> <atom.dprop.PartialCharge> (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
$$$$

BIN
openfe/tests/data/htf/t4_lysozyme_data/t4_lysozyme_solvated.pdb.gz Normal file
View File

Binary file not shown.

206
openfe/tests/protocols/openmm_rfe/helpers.py Normal file
View File

@@ -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

230
openfe/tests/protocols/openmm_rfe/test_hybrid_factory.py Normal file
View File

@@ -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
)