diff --git a/doc/userguide.md b/doc/userguide.md index 39dfe24..75c1a6e 100644 --- a/doc/userguide.md +++ b/doc/userguide.md @@ -376,6 +376,34 @@ to specify which behavior your model uses when doing mechanical embedding in a p will be raised to inform you if this information is needed and not provided; OpenMM-ML will not assume either choice automatically. +#### Molecules Spanning the ML-MM Region + +OpenMM-ML's mechanical embedding implementation supports the link-atom method for molecules having bonds crossing the +boundary between the ML and MM regions. If a molecule in the `Topology` provided contains bonds spanning the regions, +then the molecule will appear as is to the MM force field, but will have these bonds capped by hydrogen atoms when its +fragment(s) within the ML region are evaluated by the ML potential. + +The fictitious link atoms added are implemented as virtual sites which will be inserted into the `System` and `Topology` +in use. By default, `createMixedSystem()` only returns the `System`, but passing `returnInfo=True` returns a dictionary +instead, with keys `system` (the `System`), `topology` (a modified copy of the `Topology` with the added sites), and +`oldToNew` (a list of atom indices serving as a mapping from those in the original `Topology` to those in the modified +one). Since they are non-physical sites added only for implementing the method, the link atoms will be added to their +own chain in the `Topology` separate from any existing chains. + +Each link atom is maintained at a fixed distance along its respective bond crossing the boundary. By default, this +distance is chosen based on the covalent radius of the atom on the ML side of the bond. To override these distances, +pass `linkAtomDistances=[...]` to `createMixedSystem()` with a list of tuples `(atom1, atom2, distance)` for each pair +of atoms for which to use a custom distance. + +Multiple link bonds from the same atom in the ML region are supported. However, OpenMM-ML will raise an error if an ML +subset is given that would create more than one link bond to the same atom in the MM region. Such a configuration would +place the associated link atoms too close to one another. + +To avoid double-counting bonded interactions between the MM force field and ML potential, OpenMM-ML will delete: +- All MM bonds contained completely within the ML region. +- All MM angles and torsions contained completely within the ML region, but accounting for the presence of any link + atoms and bonds leaving the region. + ## Other Packages OpenMM-ML is based on a plugin architecture, allowing other packages to provide their own interfaces to it. The diff --git a/openmmml/embeddings/mechanicalembedding.py b/openmmml/embeddings/mechanicalembedding.py index 62aa9f6..9d6561a 100644 --- a/openmmml/embeddings/mechanicalembedding.py +++ b/openmmml/embeddings/mechanicalembedding.py @@ -33,6 +33,8 @@ from openmmml.embeddings import utilities import openmm import openmm.app +import copy +import typing class MechanicalEmbeddingFactory(EmbeddingFactory): """This is the factory that creates MechanicalEmbedding objects.""" @@ -71,7 +73,8 @@ def createMixedSystem(self, atoms: list[int], forceGroup: int, interpolate: bool, - **args) -> openmm.System: + returnInfo: bool = False, + **args) -> openmm.System | dict[str, typing.Any]: periodic = system.usesPeriodicBoundaryConditions() @@ -123,7 +126,8 @@ def createMixedSystem(self, # Create the new system with ML-ML interactions to be computed by the ML # potential removed. - newSystem = utilities.removeBonds(system, atoms, True) + linkBonds = utilities.findLinkBonds(topology, atoms) + newSystem = utilities.removeBonds(system, topology, atoms, True) for force in newSystem.getForces(): if isinstance(force, openmm.NonbondedForce): @@ -150,7 +154,7 @@ def createMixedSystem(self, force.setExceptionsUsePeriodicBoundaryConditions(periodic) elif isinstance(force, openmm.CustomNonbondedForce): - utilities.makeCustomNonbondedExclusions(force, atoms) + utilities.addCustomNonbondedExclusions(force, atoms) if excludeLongRange: # Prepare a force to calculate the PME energy of the ML-ML region. @@ -165,10 +169,21 @@ def createMixedSystem(self, for atom in range(newSystem.getNumParticles()): excludeForce.addParticle(mmLongRangeForce.getParticleParameters(atom)[0] if atom in atomSet else 0, 1, 0) + newTopology = copy.deepcopy(topology) + if interpolate: + # For interpolation setup to work, we need to modify the original + # system so that its NonbondedForce also has the virtual site. + system = copy.deepcopy(system) + systemList = [system, newSystem] + else: + systemList = [newSystem] + capIndices, oldToNew = utilities.addLinkAtomSites(newTopology, systemList, linkBonds, args.get("linkAtomDistances", [])) + atomsWithCaps = atoms + capIndices + if interpolate: interpolator = utilities.InterpolationHelper() - interpolator.addMLPotentialTerms(potential, topology, atoms, forceGroup, **args) - interpolator.addMMBondedTerms(system, atoms) + interpolator.addMLPotentialTerms(potential, newTopology, atomsWithCaps, forceGroup, **args) + interpolator.addMMBondedTerms(system, topology, atoms) interpolator.setupNonbonded(newSystem, system) if excludeLongRange: interpolator.addMLTerm(excludeForce, "-{}") @@ -182,6 +197,9 @@ def createMixedSystem(self, cvForce.addCollectiveVariable("excludeForce", excludeForce) newSystem.addForce(cvForce) - potential.addForces(topology, newSystem, atoms, forceGroup, **args) + potential.addForces(newTopology, newSystem, atomsWithCaps, forceGroup, **args) - return newSystem + if returnInfo: + return dict(system=newSystem, topology=newTopology, oldToNew=oldToNew) + else: + return newSystem diff --git a/openmmml/embeddings/utilities.py b/openmmml/embeddings/utilities.py index d1a9ecf..6f344cd 100644 --- a/openmmml/embeddings/utilities.py +++ b/openmmml/embeddings/utilities.py @@ -35,23 +35,170 @@ import openmm.unit as unit from openmmml.mlpotential import MLPotentialImpl -def removeBonds(system: openmm.System, atoms: list[int], removeInSet: bool) -> openmm.System: +COVALENT_RADII = [ + 0, 32, 46, 120, 94, 77, 75, 71, 63, 64, 67, 140, 125, 112, 104, 110, 102, + 99, 96, 176, 154, 133, 122, 121, 110, 107, 104, 100, 99, 101, 109, 112, 109, + 114, 110, 113, 117, 189, 167, 147, 139, 132, 124, 114, 112, 112, 108, 114, + 123, 128, 126, 126, 123, 132, 131, 209, 176, 162, 147, 158, 157, 156, 155, + 151, 152, 151, 150, 149, 149, 148, 153, 146, 137, 131, 123, 118, 115, 111, + 112, 112, 132, 130, 130, 136, 131, 138, 142, 200, 181, 167, 158, 152, 153, + 154, 155, 149, 149, 151, 151, 148, 150, 156, 158, 145, 141, 134, 129, 127, + 121, 115, 114, 109, 122, 136, 143, 146, 158, 148, 157 +] * openmm.unit.picometer +""" +Default covalent radii to use for assigning distances in the link-atom method. +This set is taken from MLIPOps, which chose them to be consistent with the +simple-dftd3 library (https://github.com/dftd3/simple-dftd3). They are taken +from Pyykko and Atsumi, Chem. Eur. J. 15, 2009, 188-197, except that the radii +of metals have been reduced by 10%. +""" + +def findLinkBonds(topology: openmm.app.Topology, atoms: list[int]) -> list[tuple[int, int]]: + """ + Finds bonds in a topology between a subset of atoms and its complement. + + Parameters + ---------- + topology: Topology + The Topology to find bonds in. + atoms: list[int] + A set of atom indices. + + Returns + ------- + A list of atom index pairs corresponding to "link bonds", i.e., bonds + between atoms in the subset and atoms not in the subset. The first index of + every pair will correspond to the atom in the subset. + """ + + atomSet = set(atoms) + linkBonds = [] + + for bond in topology.bonds(): + atom1 = bond.atom1.index + atom2 = bond.atom2.index + atom1Included = atom1 in atomSet + atom2Included = atom2 in atomSet + if atom1Included and not atom2Included: + linkBonds.append((atom1, atom2)) + if atom2Included and not atom1Included: + linkBonds.append((atom2, atom1)) + + return linkBonds + +def addLinkAtomSites(topology: openmm.app.Topology, systems: list[openmm.System], linkBonds: list[tuple[int, int]], linkAtomDistances: list[tuple[int, int, unit.Quantity]]) -> tuple[list[int], list[int]]: + """ + Adds virtual sites to systems and a topology for the link-atom method. + + Each virtual site represents a hydrogen atom capping a bond spanning the ML + and MM regions of an ML/MM simulation. By default, the distance from the + atom on the ML side of such a bond to the virtual site is calculated based + on the covalent radii of the ML atom and hydrogen, but this is overridable + for particular link bonds using `linkAtomDistances`. + + Parameters + ---------- + systems: list[System] + The list of Systems to modify in place by adding virtual sites. + topology: Topology + The Topology to look up atomic numbers from and modify in place by + adding virtual sites. + linkBonds: list[tuple[int, int]] + A list of bonds to add virtual sites to, in the format returned by + `findLinkBonds()`. + linkAtomDistances: list[tuple[int, int, Quantity]] + A list of link bonds with virtual site distances to set manually. + + Returns + ------- + A list of indices corresponding to the virtual sites added to the systems, + and a list serving as a mapping from atom indices in the original Topology + to those in the modified Topology. + + The current implementation always appends virtual sites to the end of each + System and the Topology (in a new Chain), so the mapping will always be an + identity mapping. + """ + + linkAtomDistanceTable = {} + for atom1, atom2, distance in linkAtomDistances: + linkAtomDistanceTable[min(atom1, atom2), max(atom1, atom2)] = distance + + # Update the topology with virtual sites to be added, and load data from it. + + oldToNew = list(range(topology.getNumAtoms())) + siteIndices = [] + if linkBonds: + siteChain = topology.addChain() + for site in range(len(linkBonds)): + siteIndices.append(topology.addAtom(f"V{site}", openmm.app.element.hydrogen, topology.addResidue(f"V{site}", siteChain)).index) + atomicNumbers = [atom.element.atomic_number for atom in topology.atoms()] + + # Add virtual sites to the systems. + + mmAtoms = set() + for mlAtom, mmAtom in linkBonds: + # Nothing in the implementation prevents multiple link bonds to the same + # MM atom, but this would place virtual sites too close to each other. + if mmAtom in mmAtoms: + raise ValueError(f"Multiple link bonds to MM atom {mmAtom}") + mmAtoms.add(mmAtom) + + key = min(mlAtom, mmAtom), max(mlAtom, mmAtom) + if key in linkAtomDistanceTable: + distance = linkAtomDistanceTable[key] + else: + distance = COVALENT_RADII[atomicNumbers[mlAtom]] + COVALENT_RADII[1] + + for system in systems: + site = openmm.LocalCoordinatesSite([mlAtom, mmAtom], [1.0, 0.0], [-1.0, 1.0], [0.0, 0.0], [distance, 0.0, 0.0]) + system.setVirtualSite(system.addParticle(0.0), site) + + needExclusions = False + for force in system.getForces(): + if isinstance(force, openmm.NonbondedForce): + force.addParticle(0.0, 0.0, 0.0) + elif isinstance(force, openmm.CustomNonbondedForce): + force.addParticle([0] * force.getNumPerParticleParameters()) + needExclusions = True + + # If there was a CustomNonbondedForce, the virtual site will need to + # have an exclusion with every other particle. To make the set of + # exclusions equal, this is also required for the NonbondedForce. + if needExclusions: + excludeAtom = system.getNumParticles() - 1 + for force in system.getForces(): + if isinstance(force, openmm.NonbondedForce): + for otherAtom in range(excludeAtom): + force.addException(otherAtom, excludeAtom, 0.0, 0.0, 0.0) + elif isinstance(force, openmm.CustomNonbondedForce): + for otherAtom in range(excludeAtom): + force.addExclusion(otherAtom, excludeAtom) + + return siteIndices, oldToNew + +def removeBonds(system: openmm.System, topology: openmm.app.Topology, atoms: list[int], removeInSet: bool) -> openmm.System: """ Copy a System, removing all bonded interactions between atoms in (or not in) a particular set. + Bonds spanning the set and its complement will not be removed. Angles and + torsions will be removed if they would remain in a subset of the topology + including the specified set of atoms, bonds between them, and any link atoms + and bonds that would be inserted due to bonds leaving the set. + Parameters ---------- system: System The System to copy. + topology: Topology + A corresponding Topology used to identify bonds in the System. atoms: list[int] A set of atom indices. removeInSet: bool If True, any bonded term connecting atoms in the specified set is removed. If False, any term that does *not* connect atoms in the specified set is removed. - removeConstraints: bool - If True, remove constraints between pairs of atoms in the set. Returns ------- @@ -60,6 +207,21 @@ def removeBonds(system: openmm.System, atoms: list[int], removeInSet: bool) -> o """ atomSet = set(atoms) + expandedAtomSet = set(atomSet) + + bondedToAtom = [set() for _ in topology.atoms()] + for bond in topology.bonds(): + atom1 = bond.atom1.index + atom2 = bond.atom2.index + bondedToAtom[atom1].add(atom2) + bondedToAtom[atom2].add(atom1) + if atom1 in atomSet: + expandedAtomSet.add(atom2) + if atom2 in atomSet: + expandedAtomSet.add(atom1) + + def isBondedTo(a1, a2): + return a1 in bondedToAtom[a2] # Create an XML representation of the System. @@ -67,28 +229,55 @@ def removeBonds(system: openmm.System, atoms: list[int], removeInSet: bool) -> o xml = openmm.XmlSerializer.serialize(system) root = ET.fromstring(xml) - # This function decides whether a bonded interaction should be removed. + # These functions decide whether a bonded interaction should be removed. + + def isBondInSet(a1, a2): + return a1 in atomSet and a2 in atomSet + + def isAngleInSet(a1, a2, a3): + return a1 in expandedAtomSet and a2 in atomSet and a3 in expandedAtomSet + + def isTorsionInSet(a1, a2, a3, a4): + if isBondedTo(a1, a2) and isBondedTo(a2, a3) and isBondedTo(a3, a4): + return a2 in atomSet and a3 in atomSet + elif isBondedTo(a1, a2) and isBondedTo(a1, a3) and isBondedTo(a1, a4): + return a1 in atomSet + elif isBondedTo(a2, a1) and isBondedTo(a2, a3) and isBondedTo(a2, a4): + return a2 in atomSet + elif isBondedTo(a3, a1) and isBondedTo(a3, a2) and isBondedTo(a3, a4): + return a3 in atomSet + elif isBondedTo(a4, a1) and isBondedTo(a4, a2) and isBondedTo(a4, a3): + return a4 in atomSet + else: + raise ValueError("Unrecognized torsion kind (neither proper nor improper)") - def shouldRemove(termAtoms): - return all(a in atomSet for a in termAtoms) == removeInSet + def isCMAPInSet(a1, a2, a3, a4, b1, b2, b3, b4): + return ( + a1 in expandedAtomSet and a2 in atomSet and a3 in atomSet and a4 in expandedAtomSet and + b1 in expandedAtomSet and b2 in atomSet and b3 in atomSet and b4 in expandedAtomSet + ) # Remove bonds, angles, and torsions. for bonds in root.findall('./Forces/Force/Bonds'): for bond in bonds.findall('Bond'): bondAtoms = [int(bond.attrib[p]) for p in ('p1', 'p2')] - if shouldRemove(bondAtoms): + if isBondInSet(*bondAtoms) == removeInSet: bonds.remove(bond) for angles in root.findall('./Forces/Force/Angles'): for angle in angles.findall('Angle'): angleAtoms = [int(angle.attrib[p]) for p in ('p1', 'p2', 'p3')] - if shouldRemove(angleAtoms): + if isAngleInSet(*angleAtoms) == removeInSet: angles.remove(angle) for torsions in root.findall('./Forces/Force/Torsions'): for torsion in torsions.findall('Torsion'): - torsionLabels = ('p1', 'p2', 'p3', 'p4') if 'p1' in torsion.attrib else ('a1', 'a2', 'a3', 'a4', 'b1', 'b2', 'b3', 'b4') - torsionAtoms = [int(torsion.attrib[p]) for p in torsionLabels] - if shouldRemove(torsionAtoms): + if 'p1' in torsion.attrib: + torsionAtoms = [int(torsion.attrib[p]) for p in ('p1', 'p2', 'p3', 'p4')] + inSet = isTorsionInSet(*torsionAtoms) + else: + cmapAtoms = [int(torsion.attrib[p]) for p in ('a1', 'a2', 'a3', 'a4', 'b1', 'b2', 'b3', 'b4')] + inSet = isCMAPInSet(*cmapAtoms) + if inSet == removeInSet: torsions.remove(torsion) # Create a new System from it. @@ -204,7 +393,7 @@ def addMLPotentialTerms(self, potential: MLPotentialImpl, topology: openmm.app.T for force in tempSystem.getForces(): self.addMLTerm(copy.deepcopy(force)) - def addMMBondedTerms(self, mmSystem: openmm.System, atoms: list[int]) -> None: + def addMMBondedTerms(self, mmSystem: openmm.System, topology: openmm.app.Topology, atoms: list[int]) -> None: """ Helper function to add all bonded forces removed from the ML region of an ML/MM system as MM terms for interpolation. @@ -215,11 +404,13 @@ def addMMBondedTerms(self, mmSystem: openmm.System, atoms: list[int]) -> None: A pure MM system containing all (ML and MM region) bonded terms. This will not be modified and is only used as a reference for the terms to interpolate. + topology: openmm.app.Topology + A corresponding Topology used to find the bonds in the System. atoms: list[int] The indices of the ML region atoms in the ML/MM system. """ - bondedSystem = removeBonds(mmSystem, atoms, False) + bondedSystem = removeBonds(mmSystem, topology, atoms, False) for force in bondedSystem.getForces(): if hasattr(force, "addBond") or hasattr(force, "addAngle") or hasattr(force, "addTorsion"): self.addMMTerm(copy.deepcopy(force)) diff --git a/openmmml/mlpotential.py b/openmmml/mlpotential.py index 260b20e..b9aa799 100644 --- a/openmmml/mlpotential.py +++ b/openmmml/mlpotential.py @@ -35,6 +35,7 @@ import os import shutil import tempfile +import typing import urllib.request import sys if sys.version_info < (3, 10): @@ -124,7 +125,8 @@ def createMixedSystem(self, forceGroup: int, interpolate: bool, embedding: str, - **args) -> openmm.System: + returnInfo: bool = False, + **args) -> openmm.System | dict[str, typing.Any]: """Creates a mixed system using a potential-specific embedding method. This is invoked by MLPotential.createMixedSystem(). It will only be @@ -155,13 +157,24 @@ def createMixedSystem(self, embedding: str the name of the embedding method (will always be one in the list returned by the getSupportedEmbeddings() method) + returnInfo: bool + whether to return a dictionary of key-value pairs containing a new + System with additional information, instead of the System alone args: any additional arguments for the potential or embedding method Returns ------- - a newly created System object that uses this potential function and the - requested embedding method to model the Topology + A newly created System object that uses this potential function and the + requested embedding method to model the Topology, or a dictionary of + key-value pairs if the implementation supports returnInfo and it is + True. The dictionary must contain the System as 'system', a Topology + (that may be different from the one provided) as 'topology', and a list + mapping atom indices in the original Topology to those in the returned + Topology as 'oldToNew'. If the implementation does not support + returnInfo, the System returned must be compatible with the provided + Topology. In no case may the implementation modify the given Topology + in place; it must copy it first or create a new Topology. """ raise NotImplementedError('Subclasses must implement createMixedSystem()') @@ -316,7 +329,8 @@ def createMixedSystem(self, forceGroup: int = 0, interpolate: bool = False, embedding: str = 'mechanical', - **args) -> openmm.System: + returnInfo: bool = False, + **args) -> openmm.System | dict[str, typing.Any]: """Create a System that is partly modeled with this potential and partly with a conventional force field. @@ -361,6 +375,9 @@ def createMixedSystem(self, methods may be available, as well as embedding methods specific to the ML potential selected. MLPotential.getSupportedEmbeddings() will report all embedding methods accepted by the potential. + returnInfo: bool + whether to return a dictionary of key-value pairs containing a new + System with additional information, instead of the System alone args: particular potential functions or embedding methods may define additional arguments that can be used to customize them. See the @@ -369,7 +386,12 @@ def createMixedSystem(self, Returns ------- - a newly created System object that uses this potential function to model the Topology + A newly created System object that uses this potential function and the + requested embedding method to model the Topology, or a dictionary of + key-value pairs if returnInfo is True. The dictionary will contain the + System as 'system', a (possibly modified) Topology as 'topology', and a + list mapping atom indices in the original Topology to those in the + returned Topology as 'oldToNew'. """ atomList = list(atoms) @@ -377,11 +399,21 @@ def createMixedSystem(self, # See if we are given an embedding name that the potential can handle. customEmbeddings = self._impl.getSupportedEmbeddings() if embedding in customEmbeddings: - system = self._impl.createMixedSystem(topology, system, atomList, forceGroup, interpolate, embedding, **args) + systemOrInfo = self._impl.createMixedSystem(topology, system, atomList, forceGroup, interpolate, embedding, returnInfo=returnInfo, **args) else: # Fall back on an embedding plugin. embeddingInstance = MLPotential._embeddingFactories[embedding].createEmbedding(embedding) - system = embeddingInstance.createMixedSystem(self._impl, topology, system, atomList, forceGroup, interpolate, **args) + systemOrInfo = embeddingInstance.createMixedSystem(self._impl, topology, system, atomList, forceGroup, interpolate, returnInfo=returnInfo, **args) + + if returnInfo: + if isinstance(systemOrInfo, openmm.System): + # The potential or embedding didn't support returnInfo. + info = dict(system=systemOrInfo, topology=topology, oldToNew=list(range(topology.getNumAtoms()))) + else: + info = systemOrInfo + system = info["system"] + else: + system = systemOrInfo if removeConstraints: # Remove all constraints with both atoms in the ML subset. @@ -394,7 +426,7 @@ def createMixedSystem(self, for constraint in reversed(constraintsToRemove): system.removeConstraint(constraint) - return system + return info if returnInfo else system def getSupportedEmbeddings(self) -> list[str]: """Retrieves a list of the names of all of the supported embedding @@ -482,7 +514,8 @@ def createMixedSystem(self, atoms: list[int], forceGroup: int, interpolate: bool, - **args): + returnInfo: bool = False, + **args) -> openmm.System | dict[str, typing.Any]: """Creates a mixed system using the embedding method. This is invoked by MLPotential.createMixedSystem(). It must be diff --git a/test/TestMechanicalEmbedding.py b/test/TestMechanicalEmbedding.py index 73797ef..ac01f50 100644 --- a/test/TestMechanicalEmbedding.py +++ b/test/TestMechanicalEmbedding.py @@ -12,25 +12,54 @@ # Get the path to the test data test_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") -atol = 0.01 +atol = 0.02 @pytest.mark.parametrize("platform_int", list(platform_ints)) class TestMechanicalEmbedding: - def getTopologyPositionsSubset(self, topology, positions, subset): modeller = openmm.app.Modeller(topology, positions) modeller.delete([atom for atom in topology.atoms() if atom.index not in subset]) return modeller.getTopology(), modeller.getPositions() + def getBondedTerms(self, system): + bonds = set() + angles = set() + torsions = set() + cmaps = set() + + for force in system.getForces(): + if isinstance(force, openmm.HarmonicBondForce): + for i in range(force.getNumBonds()): + bond = tuple(force.getBondParameters(i)[:2]) + bonds.add(min(bond, bond[::-1])) + elif isinstance(force, openmm.HarmonicAngleForce): + for i in range(force.getNumAngles()): + angle = tuple(force.getAngleParameters(i)[:3]) + angles.add(min(angle, angle[::-1])) + elif isinstance(force, openmm.PeriodicTorsionForce): + for i in range(force.getNumTorsions()): + torsion = tuple(force.getTorsionParameters(i)[:4]) + torsions.add(min(torsion, torsion[::-1])) + elif isinstance(force, openmm.CMAPTorsionForce): + for i in range(force.getNumTorsions()): + cmap = tuple(force.getTorsionParameters(i)[1:]) + cmaps.add((min(cmap[:4], cmap[:4][::-1]), min(cmap[4:], cmap[4:][::-1]))) + + return bonds, angles, torsions, cmaps + @pytest.mark.parametrize("periodic", (False, True)) @pytest.mark.parametrize("interpolate", (False, True)) - def testEmbedding(self, platform_int, periodic, interpolate): + @pytest.mark.parametrize("ff_family", ("amber", "charmm")) + def testEmbedding(self, platform_int, periodic, interpolate, ff_family): """ Mechanical embedding for a non-periodic system, or for a periodic long-range system (in both cases, all periodic images if any are present are included or excluded, so the verification calculation is the same). """ + if ff_family == "charmm" and interpolate: + pytest.skip("Interpolation not yet supported with CustomNonbondedForce") + pdb = openmm.app.PDBFile(os.path.join(test_data_dir, "alanine-dipeptide", "alanine-dipeptide-explicit.pdb")) topology_ml_mm = pdb.topology positions_ml_mm = pdb.positions @@ -38,7 +67,14 @@ def testEmbedding(self, platform_int, periodic, interpolate): subset = [atom.index for atom in topology_ml_mm.atoms() if atom.residue.chain.index == 0] topology_ml, positions_ml = self.getTopologyPositionsSubset(topology_ml_mm, positions_ml_mm, set(subset)) - mm_force_field = openmm.app.ForceField("amber19-all.xml", "amber19/tip3pfb.xml") + if ff_family == "amber": + # Amber will result in an ordinary NonbondedForce only + mm_force_field = openmm.app.ForceField("amber19-all.xml", "amber19/tip3pfb.xml") + elif ff_family == "charmm": + # CHARMM has NBFix and so a CustomNonbondedForce will also be used + mm_force_field = openmm.app.ForceField("charmm36_2024.xml", "charmm36_2024/water.xml") + else: + raise NotImplementedError ml_potential = MLPotential("ase") from mace.calculators.foundations_models import mace_off @@ -54,6 +90,8 @@ def testEmbedding(self, platform_int, periodic, interpolate): for force in mm_system_ml.getForces(): if isinstance(force, openmm.NonbondedForce): force.setUseDispersionCorrection(False) + elif isinstance(force, openmm.CustomNonbondedForce): + force.setUseLongRangeCorrection(False) platform = openmm.Platform.getPlatform(platform_int) mm_context_ml_mm = openmm.Context(mm_system_ml_mm, openmm.VerletIntegrator(0.001), platform) @@ -210,3 +248,244 @@ def testRemoveConstraints(self, platform_int, remove): assert (atom_1, atom_2) in mm_constraints or (atom_2, atom_1) in mm_constraints if atom_1 in subset_set and atom_2 in subset_set: assert ((atom_1, atom_2) in mixed_constraints or (atom_2, atom_1) in mixed_constraints) != remove + + @pytest.mark.parametrize("override_distance", (False, True)) + @pytest.mark.parametrize("ff_name", ("ethanol.xml", "ethanol_ljforce.xml")) + def testLinkAtomTerms(self, platform_int, override_distance, ff_name): + """ + Test for presence of the appropriate terms and positions of the virtual + sites in the link-atom method. + """ + + pdb = openmm.app.PDBFile(os.path.join(test_data_dir, "ethanol", "ethanol.pdb")) + """ + H4 H6 + | | + H3 - O0 - C1 - C2 - H8 + | | + H5 H7 + """ + + # Expected distances are in nanometers. + expected_cc_distance = 0.1525970013793 # From force field. + if override_distance: + expected_ch_distance = 0.12 + else: + expected_ch_distance = 0.107 # From default covalent radii. + + mm_force_field = openmm.app.ForceField(os.path.join(test_data_dir, "ethanol", ff_name)) + ml_potential = MLPotential("mace-off23-small") + + mm_system = mm_force_field.createSystem(pdb.topology) + args = {} + if override_distance: + args["linkAtomDistances"] = [(1, 2, 0.12)] + mixed_system = ml_potential.createMixedSystem(pdb.topology, mm_system, [0, 1, 3, 4, 5], interpolate=False, **args) + + # Get all of the bonded terms in both systems. + mm_bonds, mm_angles, mm_torsions, _ = self.getBondedTerms(mm_system) + mixed_bonds, mixed_angles, mixed_torsions, _ = self.getBondedTerms(mixed_system) + + # No bonded terms should be added to the mixed system. + assert not mixed_bonds - mm_bonds + assert not mixed_angles - mm_angles + assert not mixed_torsions - mm_torsions + + # The appropriate terms should be removed from the mixed system. + assert mm_bonds - mixed_bonds == {(0, 1), (0, 3), (1, 4), (1, 5)} + assert mm_angles - mixed_angles == {(0, 1, 2), (0, 1, 4), (0, 1, 5), (1, 0, 3), (2, 1, 4), (2, 1, 5), (4, 1, 5)} + assert mm_torsions - mixed_torsions == {(2, 1, 0, 3), (3, 0, 1, 4), (3, 0, 1, 5)} + + platform = openmm.Platform.getPlatform(platform_int) + context = openmm.Context(mixed_system, openmm.LangevinIntegrator(300, 1, 0.001), platform) + context.setPositions(pdb.positions + [openmm.Vec3(0, 0, 0)] * openmm.unit.nanometer) + context.computeVirtualSites() + + def check_positions(): + positions = context.getState(positions=True).getPositions(asNumpy=True) / openmm.unit.nanometer + delta_c1_c2 = positions[2] - positions[1] + delta_c1_vs = positions[9] - positions[1] + dist_c1_c2 = np.linalg.norm(delta_c1_c2) + dist_c1_vs = np.linalg.norm(delta_c1_vs) + + # Virtual site should be the appropriate distance from C1. + assert np.isclose(dist_c1_vs, expected_ch_distance) + # Virtual site should be in line with C1-C2. + assert np.isclose(delta_c1_c2 @ delta_c1_vs, dist_c1_c2 * dist_c1_vs) + # C1-C2 distance should be appropriate. + assert dist_c1_c2 < 1.5 * expected_cc_distance + + # Check positions, run some dynamics, and check again. + check_positions() + openmm.LocalEnergyMinimizer.minimize(context) + context.getIntegrator().step(1000) + check_positions() + + def testLinkAtomForbidden(self, platform_int): + """ + Ensure that multiple ML-MM bonds to the same MM atom are disallowed. + """ + + pdb = openmm.app.PDBFile(os.path.join(test_data_dir, "ethanol", "ethanol.pdb")) + mm_force_field = openmm.app.ForceField(os.path.join(test_data_dir, "ethanol", "ethanol.xml")) + ml_potential = MLPotential("mace-off23-small") + + mm_system = mm_force_field.createSystem(pdb.topology) + with pytest.raises(ValueError, match="Multiple link bonds to MM atom 1"): + ml_potential.createMixedSystem(pdb.topology, mm_system, [0, 3, 4, 5]) + + def testLinkAtomMultipleRegions(self, platform_int): + """ + Check that the correct bonded terms are present in a molecule with + multiple ML and MM subregions. + """ + + topology = openmm.app.Topology() + chain = topology.addChain() + atoms = [topology.addAtom("X", openmm.app.element.carbon, topology.addResidue("X", chain)) for _ in range(22)] + for pair in zip(atoms[:-1], atoms[1:]): + topology.addBond(*pair) + + mm_system = openmm.System() + + bond_force = openmm.HarmonicBondForce() + for i in range(len(atoms) - 1): + bond_force.addBond(i, i + 1, 1, 1) + mm_system.addForce(bond_force) + + angle_force = openmm.HarmonicAngleForce() + for i in range(len(atoms) - 2): + angle_force.addAngle(i, i + 1, i + 2, 1, 1) + mm_system.addForce(angle_force) + + torsion_force = openmm.PeriodicTorsionForce() + for i in range(len(atoms) - 3): + torsion_force.addTorsion(i, i + 1, i + 2, i + 3, 1, 0, 1) + mm_system.addForce(torsion_force) + + cmap_force = openmm.CMAPTorsionForce() + for i in range(len(atoms) - 4): + cmap_force.addTorsion(0, i, i + 1, i + 2, i + 3, i + 1, i + 2, i + 3, i + 4) + mm_system.addForce(cmap_force) + + """ + The following should cover all the possible configurations of up to + five atoms forward or in reverse (excluding forbidden ML-MM-ML). + + 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 + ML-ML-ML-ML-ML-MM-MM-MM-MM-MM-ML-ML-ML-MM-MM-MM-ML-MM-MM-ML-ML-MM + """ + mixed_system = MLPotential("mace-off23-small").createMixedSystem(topology, mm_system, [0, 1, 2, 3, 4, 10, 11, 12, 16, 19, 20]) + mixed_bonds, mixed_angles, mixed_torsions, mixed_cmaps = self.getBondedTerms(mixed_system) + + assert mixed_bonds == {(i, i + 1) for i in [4, 5, 6, 7, 8, 9, 12, 13, 14, 15, 16, 17, 18, 20]} + assert mixed_angles == {(i, i + 1, i + 2) for i in [4, 5, 6, 7, 8, 12, 13, 14, 16, 17]} + assert mixed_torsions == {(i, i + 1, i + 2, i + 3) for i in [3, 4, 5, 6, 7, 8, 11, 12, 13, 14, 15, 16, 17]} + assert mixed_cmaps == {((i, i + 1, i + 2, i + 3), (i + 1, i + 2, i + 3, i + 4)) for i in [2, 3, 4, 5, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 17]} + + def testLinkAtomImpropers(self, platform_int): + """ + Check that the correct improper terms are present in a molecule with + multiple ML and MM subregions. + """ + + topology = openmm.app.Topology() + chain = topology.addChain() + atoms = [topology.addAtom("X", openmm.app.element.carbon, topology.addResidue("X", chain)) for _ in range(8)] + topology.addBond(atoms[0], atoms[1]) + topology.addBond(atoms[1], atoms[2]) + topology.addBond(atoms[2], atoms[3]) + topology.addBond(atoms[3], atoms[4]) + topology.addBond(atoms[1], atoms[5]) + topology.addBond(atoms[2], atoms[6]) + topology.addBond(atoms[3], atoms[7]) + + mm_system = openmm.System() + torsion_force = openmm.PeriodicTorsionForce() + torsion_force.addTorsion(0, 1, 2, 5, 1, 0, 1) + torsion_force.addTorsion(1, 0, 2, 5, 1, 0, 1) + torsion_force.addTorsion(2, 0, 1, 5, 1, 0, 1) + torsion_force.addTorsion(5, 0, 1, 2, 1, 0, 1) + torsion_force.addTorsion(1, 2, 3, 6, 1, 0, 1) + torsion_force.addTorsion(2, 1, 3, 6, 1, 0, 1) + torsion_force.addTorsion(3, 1, 2, 6, 1, 0, 1) + torsion_force.addTorsion(6, 1, 2, 3, 1, 0, 1) + torsion_force.addTorsion(2, 3, 4, 7, 1, 0, 1) + torsion_force.addTorsion(3, 2, 4, 7, 1, 0, 1) + torsion_force.addTorsion(4, 2, 3, 7, 1, 0, 1) + torsion_force.addTorsion(7, 2, 3, 4, 1, 0, 1) + mm_system.addForce(torsion_force) + + """ + An improper is added for each of the three improper centers and with + the central atom as each of the four possible atoms. + + MM5 ML7 + | | + MM0 - MM1 - ML2 - ML3 - MM4 + | + MM6 + """ + mixed_system = MLPotential("mace-off23-small").createMixedSystem(topology, mm_system, [2, 3, 7]) + _, _, mixed_torsions, _ = self.getBondedTerms(mixed_system) + + assert mixed_torsions == {(0, 1, 2, 5), (1, 0, 2, 5), (2, 0, 1, 5), (2, 1, 0, 5)} + + def testLinkAtomInterpolation(self, platform_int): + """ + Ensure interpolation works as expected with the link-atom method. + """ + + pdb = openmm.app.PDBFile(os.path.join(test_data_dir, "ethanol", "ethanol.pdb")) + + mm_force_field = openmm.app.ForceField(os.path.join(test_data_dir, "ethanol", "ethanol.xml")) + ml_potential = MLPotential("mace-off23-small") + + mm_system = mm_force_field.createSystem(pdb.topology) + mixed_system = ml_potential.createMixedSystem(pdb.topology, mm_system, [0, 1, 3, 4, 5], interpolate=False) + interpolate_system = ml_potential.createMixedSystem(pdb.topology, mm_system, [0, 1, 3, 4, 5], interpolate=True) + + platform = openmm.Platform.getPlatform(platform_int) + mm_context = openmm.Context(mm_system, openmm.VerletIntegrator(0.001), platform) + mixed_context = openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) + interpolate_context = openmm.Context(interpolate_system, openmm.VerletIntegrator(0.001), platform) + + mm_context.setPositions(pdb.positions) + for context in (mixed_context, interpolate_context): + context.setPositions(pdb.positions + [openmm.Vec3(0, 0, 0)] * openmm.unit.nanometer) + context.computeVirtualSites() + + mm_energy = mm_context.getState(energy=True).getPotentialEnergy().value_in_unit(openmm.unit.kilojoule_per_mole) + mixed_energy = mixed_context.getState(energy=True).getPotentialEnergy().value_in_unit(openmm.unit.kilojoule_per_mole) + + for lambda_value in (0.0, 0.25, 0.5, 0.75, 1.0): + interpolate_context.setParameter("lambda_interpolate", lambda_value) + interpolate_energy = interpolate_context.getState(energy=True).getPotentialEnergy().value_in_unit(openmm.unit.kilojoule_per_mole) + assert np.isclose(interpolate_energy, mixed_energy * lambda_value + mm_energy * (1 - lambda_value), rtol=0, atol=atol) + + def testLinkAtomInfo(self, platform_int): + """ + Ensure the returnInfo keyword works with the link-atom method. + """ + + pdb = openmm.app.PDBFile(os.path.join(test_data_dir, "ethanol", "ethanol.pdb")) + mm_force_field = openmm.app.ForceField(os.path.join(test_data_dir, "ethanol", "ethanol.xml")) + ml_potential = MLPotential("mace-off23-small") + mm_system = mm_force_field.createSystem(pdb.topology) + + original_count = mm_system.getNumParticles() + mixed_system = ml_potential.createMixedSystem(pdb.topology, mm_system, [0, 1, 3, 4, 5], returnInfo=False) + mixed_info = ml_potential.createMixedSystem(pdb.topology, mm_system, [0, 1, 3, 4, 5], returnInfo=True) + + assert isinstance(mixed_system, openmm.System) + assert isinstance(mixed_info["system"], openmm.System) + assert isinstance(mixed_info["topology"], openmm.app.Topology) + + # Make sure the inputs were not modified. + assert mm_system.getNumParticles() == pdb.topology.getNumAtoms() == original_count + # Make sure the outputs have been modified and match. + assert mixed_system.getNumParticles() == mixed_info["system"].getNumParticles() == mixed_info["topology"].getNumAtoms() > original_count + # Make sure the virtual sites were appended to the end. + assert mixed_info["oldToNew"] == list(range(original_count)) + for i in range(mixed_system.getNumParticles()): + assert mixed_system.isVirtualSite(i) == (i >= original_count) diff --git a/test/data/ethanol/ethanol.pdb b/test/data/ethanol/ethanol.pdb new file mode 100644 index 0000000..8f79ea8 --- /dev/null +++ b/test/data/ethanol/ethanol.pdb @@ -0,0 +1,21 @@ +REMARK 1 CREATED WITH OPENMM 8.5.2, 2026-07-30 +HETATM 1 O1x UNK A 1 1.412 0.565 -0.264 1.00 0.00 O +HETATM 2 C1x UNK A 1 0.459 -0.319 0.200 1.00 0.00 C +HETATM 3 C2x UNK A 1 -0.959 0.211 0.058 1.00 0.00 C +HETATM 4 H1x UNK A 1 1.732 0.251 -1.133 1.00 0.00 H +HETATM 5 H2x UNK A 1 0.660 -0.459 1.292 1.00 0.00 H +HETATM 6 H3x UNK A 1 0.568 -1.306 -0.296 1.00 0.00 H +HETATM 7 H4x UNK A 1 -1.090 0.835 -0.846 1.00 0.00 H +HETATM 8 H5x UNK A 1 -1.627 -0.670 0.064 1.00 0.00 H +HETATM 9 H6x UNK A 1 -1.154 0.890 0.926 1.00 0.00 H +TER 10 UNK A 1 +CONECT 1 2 4 +CONECT 2 1 3 5 6 +CONECT 3 2 7 8 9 +CONECT 4 1 +CONECT 5 2 +CONECT 6 2 +CONECT 7 3 +CONECT 8 3 +CONECT 9 3 +END diff --git a/test/data/ethanol/ethanol.xml b/test/data/ethanol/ethanol.xml new file mode 100644 index 0000000..5cf4a6c --- /dev/null +++ b/test/data/ethanol/ethanol.xml @@ -0,0 +1,85 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/data/ethanol/ethanol_ljforce.xml b/test/data/ethanol/ethanol_ljforce.xml new file mode 100644 index 0000000..69a69df --- /dev/null +++ b/test/data/ethanol/ethanol_ljforce.xml @@ -0,0 +1,96 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +