From ee2dde4973ea7e9bb714403eff3faf3a2b3de6eb Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 12:08:17 +0000 Subject: [PATCH 01/20] Sync main with upstream OpenMM-ML --- README.md | 5 +- devtools/requirements/mace-aimnet-torchmd.txt | 2 + doc/userguide.md | 39 ++- openmmml/embeddings/mechanicalembedding.py | 32 +- openmmml/embeddings/utilities.py | 217 ++++++++++++- openmmml/mlpotential.py | 51 +++- openmmml/models/macepotential.py | 47 +-- setup.py | 4 + test/TestMACEPotential.py | 55 +++- test/TestMechanicalEmbedding.py | 287 +++++++++++++++++- test/data/ethanol/ethanol.pdb | 21 ++ test/data/ethanol/ethanol.xml | 85 ++++++ test/data/ethanol/ethanol_ljforce.xml | 96 ++++++ test/data/mace_energies.py | 18 +- 14 files changed, 893 insertions(+), 66 deletions(-) create mode 100644 test/data/ethanol/ethanol.pdb create mode 100644 test/data/ethanol/ethanol.xml create mode 100644 test/data/ethanol/ethanol_ljforce.xml diff --git a/README.md b/README.md index 32e98c1..de7c54c 100644 --- a/README.md +++ b/README.md @@ -8,8 +8,9 @@ can set up a simulation that uses a standard, pretrained model to represent some The current release of OpenMM-ML supports the following potential functions. - [MACE](https://arxiv.org/abs/2206.07697) models, including the pre-trained [MACE-OFF23](https://arxiv.org/abs/2312.15211), - [MACE-MPA-0](https://github.com/ACEsuit/mace-foundations), [MACE-OMAT-0](https://github.com/ACEsuit/mace-foundations), - and MACE-OMOL-0 models, utilizing the [MACE implementation](https://github.com/ACEsuit/mace). + [MACE-MPA-0](https://github.com/ACEsuit/mace-foundations), [MACE-OMAT-0](https://github.com/ACEsuit/mace-foundations), MACE-OMOL-0, + [MACELES-OFF](https://github.com/ChengUCB/les_fit), and [MACE-POLAR-1](https://mace-docs.readthedocs.io/en/latest/guide/polar_mace.html) + models, utilizing the [MACE implementation](https://github.com/ACEsuit/mace). - The pretrained [AIMNet2](https://doi.org/10.1039/D4SC08572H) model. diff --git a/devtools/requirements/mace-aimnet-torchmd.txt b/devtools/requirements/mace-aimnet-torchmd.txt index 86e885f..0bcef3b 100644 --- a/devtools/requirements/mace-aimnet-torchmd.txt +++ b/devtools/requirements/mace-aimnet-torchmd.txt @@ -5,3 +5,5 @@ torchmd-net ase huggingface-hub torchani +git+https://github.com/ChengUCB/les.git@v0.2.0 +git+https://github.com/WillBaldwin0/graph_electrostatics.git@v0.4.0 diff --git a/doc/userguide.md b/doc/userguide.md index 826034d..2c579b5 100644 --- a/doc/userguide.md +++ b/doc/userguide.md @@ -89,6 +89,8 @@ are supported. | `mace-mpa-0-medium` | Pretrained [MACE-MPA-0](https://github.com/ACEsuit/mace-foundations) model | | `mace-omat-0-small`
`mace-omat-0-medium` | Pretrained [MACE-OMAT-0](https://github.com/ACEsuit/mace-foundations) models | | `mace-omol-0-extra-large` | Pretrained MACE-OMOL-0 model | +| `mace-les-off-small` | Pretrained [MACELES-OFF](https://github.com/ChengUCB/les_fit) model (requires [LES](https://github.com/ChengUCB/les) plugin) | +| `mace-polar-1-small`
`mace-polar-1-medium`
`mace-polar-1-large` | Pretrained [MACE-POLAR-1](https://mace-docs.readthedocs.io/en/latest/guide/polar_mace.html) models (requires [`graph_longrange`](https://github.com/WillBaldwin0/graph_electrostatics)) | | `mace` | Custom MACE models specified with the `modelPath` argument | When creating MACE models, the following keyword arguments to the `MLPotential` constructor are supported. @@ -307,14 +309,19 @@ When using Orb models, the following extra keyword arguments to `createSystem()` OpenMM-ML can use an arbitrary [ASE](https://ase-lib.org/) Calculator to perform calculations. This allows it to use any model or code for which a Calculator is available, including a wide variety of MLIPs and quantum chemistry programs. -Simply pass the [Calculator](https://ase-lib.org/ase/calculators/calculators.html) to `createSystem()`: +Simply pass the [Calculator](https://docs.ase-lib.org/ase/calculators/calculators.html) to `createSystem()`. For example, +the following uses [fairchem](https://github.com/facebookresearch/fairchem) to simulate a system with the UMA-s-1.2.1 +model. ```python +from fairchem.core import pretrained_mlip, FAIRChemCalculator +predictor = pretrained_mlip.get_predict_unit("uma-s-1p2p1", device="cuda") +calculator = FAIRChemCalculator(predictor, task_name="omol") potential = MLPotential('ase') system = potential.createSystem(topology, calculator=calculator) ``` -An ASE [Atoms](https://ase-lib.org/ase/atoms.html) object is created automatically based on the OpenMM Topology. You +An ASE [Atoms](https://docs.ase-lib.org/ase/atoms.html) object is created automatically based on the OpenMM Topology. You can optionally provide values to add to its `info` dict. Some Calculators use this as a mechanism to specify parameters like total charge and spin multiplicity: @@ -374,6 +381,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/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index e9efea6..829141b 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -57,7 +57,8 @@ class MACEPotentialImpl(MLPotentialImpl): >>> potential = MLPotential('mace-off23-small') Other available models include 'mace-off23-medium', 'mace-off23-large', 'mace-off24-medium', - 'mace-mpa-0-medium', 'mace-omat-0-small', 'mace-omat-0-medium', and 'mace-omol-0-extra-large'. + 'mace-mpa-0-medium', 'mace-omat-0-small', 'mace-omat-0-medium', 'mace-omol-0-extra-large', + 'mace-les-off-small', 'mace-polar-1-small', 'mace-polar-1-medium', and 'mace-polar-1-large'. To use a locally trained MACE model, provide the path to the model file. For example: @@ -89,16 +90,20 @@ class MACEPotentialImpl(MLPotentialImpl): The path to the locally trained MACE model if ``name`` is 'mace'. """ - # (Function name, model name, restrictive license, long-range) + # (Function name, model name, restrictive license name or None, long-range) KNOWN_MODELS = { - 'mace-off23-small': ('mace_off', 'small', True, False), - 'mace-off23-medium': ('mace_off', 'medium', True, False), - 'mace-off23-large': ('mace_off', 'large', True, False), - 'mace-off24-medium': ('mace_off', 'https://github.com/ACEsuit/mace-off/blob/main/mace_off24/MACE-OFF24_medium.model?raw=true', True, False), - 'mace-mpa-0-medium': ('mace_mp', 'medium-mpa-0', False, False), - 'mace-omat-0-small': ('mace_mp', 'small-omat-0', True, False), - 'mace-omat-0-medium': ('mace_mp', 'medium-omat-0', True, False), - 'mace-omol-0-extra-large': ('mace_omol', 'extra_large', True, False), + 'mace-off23-small': ('mace_off', 'small', 'ASL', False), + 'mace-off23-medium': ('mace_off', 'medium', 'ASL', False), + 'mace-off23-large': ('mace_off', 'large', 'ASL', False), + 'mace-off24-medium': ('mace_off', 'https://github.com/ACEsuit/mace-off/blob/main/mace_off24/MACE-OFF24_medium.model?raw=true', 'ASL', False), + 'mace-mpa-0-medium': ('mace_mp', 'medium-mpa-0', None, False), + 'mace-omat-0-small': ('mace_mp', 'small-omat-0', 'ASL', False), + 'mace-omat-0-medium': ('mace_mp', 'medium-omat-0', 'ASL', False), + 'mace-omol-0-extra-large': ('mace_omol', 'extra_large', 'ASL', False), + 'mace-les-off-small': ('mace_off', 'https://github.com/ChengUCB/les_fit/blob/main/MACELES-OFF/MACELES-OFF_small_converted.model?raw=true', 'CC BY-NC 4.0', True), + 'mace-polar-1-small': ('mace_polar', 'polar-1-s', None, True), + 'mace-polar-1-medium': ('mace_polar', 'polar-1-m', None, True), + 'mace-polar-1-large': ('mace_polar', 'polar-1-l', None, True), } def __init__(self, name: str, modelPath) -> None: @@ -111,7 +116,8 @@ def __init__(self, name: str, modelPath) -> None: The name of the MACE model. Options include 'mace-off23-small', 'mace-off23-medium', 'mace-off23-large', 'mace-off24-medium', 'mace-mpa-0-medium', 'mace-omat-0-small', 'mace-omat-0-medium', - 'mace-omol-0-extra-large', and 'mace'. + 'mace-omol-0-extra-large', 'mace-les-off-small', 'mace-polar-1-small', + 'mace-polar-1-medium', 'mace-polar-1-large', and 'mace'. modelPath : str, optional The path to the locally trained MACE model if ``name`` is 'mace'. """ @@ -151,7 +157,7 @@ def addForces( import torch try: from mace.tools import utils, to_one_hot, atomic_numbers_to_indices - from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol + from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol, mace_polar except ImportError as e: raise ImportError(f"Failed to import mace with error: {e}. Install mace with 'pip install mace-torch'.") try: @@ -169,12 +175,13 @@ def addForces( 'mace_off': mace_off, 'mace_mp': mace_mp, 'mace_omol': mace_omol, + 'mace_polar': mace_polar, } - fnName, name, warn, _ = MACEPotentialImpl.KNOWN_MODELS[self.name] + fnName, name, restrictiveLicense, _ = MACEPotentialImpl.KNOWN_MODELS[self.name] model = functions[fnName](model=name, device=device, return_raw_model=True).to(device) - if warn: + if restrictiveLicense is not None: import logging - logging.warning(f'The model {self.name} is distributed under the restrictive ASL license. Commercial use is not permitted.') + logging.warning(f'The model {self.name} is distributed under the restrictive {restrictiveLicense} license. Commercial use is not permitted.') elif self.name == "mace": if self.modelPath is not None: model = torch.load(self.modelPath, map_location=device).to(device) @@ -203,6 +210,7 @@ def addForces( raise ValueError(f"Unsupported precision {precision} for the model. Supported values are 'single' and 'double'.") if dtype != modelDefaultDtype: print(f"Model dtype is {modelDefaultDtype} and requested dtype is {dtype}. The model will be converted to the requested dtype.") + model = model.to(dtype) # One hot encoding of atomic numbers @@ -258,6 +266,7 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch dtype = node_attrs.dtype cutoff = float(model.r_max.detach()) edgeIndex, shifts, _, _ = get_neighborhood(positions, cutoff, [periodic, periodic, periodic], cell) + cellTensor = torch.tensor(cell, dtype=dtype, device=ptr.device) inputDict = { "ptr": ptr, "node_attrs": node_attrs, @@ -266,9 +275,13 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch "positions": torch.tensor(positions, dtype=dtype, device=ptr.device), "edge_index": torch.tensor(edgeIndex, dtype=torch.int64, device=ptr.device), "shifts": torch.tensor(shifts, dtype=dtype, device=ptr.device), - "cell": torch.tensor(cell, dtype=dtype, device=ptr.device), + "cell": cellTensor, + "rcell": 2 * torch.pi * torch.linalg.inv(cellTensor.mT), + "volume": torch.linalg.det(cellTensor), "total_charge": charge, - "total_spin": multiplicity + "total_spin": multiplicity, + "external_field": torch.zeros((1, 3), dtype=dtype, device=ptr.device), + "fermi_level": torch.zeros((1,), dtype=dtype, device=ptr.device) } results = model(inputDict, compute_force=True) energy = float(results[returnEnergyType].detach())*energyScale diff --git a/setup.py b/setup.py index f939e3f..57b776b 100644 --- a/setup.py +++ b/setup.py @@ -69,6 +69,10 @@ 'mace-omat-0-small = openmmml.models.macepotential:MACEPotentialImplFactory', 'mace-omat-0-medium = openmmml.models.macepotential:MACEPotentialImplFactory', 'mace-omol-0-extra-large = openmmml.models.macepotential:MACEPotentialImplFactory', + 'mace-les-off-small = openmmml.models.macepotential:MACEPotentialImplFactory', + 'mace-polar-1-small = openmmml.models.macepotential:MACEPotentialImplFactory', + 'mace-polar-1-medium = openmmml.models.macepotential:MACEPotentialImplFactory', + 'mace-polar-1-large = openmmml.models.macepotential:MACEPotentialImplFactory', 'nequip = openmmml.models.nequippotential:NequIPPotentialImplFactory', 'orb-v3-conservative-inf-omat = openmmml.models.orbpotential:OrbPotentialImplFactory', 'orb-v3-conservative-omol = openmmml.models.orbpotential:OrbPotentialImplFactory', diff --git a/test/TestMACEPotential.py b/test/TestMACEPotential.py index 568d9f4..4470278 100644 --- a/test/TestMACEPotential.py +++ b/test/TestMACEPotential.py @@ -15,9 +15,31 @@ @pytest.mark.parametrize("platform_int", list(platform_ints)) class TestMACE: + + # Reference energies are calculated with MACECalculator + refEnergy = { + 'mace-off23-small': -713468.6327560507, + 'mace-off23-medium': -713468.0563706581, + 'mace-off23-large': -713467.7476380612, + 'mace-off24-medium': -713467.9394350434, + 'mace-mpa-0-medium': -8839.299589829867, + 'mace-omat-0-small': -8726.63865431241, + 'mace-omat-0-medium': -8679.026847088873, + 'mace-omol-0-extra-large': -712903.4934289698, + 'mace-les-off-small': -713467.9354591698, + 'mace-polar-1-small': -712903.1710073923, + 'mace-polar-1-medium': -712903.4536792638, + 'mace-polar-1-large': -712903.7834631138, + } + @pytest.mark.parametrize("model", ['mace-off23-small', 'mace-off23-medium', 'mace-off23-large', 'mace-off24-medium', - 'mace-mpa-0-medium', 'mace-omat-0-small', 'mace-omat-0-medium', 'mace-omol-0-extra-large']) + 'mace-mpa-0-medium', 'mace-omat-0-small', 'mace-omat-0-medium', 'mace-omol-0-extra-large', + 'mace-les-off-small', 'mace-polar-1-small', 'mace-polar-1-medium', 'mace-polar-1-large']) def testCreatePureMLSystem(self, platform_int, model): + if 'mace-les' in model: + pytest.importorskip("les", reason="les is not installed") + if 'mace-polar' in model: + pytest.importorskip("graph_longrange", reason="graph_electrostatics is not installed") pdb = app.PDBFile(os.path.join(test_data_dir, "toluene", "toluene.pdb")) potential = MLPotential(model) system = potential.createSystem(pdb.topology, returnEnergyType='energy') @@ -25,16 +47,7 @@ def testCreatePureMLSystem(self, platform_int, model): context = mm.Context(system, mm.VerletIntegrator(0.001), platform) context.setPositions(pdb.getPositions(asNumpy=True)) energyML = context.getState(energy=True).getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole) - # Reference energies are calculated with MACECalculator - refEnergy = {'mace-off23-small': -713468.6327560507, - 'mace-off23-medium': -713468.0563706581, - 'mace-off23-large': -713467.7476380612, - 'mace-off24-medium': -713467.9394350434, - 'mace-mpa-0-medium': -8839.299589829867, - 'mace-omat-0-small': -8726.63865431241, - 'mace-omat-0-medium': -8679.026847088873, - 'mace-omol-0-extra-large': -712903.4934289698} - assert np.isclose(refEnergy[model], energyML, rtol=1e-6) + assert np.isclose(self.refEnergy[model], energyML, rtol=1e-6) def testPeriodicSystem(self, platform_int): pdb = app.PDBFile(os.path.join(test_data_dir, "alanine-dipeptide", "alanine-dipeptide-explicit.pdb")) @@ -72,3 +85,23 @@ def testCreateMixedSystem(self, platform_int): interpEnergy2 = interpContext.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole) assert np.isclose(mixedEnergy, interpEnergy1, rtol=1e-5) assert np.isclose(mmEnergy, interpEnergy2, rtol=1e-5) + + @pytest.mark.parametrize("precision", ["single", "double"]) + def testPrecisionApplied(self, platform_int, precision): + pdb = app.PDBFile(os.path.join(test_data_dir, "toluene", "toluene.pdb")) + potential = MLPotential('mace-off23-small') + + # Specifying precision single/double. + system = potential.createSystem(pdb.topology, returnEnergyType='energy', precision=precision) + platform = mm.Platform.getPlatform(platform_int) + context = mm.Context(system, mm.VerletIntegrator(0.001), platform) + context.setPositions(pdb.positions) + + # Inconsistent dtypes will crash the simulation. + energyML = context.getState(energy=True, forces=True).getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole) + + # The energy should be physically meaningful under both precisions + assert np.isfinite(energyML), \ + "Energy is not finite under precision {}".format(precision) + assert np.isclose(energyML, self.refEnergy['mace-off23-small'], rtol=1e-6),\ + "Energy is not close to reference under precision {}".format(precision) 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 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/data/mace_energies.py b/test/data/mace_energies.py index 4febeb2..72ee14f 100644 --- a/test/data/mace_energies.py +++ b/test/data/mace_energies.py @@ -1,7 +1,7 @@ # This script computes reference energies for the MACE foundation models. import ase.io -from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol +from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol, mace_polar from openmm.unit import kilojoules_per_mole, ev, item atoms = ase.io.read('toluene/toluene.pdb') @@ -22,6 +22,22 @@ results['mace-omat-0-medium'] = atoms.get_potential_energy() atoms.calc = mace_omol('extra_large') results['mace-omol-0-extra-large'] = atoms.get_potential_energy() +try: + import les + atoms.calc = mace_off('https://github.com/ChengUCB/les_fit/blob/main/MACELES-OFF/MACELES-OFF_small_converted.model?raw=true') + results['mace-les-off-small'] = atoms.get_potential_energy() +except ImportError: + print('No MACELES reference energies will be generated; you must first install LES from https://github.com/ChengUCB/les') +try: + import graph_longrange + atoms.calc = mace_polar('polar-1-s') + results['mace-polar-1-small'] = atoms.get_potential_energy() + atoms.calc = mace_polar('polar-1-m') + results['mace-polar-1-medium'] = atoms.get_potential_energy() + atoms.calc = mace_polar('polar-1-l') + results['mace-polar-1-large'] = atoms.get_potential_energy() +except ImportError: + print('No MACE-POLAR reference energies will be generated; you must first install graph_longrange from https://github.com/WillBaldwin0/graph_electrostatics') atoms = ase.io.read('alanine-dipeptide/alanine-dipeptide-explicit.pdb') atoms.calc = mace_off('small') results['alanine-dipeptide'] = atoms.get_potential_energy() From 94f7fb116949ddc0706be568e8c0c574ed8c6ce5 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 12:10:05 +0000 Subject: [PATCH 02/20] Add graph-longrange PolarMACE electrostatic embedding --- doc/userguide.md | 59 +- openmmml/models/_links.py | 331 +++++++++ openmmml/models/_polarmace_external.py | 348 +++++++++ openmmml/models/macepotential.py | 739 ++++++++++++++++++-- test/TestElectrostaticEmbeddingPME.py | 506 ++++++++++++++ test/TestLinkChargeRedistribution.py | 222 ++++++ test/TestLinkRecordsTargetDistUnits.py | 147 ++++ test/TestMACEPotential.py | 177 +++++ test/TestMechanicalEmbedding.py | 37 + test/TestPolarMACEElectrostaticEmbedding.py | 658 +++++++++++++++++ 10 files changed, 3153 insertions(+), 71 deletions(-) create mode 100644 openmmml/models/_links.py create mode 100644 openmmml/models/_polarmace_external.py create mode 100644 test/TestElectrostaticEmbeddingPME.py create mode 100644 test/TestLinkChargeRedistribution.py create mode 100644 test/TestLinkRecordsTargetDistUnits.py create mode 100644 test/TestPolarMACEElectrostaticEmbedding.py diff --git a/doc/userguide.md b/doc/userguide.md index 2c579b5..29efd6f 100644 --- a/doc/userguide.md +++ b/doc/userguide.md @@ -104,7 +104,10 @@ When using MACE models, the following extra keyword arguments to `createSystem() | Argument | Description | | --- | --- | | `precision` | The numerical precision of the model. Supported options are `'single'` and `'double'`. If `None`, the default precision of the model is used. | -| `returnEnergyType` | Whether to return the interaction energy or the energy including the self-energy. The default is `'interaction_energy'`. Supported options are `'interaction_energy'` and `'energy'`. | +| `returnEnergyType` | Whether to return the interaction energy or the energy including the self-energy. The default is `'energy'`, which is the scalar the reported forces are the gradient of, so the potential is conservative. Supported options are `'interaction_energy'` and `'energy'`. | +| `linkRecords` | Hydrogen link atoms to cap covalent bonds crossing the ML/MM boundary; see *Link atoms* below. Only meaningful for `createMixedSystem()`. | +| `linkChargeScheme` | How to treat the MM charges on the MM-side boundary atoms when `linkRecords` is given; see *Link atoms* below. The default is `'dz1'`. | +| `embedding` | Passed by the embedding method rather than by you; see *Embeddings* below. | | `device` | The PyTorch device to perform calculations on, either a `torch.device` object or a string (such as `'cuda'` or `'cpu'`.) If omitted, a device is chosen automatically. | | `charge` | The total charge of the system. If omitted, it is assumed to be 0. This is only used by MACE-OMOL-0. For other models it is ignored. | | `multiplicity` | The spin multiplicity of the system. If omitted, it is assumed to be 1. This is only used by MACE-OMOL-0. For other models it is ignored. | @@ -381,7 +384,59 @@ 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 +### Electrostatic Embedding + +This is a potential-specific embedding method provided by the MACE interface rather than a generic one, and it is +selected with the embedding name `electrostatic`. The MLIP, rather than the MM force field, computes the electrostatic +interactions between the ML and MM atoms: it is given the positions and MM force field charges of the MM atoms, and +returns forces on them alongside the forces on the ML atoms. The ML subset can therefore polarise in response to its +surroundings, which mechanical embedding does not allow. Lennard-Jones interactions between the ML and MM atoms are +still computed by the MM force field, as are all bonded terms that cross the ML/MM boundary. + +This is implemented by setting the MM force field charge of every ML atom to zero, so that every Coulomb term involving +an ML atom vanishes, including the reciprocal space part of PME. As a result, the MM force field's own charges are +untouched and MM-MM electrostatics are unchanged. + +Only PolarMACE models can be used with this embedding method. This includes the pretrained +`mace-polar-1-small`, `mace-polar-1-medium`, and `mace-polar-1-large` models, as well as a PolarMACE checkpoint loaded +with the model name `mace` and a `modelPath`. For a custom checkpoint the method is listed because its type cannot be +known until it is loaded; a non-PolarMACE checkpoint is rejected at that point rather than silently falling back to +mechanical embedding, since the mixed system has already had its ML-MM electrostatics removed. + +The damped real-space and reciprocal-space ML-MM cross energy, including slab and molecule-in-box corrections, is +implemented by `GTOElectrostaticCrossEnergy` in `graph_longrange`. OpenMM-ML only adapts the external positions and +charges to the stock PolarMACE electrostatic feature and energy blocks, then returns the corresponding MM forces. The +normalization is therefore inherited from the checkpoint's own GTO density and feature bases rather than duplicated in +the OpenMM interface. +Interpolation is not supported, because at `lambda_interpolate=0` the ML-MM electrostatics would be missing from the MM +endpoint. + +```python +system = potential.createMixedSystem(topology, mm_system, ml_atoms, embedding='electrostatic') +``` + +#### Link atoms + +When the ML subset is carved out of the middle of a molecule, the covalent bonds crossing the boundary leave the ML +atoms with dangling valences. The MACE interface can cap each such bond with a hydrogen link atom, placed along the +bond vector at every force evaluation and with its forces redistributed onto the two real atoms it sits between, so the +capped system stays conservative. Caps are requested with the `linkRecords` argument to `createMixedSystem()`, which +takes a sequence of `(q, m, target_dist)` tuples — `q` is the ML-side atom index, `m` the MM-side atom index, and +`target_dist` the desired cap distance from `q` in Ångströms — or the path to a CSV file with `q_idx1`, `m_idx1`, and +`target_dist_ang` columns. Note that the tuples are 0-based, matching OpenMM's atom indices, while the CSV columns are +1-based, hence their names. + +The MM partial charge on each MM-side boundary atom would otherwise sit about 1.5 Å from the nearest ML atom, through +the cap, and over-polarise it. The `linkChargeScheme` argument controls how that charge is handled, affecting only the +charges passed to the MLIP; the MM `NonbondedForce` is left alone, so MM-MM electrostatics remain exact. + +| `linkChargeScheme` | Description | +| --- | --- | +| `'dz1'` | The default. Sets the MM-side boundary atom charge to zero and spreads it evenly over that atom's MM neighbours, conserving total charge. Falls back to `'z1'`, with a warning, for a boundary atom with no MM neighbours. | +| `'z1'` | Sets the MM-side boundary atom charge to zero without redistributing it, which changes the total charge. | +| `'none'` | Leaves the MM charges untouched. | + +### 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, diff --git a/openmmml/models/_links.py b/openmmml/models/_links.py new file mode 100644 index 0000000..25a93d1 --- /dev/null +++ b/openmmml/models/_links.py @@ -0,0 +1,331 @@ +"""Hydrogen link-atom (cap) primitives shared across the codebase. + +Caps are virtual hydrogen atoms placed on the QM/ML side of a Q-M +boundary bond. Their positions are deterministic functions of the +Q (ML-side) and M (MM-side) atom positions; their forces from the +ML evaluator must be redistributed onto Q and M before being handed +back to OpenMM. + +The caller is `openmmml/models/macepotential.py::_computeMACE`, which +places caps in MACE's input each step and redistributes forces from +MACE's output. The math lives here rather than there so that any +further evaluator needing caps uses the *same* placement formula and +the *same* redistribution Jacobian: + + r_cap = (1 - C_L) * r_Q + C_L * r_M where C_L = target_dist / |r_M - r_Q| + + F_Q ← F_Q + (1 - C_L) F_cap + C_L (F_cap · ê_b) ê_b + F_M ← F_M + C_L F_cap - C_L (F_cap · ê_b) ê_b + where ê_b = (r_M - r_Q) / |r_M - r_Q| + +The redistribution preserves total force and total torque (the cap is +constrained to lie along the Q-M axis at fixed fraction `C_L`, so +its motion is entirely determined by Q and M motion; the chain rule +of energy w.r.t. Q,M gives the formulas above). +""" +from __future__ import annotations + +import numpy as np + + +# All 27 neighbouring lattice-cell offsets (coefficients in {-1, 0, 1}^3). +_OFFSETS_27 = np.array( + [[i, j, k] for i in (-1, 0, 1) for j in (-1, 0, 1) for k in (-1, 0, 1)], + dtype=np.float64, +) + + +def minimum_image_M( + r_Q: np.ndarray, + r_M: np.ndarray, + cell: np.ndarray, + max_bond: float = 2.5, +) -> np.ndarray: + """Return each M atom shifted to its minimum image relative to its Q partner. + + When a Q-M boundary bond straddles a periodic boundary the raw M position is + in a different image than Q, and both cap *placement* and cap-force + *redistribution* must use the same imaged M (otherwise the redistributed + force is not the gradient of the energy the cap produced). Both call sites in + ``macepotential._computeMACE`` therefore route ``r_M`` through this helper. + + The nearest image is found by a 27-cell search: round the fractional + displacement to the nearest lattice cell, then pick the shortest candidate + over that cell and its 26 neighbours. This is exact for a real (weakly skewed) + simulation box; it is *not* a certified closest-vector solver for + pathologically skewed triclinic cells, which do not occur for physical MD + systems. As a sanity check, the imaged Q-M distance must stay below + ``max_bond`` (Angstrom); a longer bond means a genuinely wrapped / broken pair + (or a box thinner than the bond) and raises rather than placing a bad cap. + + Parameters + ---------- + r_Q, r_M : arrays of shape (K, 3) + Q (ML-side) and M (MM-side) positions, one row per cap, in Angstrom + (the unit ``cell`` is given in). + cell : array of shape (3, 3) + Periodic box vectors as rows (OpenMM convention), in Angstrom. + max_bond : float, optional + Chemical sanity ceiling on the imaged Q-M bond length (Angstrom). Frontier + bonds are ~1.0-1.6 A; the default 2.5 A leaves margin while still catching + a wrongly-imaged / wrapped pair. + + Returns + ------- + r_M_imaged : array of shape (K, 3) + M positions shifted to the image nearest their Q partner. + """ + r_Q = np.atleast_2d(np.asarray(r_Q, dtype=np.float64)) + r_M = np.atleast_2d(np.asarray(r_M, dtype=np.float64)) + cell = np.asarray(cell, dtype=np.float64) + + if cell.shape != (3, 3): + raise ValueError("cell must have shape (3, 3)") + if abs(np.linalg.det(cell)) < 1e-12: + raise ValueError("cell is singular or nearly singular") + if r_Q.shape != r_M.shape: + raise ValueError("r_Q and r_M must have the same shape") + + dr = r_M - r_Q + base = np.round(dr @ np.linalg.inv(cell)) # nearest lattice cell (orthorhombic guess) + shifts = base[:, None, :] + _OFFSETS_27[None, :, :] # (K, 27, 3) integer coefficients + cand = dr[:, None, :] - shifts @ cell # (K, 27, 3) candidate displacements + best = np.argmin(np.einsum("kij,kij->ki", cand, cand), axis=1) + dr_mi = cand[np.arange(cand.shape[0]), best] + + bond = np.linalg.norm(dr_mi, axis=-1) + if np.any(bond >= max_bond): + raise ValueError( + f"imaged Q-M link bond length {bond.max():.2f} A exceeds max_bond " + f"{max_bond:.2f} A; the pair is genuinely wrapped/broken (or the box is " + "thinner than the bond). Increase max_bond only if this bond is real." + ) + return r_Q + dr_mi + + +def compute_cap_positions( + r_Q: np.ndarray, + r_M: np.ndarray, + target_dist: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Compute cap positions and the per-cap C_L ratio. + + Parameters + ---------- + r_Q : array of shape (K, 3) + Position of each Q atom (ML side of the boundary bond), one + row per cap, in the same length unit as `target_dist`. + r_M : array of shape (K, 3) + Position of each M atom (MM side), aligned 1:1 with `r_Q`. + target_dist : array of shape (K,) + Target distance from Q to the cap. Typically a hydrogen bond + length (e.g., 1.09 Å for C-H). + + Returns + ------- + r_cap : array of shape (K, 3) + Cap positions. + C_L : array of shape (K,) + The placement ratio used for each cap. The caller usually + also needs `C_L` for the corresponding force redistribution. + + Raises + ------ + ValueError + If any C_L falls outside (0, 1) — meaning the requested + target distance exceeds (or equals) the actual Q-M distance, + or Q and M are coincident. Both cases are user-error. + """ + r_Q = np.asarray(r_Q, dtype=np.float64) + r_M = np.asarray(r_M, dtype=np.float64) + target_dist = np.asarray(target_dist, dtype=np.float64) + v = r_M - r_Q + s = np.linalg.norm(v, axis=-1) + if np.any(s == 0.0): + raise ValueError("Cap Q and M atoms are coincident.") + C_L = target_dist / s + if np.any(~np.isfinite(C_L)) or np.any(C_L <= 0.0) or np.any(C_L >= 1.0): + raise ValueError( + f"Link-atom C_L out of (0, 1): {C_L}. " + "Check that target_dist < |r_M - r_Q| for every cap." + ) + r_cap = (1.0 - C_L)[:, None] * r_Q + C_L[:, None] * r_M + return r_cap, C_L + + +def redistribute_cap_force( + F_cap: np.ndarray, + r_Q: np.ndarray, + r_M: np.ndarray, + C_L: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """Map per-cap forces onto the corresponding Q and M atoms. + + Parameters + ---------- + F_cap : array of shape (K, 3) + Force on each cap from the ML evaluator. + r_Q, r_M : arrays of shape (K, 3) + Positions of the Q (ML side) and M (MM side) atoms. + C_L : array of shape (K,) + Placement ratio for each cap (same as returned by + `compute_cap_positions`). + + Returns + ------- + F_Q_add : array of shape (K, 3) + Force contribution to add onto each Q atom. + F_M_add : array of shape (K, 3) + Force contribution to add onto each M atom. + + Notes + ----- + The caller is responsible for accumulating the returned + contributions onto the right global atom indices (e.g., via + `f[q_global] += F_Q_add` if multiple caps share a Q or M atom + is forbidden by `_prepareLinkRecords`). + """ + F_cap = np.asarray(F_cap, dtype=np.float64) + r_Q = np.asarray(r_Q, dtype=np.float64) + r_M = np.asarray(r_M, dtype=np.float64) + C_L = np.asarray(C_L, dtype=np.float64) + v = r_M - r_Q + s = np.linalg.norm(v, axis=-1) + e_b = v / s[:, None] + proj = np.einsum("ki,ki->k", F_cap, e_b) + F_Q_add = (1.0 - C_L)[:, None] * F_cap + (C_L * proj)[:, None] * e_b + F_M_add = C_L[:, None] * F_cap - (C_L * proj)[:, None] * e_b + return F_Q_add, F_M_add + + +# ----------------------------------------------------------------------------- +# Link-atom CHARGE redistribution (Z1 / DZ1) +# +# Standard QM/MM correction. When a Q–M bond is cut and the M atom is replaced +# (on the QM side) by an H link atom, the M atom's full partial charge sits +# ~1.5 Å from the nearest QM atom. With electrostatic embedding the QM region +# sees that close-in partial charge through the link H and gets +# over-polarised. The "charge shift" / "redistributed charge" schemes move +# the M partial charge away from the boundary. +# +# Z1 — set q_M ← 0. Cheapest; breaks total MM-charge neutrality by -q_M_orig. +# DZ1 — Z1 + redistribute q_M_orig / N(M1) onto each MM neighbour (M1 atom). +# Preserves total MM charge to round-off. +# +# Both schemes only modify the *constant* mm_charges array fed into the ML +# potential's electrostatic embedding — they do not touch the OpenMM MM +# scaffold's NonbondedForce, which keeps the MM-MM Coulomb sum bit-exact +# with the original force field (matching standard QM/MM practice). +# ----------------------------------------------------------------------------- + +def apply_link_charge_redistribution( + mm_atoms: np.ndarray, + mm_charges: np.ndarray, + link_info: dict, + topology, + scheme: str, +) -> np.ndarray: + """Apply Z1 / DZ1 link-atom charge redistribution to an MM charges array. + + Parameters + ---------- + mm_atoms : np.ndarray, shape (n_mm,) + Global atom indices of the MM atoms, in the order their charges appear + in ``mm_charges``. + mm_charges : np.ndarray, shape (n_mm,) + Original MM partial charges aligned with ``mm_atoms``. + link_info : dict + Output of ``MACEPotentialImpl._prepareLinkRecords``; must contain + ``q_global`` and ``m_global`` arrays. + topology : openmm.app.Topology + The full system topology — needed to look up M's MM neighbours (M1) for + DZ1. + scheme : str + One of ``"none"``, ``"z1"``, ``"dz1"``. + + Returns + ------- + np.ndarray, shape (n_mm,) + New MM charges array (always a copy of the input; never an alias). + + Notes + ----- + For ``"dz1"``, if any M atom has zero MM neighbours the call falls back to + Z1 for that atom (charge zeroed, nothing to redistribute) and emits a + ``UserWarning``. Total MM charge will then drift by the orphan q_M_orig. + """ + scheme = scheme.lower() + if scheme not in {"none", "z1", "dz1"}: + raise ValueError( + f"Unsupported linkChargeScheme {scheme!r}. " + "Supported in this iteration: 'none', 'z1', 'dz1'." + ) + new_charges = np.asarray(mm_charges, dtype=np.float64).copy() + if scheme == "none": + return new_charges + + mm_idx_to_row = {int(g): k for k, g in enumerate(mm_atoms)} + qm_set = set(int(g) for g in link_info["q_global"]) + m_globals = [int(g) for g in link_info["m_global"]] + + # Snapshot original q_M before any zeroing so DZ1 redistributes the + # *original* value even if two M atoms happened to be the same row (they + # are validated unique upstream, but the snapshot is also free insurance). + orig_qM = {m: float(new_charges[mm_idx_to_row[m]]) for m in m_globals} + + if scheme == "z1": + for m in m_globals: + new_charges[mm_idx_to_row[m]] = 0.0 + return new_charges + + # ---- scheme == "dz1" ---- + # Build the MM-neighbour list for each M atom by walking topology bonds. + # Filter: the neighbour must (a) be in mm_atoms, (b) not be a Q atom. + # The QM filter is redundant against (a) when q_global ⊂ atoms ⊂ mm complement, + # but cheap and explicit. + # A neighbour that is itself an M atom is excluded: it is about to be zeroed + # too, so charge handed to it would simply be discarded. Without that filter + # the result depends on the order the link records happen to be listed in, + # and two M atoms bonded to each other leave one of them holding charge that + # DZ1 exists to remove. + m_set = set(m_globals) + mm_neighbors_of_M: dict[int, list[int]] = {m: [] for m in m_globals} + for bond in topology.bonds(): + a = bond.atom1.index + b = bond.atom2.index + if a in m_set and b in mm_idx_to_row and b not in qm_set and b not in m_set: + mm_neighbors_of_M[a].append(b) + if b in m_set and a in mm_idx_to_row and a not in qm_set and a not in m_set: + mm_neighbors_of_M[b].append(a) + + # Zero every M atom before redistributing any charge, so a share can never + # be written to a row that a later iteration then clears. + import warnings + for m in m_globals: + new_charges[mm_idx_to_row[m]] = 0.0 + + for m in m_globals: + m1s = mm_neighbors_of_M[m] + if not m1s: + warnings.warn( + f"DZ1: M atom {m} has no MM neighbours to redistribute " + f"q={orig_qM[m]:+.4f} e onto. Falling back to Z1 for this atom.", + stacklevel=2, + ) + continue + share = orig_qM[m] / len(m1s) + for m1 in m1s: + new_charges[mm_idx_to_row[m1]] += share + + # Conservation check: warn if numerically non-trivial drift (DZ1 should + # preserve total to f64 round-off). + delta = float(new_charges.sum() + - np.asarray(mm_charges, dtype=np.float64).sum()) + if abs(delta) > 1e-9: + import warnings as _w + _w.warn( + f"DZ1: post-redistribution total MM charge drift = {delta:+.3e} e " + "(likely due to M atoms with zero MM neighbours).", + stacklevel=2, + ) + return new_charges diff --git a/openmmml/models/_polarmace_external.py b/openmmml/models/_polarmace_external.py new file mode 100644 index 0000000..f690743 --- /dev/null +++ b/openmmml/models/_polarmace_external.py @@ -0,0 +1,348 @@ +"""Small eager-mode adapter for PolarMACE external electrostatic sources. + +The electrostatic mathematics lives in graph_longrange. This module only +bridges dynamic MM sources into an otherwise unchanged PolarMACE forward pass. +It is intended for OpenMM's PythonForce execution path, not TorchScript export. +""" + +from __future__ import annotations + + +def _floating_reference(module): + for tensor in module.buffers(): + if tensor.is_floating_point(): + return tensor + for tensor in module.parameters(): + if tensor.is_floating_point(): + return tensor + return None + + +def _rebuild_feature_block(block, pbc_handling: str = "auto"): + """Rebuild a deterministic graph block saved by an older graph release.""" + from graph_longrange.features import GTOElectrostaticFeatures + + realspace = block.realspace_features + quadrupoles = bool( + getattr( + getattr(block.non_periodic_correction_terms, "self_field", None), + "include_quadrupole_corrections", + False, + ) + ) + rebuilt = GTOElectrostaticFeatures( + density_max_l=int(realspace.density_max_l), + density_smearing_width=float(realspace.density_smearing_width), + feature_max_l=int(realspace.projection_max_l), + feature_smearing_widths=[ + float(x) for x in realspace.projection_smearing_widths + ], + include_self_interaction=bool(block.include_self_interaction), + kspace_cutoff=float(block.kspace_cutoff), + quadrupole_feature_corrections=quadrupoles, + integral_normalization=str(block.feature_basis.normalize), + pbc_handling=pbc_handling, + ) + reference = _floating_reference(block) + if reference is not None: + rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) + return rebuilt + + +def _rebuild_energy_block(block, pbc_handling: str = "auto"): + from graph_longrange.energy import GTOElectrostaticEnergy + + rebuilt = GTOElectrostaticEnergy( + density_max_l=int(block.density_max_l), + density_smearing_width=float(block.density_smearing_width), + kspace_cutoff=float(block.kspace_cutoff), + include_self_interaction=bool(block.include_self_interaction), + pbc_handling=pbc_handling, + ) + reference = _floating_reference(block) + if reference is not None: + rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) + return rebuilt + + +class _ExternalFeatureBlock: + """Mixin-like implementation installed around a graph feature module.""" + + def __init__(self, base): + import torch + + class FeatureBlock(torch.nn.Module): + def __init__(inner_self, wrapped): + super().__init__() + inner_self.base = wrapped + inner_self._external = None + + def __getattr__(inner_self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(inner_self.base, name) + + def set_external_sources(inner_self, external): + inner_self._external = external + + def precompute_geometry(inner_self, **kwargs): + base_kwargs = dict(kwargs) + base_kwargs.pop("force_pbc_evaluator", None) + cache = inner_self.base.precompute_geometry(**base_kwargs) + external = inner_self._external + if external is None: + return cache + external_cache = inner_self.base.precompute_geometry_source_target( + k_vectors=base_kwargs["k_vectors"], + k_norm2=base_kwargs["k_norm2"], + k_vector_batch=base_kwargs["k_vector_batch"], + k0_mask=base_kwargs["k0_mask"], + src_positions=external["positions"], + src_batch=external["batch"], + tgt_positions=base_kwargs["node_positions"], + tgt_batch=base_kwargs["batch"], + volume=base_kwargs["volume"], + pbc=base_kwargs["pbc"], + ) + external_field = inner_self.base.forward_dynamic_source_target( + cache=external_cache, + source_feats=external["features"], + ) + result = dict(cache) + result["_openmmml_external_field"] = external_field + return result + + def forward_dynamic(inner_self, cache, source_feats, pbc=None): + if source_feats.dim() == 3 and source_feats.shape[-2] == 1: + source_feats = source_feats.squeeze(-2) + value = inner_self.base.forward_dynamic( + cache=cache, source_feats=source_feats + ) + external_field = cache.get("_openmmml_external_field") + if external_field is not None: + # PolarMACE has two spin channels. Each channel receives half + # of the physical external potential. + value = value + 0.5 * external_field + return value + + self.module = FeatureBlock(base) + + +class _ExternalEnergyBlock: + def __init__(self, base): + import torch + try: + from graph_longrange.external_source_energy import ( + GTOElectrostaticCrossEnergy, + ) + except ImportError as exc: + raise ImportError( + "PolarMACE electrostatic embedding requires a graph_longrange " + "release that provides GTOElectrostaticCrossEnergy." + ) from exc + + class EnergyBlock(torch.nn.Module): + def __init__(inner_self, wrapped): + super().__init__() + inner_self.base = wrapped + inner_self.cross = GTOElectrostaticCrossEnergy.from_energy(wrapped) + inner_self._external = None + + def __getattr__(inner_self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(inner_self.base, name) + + def set_external_sources(inner_self, external): + inner_self._external = external + + def forward(inner_self, **kwargs): + base_kwargs = dict(kwargs) + base_kwargs.pop("force_pbc_evaluator", None) + energy = inner_self.base(**base_kwargs) + external = inner_self._external + if external is None: + return energy + cross = inner_self.cross( + k_vectors=base_kwargs["k_vectors"], + k_norm2=base_kwargs["k_norm2"], + k_vector_batch=base_kwargs["k_vector_batch"], + k0_mask=base_kwargs["k0_mask"], + source_feats=base_kwargs["source_feats"], + source_positions=base_kwargs["node_positions"], + source_batch=base_kwargs["batch"], + target_feats=external["features"], + target_positions=external["positions"], + target_batch=external["batch"], + volume=base_kwargs["volume"], + pbc=base_kwargs["pbc"], + ) + return energy + cross + + self.module = EnergyBlock(base) + + +def _prepare_external_sources(model, data, compute_force: bool): + import torch + + positions = data.get("mm_positions") + charges = data.get("mm_charges") + multipoles = data.get("mm_multipoles") + if charges is not None and multipoles is not None: + raise ValueError("mm_charges and mm_multipoles are mutually exclusive.") + values = multipoles if multipoles is not None else charges + if positions is None or values is None or positions.numel() == 0 or values.numel() == 0: + return None + + ml_positions = data["positions"] + positions = positions.to(device=ml_positions.device, dtype=ml_positions.dtype) + positions = positions.clone().requires_grad_(compute_force) + width = (int(model.atomic_multipoles_max_l) + 1) ** 2 + if multipoles is None: + features = torch.zeros( + (charges.numel(), width), dtype=ml_positions.dtype, device=ml_positions.device + ) + features[:, 0] = charges.to(features).reshape(-1) + else: + features = multipoles.to(device=ml_positions.device, dtype=ml_positions.dtype).clone() + if features.dim() != 2 or features.shape != (positions.shape[0], width): + raise ValueError(f"mm_multipoles must have shape [N_mm, {width}].") + if width >= 4: + # Public Cartesian (q, px, py, pz) -> graph/e3nn (q, py, pz, px). + features[:, 1:4] = features[:, [2, 3, 1]] + if positions.shape[0] != features.shape[0]: + raise ValueError("MM positions and electrostatic sources must have the same length.") + + transform = getattr(model, "_charges_to_mul_ir", None) + if transform is not None: + features = transform(features) + + batch = data.get("mm_source_batch") + if batch is None: + if int(data["pbc"].reshape(-1, 3).shape[0]) != 1: + raise ValueError("mm_source_batch is required for batched PolarMACE inputs.") + batch = torch.zeros(positions.shape[0], dtype=torch.long, device=positions.device) + else: + batch = batch.to(device=positions.device, dtype=torch.long).reshape(-1) + if batch.shape[0] != positions.shape[0]: + raise ValueError("mm_source_batch and mm_positions must have the same length.") + return {"positions": positions, "features": features, "batch": batch} + + +def enable_polarmace_external_sources(model): + """Return an eager wrapper that adds dynamic MM electrostatic sources. + + Models that have already been wrapped are returned unchanged. Non-PolarMACE + models are rejected rather than silently dropping the external interaction. + """ + import torch + + if getattr(model, "supports_external_electrostatics", False): + return model + if model.__class__.__name__ != "PolarMACE": + raise TypeError( + "External electrostatic sources require a PolarMACE model; got " + f"{model.__class__.__name__}." + ) + + feature_base = _rebuild_feature_block(model.electric_potential_descriptor) + energy_base = _rebuild_energy_block(model.coulomb_energy) + feature_block = _ExternalFeatureBlock(feature_base).module + energy_block = _ExternalEnergyBlock(energy_base).module + model.electric_potential_descriptor = feature_block + model.coulomb_energy = energy_block + + class PolarMACEExternalSources(torch.nn.Module): + supports_external_electrostatics = True + + def __init__(self, wrapped): + super().__init__() + self.model = wrapped + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.model, name) + + def forward( + self, + data, + training: bool = False, + compute_force: bool = True, + compute_virials: bool = False, + compute_stress: bool = False, + compute_displacement: bool = False, + compute_hessian: bool = False, + compute_edge_forces: bool = False, + compute_atomic_stresses: bool = False, + **kwargs, + ): + external = _prepare_external_sources(self.model, data, compute_force) + if external is None: + return self.model( + data, + training=training, + compute_force=compute_force, + compute_virials=compute_virials, + compute_stress=compute_stress, + compute_displacement=compute_displacement, + compute_hessian=compute_hessian, + compute_edge_forces=compute_edge_forces, + compute_atomic_stresses=compute_atomic_stresses, + **kwargs, + ) + if any( + (compute_virials, compute_stress, compute_displacement, compute_hessian, + compute_edge_forces, compute_atomic_stresses) + ): + raise NotImplementedError( + "The OpenMM external-source adapter currently supports energies and " + "Cartesian forces only." + ) + + self.model.electric_potential_descriptor.set_external_sources(external) + self.model.coulomb_energy.set_external_sources(external) + try: + result = self.model( + data, + training=training, + compute_force=False, + compute_virials=False, + compute_stress=False, + compute_displacement=False, + compute_hessian=False, + compute_edge_forces=False, + compute_atomic_stresses=False, + **kwargs, + ) + if compute_force: + ml_gradient, mm_gradient = torch.autograd.grad( + outputs=[result["energy"]], + inputs=[data["positions"], external["positions"]], + grad_outputs=[torch.ones_like(result["energy"])], + create_graph=training, + retain_graph=training, + allow_unused=True, + ) + result["forces"] = ( + torch.zeros_like(data["positions"]) + if ml_gradient is None else -ml_gradient + ) + result["mm_forces"] = ( + torch.zeros_like(external["positions"]) + if mm_gradient is None else -mm_gradient + ) + else: + result["mm_forces"] = None + return result + finally: + self.model.electric_potential_descriptor.set_external_sources(None) + self.model.coulomb_energy.set_external_sources(None) + + return PolarMACEExternalSources(model) + + +__all__ = ["enable_polarmace_external_sources"] diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 829141b..6eeab59 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -28,13 +28,19 @@ OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +import os import openmm from openmm import unit from openmmml.mlpotential import MLPotential, MLPotentialImpl, MLPotentialImplFactory -from typing import Iterable, Optional +from openmmml.embeddings import utilities +from typing import Iterable, Optional, Sequence, Tuple, Union from functools import partial +from pathlib import Path import numpy as np +LinkRecordTuple = Tuple[int, int, float] # (q_global, m_global, target_dist) +LinkRecordsArg = Union[str, "os.PathLike[str]", Sequence[LinkRecordTuple], None] + class MACEPotentialImplFactory(MLPotentialImplFactory): """This is the factory that creates MACEPotentialImpl objects.""" @@ -58,7 +64,9 @@ class MACEPotentialImpl(MLPotentialImpl): Other available models include 'mace-off23-medium', 'mace-off23-large', 'mace-off24-medium', 'mace-mpa-0-medium', 'mace-omat-0-small', 'mace-omat-0-medium', 'mace-omol-0-extra-large', - 'mace-les-off-small', 'mace-polar-1-small', 'mace-polar-1-medium', and 'mace-polar-1-large'. + 'mace-les-off-small', and the PolarMACE models 'mace-polar-1-small', + 'mace-polar-1-medium', and 'mace-polar-1-large'. The PolarMACE models are + the ones that support electrostatic embedding. To use a locally trained MACE model, provide the path to the model file. For example: @@ -73,14 +81,30 @@ class MACEPotentialImpl(MLPotentialImpl): According to the MACE documentation, 'single' precision is recommended for MD (faster but less accurate), while 'double' precision is recommended for geometry optimization. - Additionally, you can request computation of the full atomic energy, including the atom - self-energy, instead of the default interaction energy, by setting ``returnEnergyType`` to - 'energy'. For example: - - >>> system = potential.createSystem(topology, returnEnergyType='energy') - - The default is to compute the interaction energy, which can be made explicit by setting - ``returnEnergyType='interaction_energy'``. + By default the reported energy is the full ``energy`` returned by the MACE + model — the same scalar whose gradient w.r.t. positions is reported as the + force, so the resulting potential is exactly conservative. To get only the + message-passing readout component, set ``returnEnergyType='interaction_energy'``: + + >>> system = potential.createSystem(topology, returnEnergyType='interaction_energy') + + Note: ``returnEnergyType='interaction_energy'`` is **not** energy/force + consistent for the PolarMACE family, which adds Coulomb / dipole / local- + electron terms to ``total_energy`` whose gradients are in ``forces`` but + which are not in ``interaction_energy``. + + Precision caveat for ``returnEnergyType='energy'``: this key returns the + full ``total_energy = e0 + inter_e + extras`` where ``e0`` are the + model's per-atom reference energies. For foundation models (mace-mp, + mace-off, mace-omat, ...) ``e0`` is typically tens of eV per atom, so the + reported scalar for a large ML region can be 10⁴–10⁶ eV in magnitude. + The **forces** stay exact at any scale (they are gradients of this same + scalar), but the **energy** column written into single-precision OpenMM + state files / log lines may carry only ~6–7 significant digits at that + magnitude — meV resolution is lost. Use ``precision='double'`` if you + need accurate absolute energies, or note that energy differences (e.g. + NVE drift) still resolve cleanly because the e0 contribution cancels in + the difference. A runtime warning fires when this regime is detected. Attributes ---------- @@ -90,20 +114,25 @@ class MACEPotentialImpl(MLPotentialImpl): The path to the locally trained MACE model if ``name`` is 'mace'. """ - # (Function name, model name, restrictive license name or None, long-range) + # (Function name, model name, restrictive license name or None, long-range, + # accepts MM charges) + # + # The last flag records whether the model can be given the charges and + # positions of the atoms outside the ML subset, which is what electrostatic + # embedding requires. Only the PolarMACE family can. KNOWN_MODELS = { - 'mace-off23-small': ('mace_off', 'small', 'ASL', False), - 'mace-off23-medium': ('mace_off', 'medium', 'ASL', False), - 'mace-off23-large': ('mace_off', 'large', 'ASL', False), - 'mace-off24-medium': ('mace_off', 'https://github.com/ACEsuit/mace-off/blob/main/mace_off24/MACE-OFF24_medium.model?raw=true', 'ASL', False), - 'mace-mpa-0-medium': ('mace_mp', 'medium-mpa-0', None, False), - 'mace-omat-0-small': ('mace_mp', 'small-omat-0', 'ASL', False), - 'mace-omat-0-medium': ('mace_mp', 'medium-omat-0', 'ASL', False), - 'mace-omol-0-extra-large': ('mace_omol', 'extra_large', 'ASL', False), - 'mace-les-off-small': ('mace_off', 'https://github.com/ChengUCB/les_fit/blob/main/MACELES-OFF/MACELES-OFF_small_converted.model?raw=true', 'CC BY-NC 4.0', True), - 'mace-polar-1-small': ('mace_polar', 'polar-1-s', None, True), - 'mace-polar-1-medium': ('mace_polar', 'polar-1-m', None, True), - 'mace-polar-1-large': ('mace_polar', 'polar-1-l', None, True), + 'mace-off23-small': ('mace_off', 'small', 'ASL', False, False), + 'mace-off23-medium': ('mace_off', 'medium', 'ASL', False, False), + 'mace-off23-large': ('mace_off', 'large', 'ASL', False, False), + 'mace-off24-medium': ('mace_off', 'https://github.com/ACEsuit/mace-off/blob/main/mace_off24/MACE-OFF24_medium.model?raw=true', 'ASL', False, False), + 'mace-mpa-0-medium': ('mace_mp', 'medium-mpa-0', None, False, False), + 'mace-omat-0-small': ('mace_mp', 'small-omat-0', 'ASL', False, False), + 'mace-omat-0-medium': ('mace_mp', 'medium-omat-0', 'ASL', False, False), + 'mace-omol-0-extra-large': ('mace_omol', 'extra_large', 'ASL', False, False), + 'mace-les-off-small': ('mace_off', 'https://github.com/ChengUCB/les_fit/blob/main/MACELES-OFF/MACELES-OFF_small_converted.model?raw=true', 'CC BY-NC 4.0', True, False), + 'mace-polar-1-small': ('mace_polar', 'polar-1-s', None, True, True), + 'mace-polar-1-medium': ('mace_polar', 'polar-1-m', None, True, True), + 'mace-polar-1-large': ('mace_polar', 'polar-1-l', None, True, True), } def __init__(self, name: str, modelPath) -> None: @@ -123,6 +152,56 @@ def __init__(self, name: str, modelPath) -> None: """ self.name = name self.modelPath = modelPath + self._preloadedModel = None + + def _loadModel(self, args): + """Load the MACE model, returning it along with the device it is on. + + If createMixedSystem() has already loaded a model in order to inspect + it, that one is handed over here rather than the checkpoint being read a + second time. The handover is consumed on use, so each call after that + loads afresh; holding the model indefinitely would mean a later call + with a different precision converting an already converted model. + """ + import torch + try: + from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol, mace_polar + except ImportError as e: + raise ImportError(f"Failed to import mace with error: {e}. Install mace with 'pip install mace-torch'.") + + device = self._getTorchDevice(args) + preloaded, self._preloadedModel = self._preloadedModel, None + if preloaded is not None and preloaded[1] == device: + return preloaded[0], device + + if self.name in MACEPotentialImpl.KNOWN_MODELS: + functions = { + 'mace_off': mace_off, + 'mace_mp': mace_mp, + 'mace_omol': mace_omol, + 'mace_polar': mace_polar, + } + fnName, name, restrictiveLicense, _, _ = MACEPotentialImpl.KNOWN_MODELS[self.name] + model = functions[fnName](model=name, device=device, return_raw_model=True).to(device) + if restrictiveLicense is not None: + import logging + logging.warning(f'The model {self.name} is distributed under the restrictive {restrictiveLicense} license. Commercial use is not permitted.') + elif self.name == "mace": + if self.modelPath is not None: + model = torch.load(self.modelPath, map_location=device) + if hasattr(model, "to"): + model = model.to(device) + else: + raise ValueError("No modelPath provided for local MACE model.") + else: + raise ValueError(f"Unsupported MACE model: {self.name}") + if model.__class__.__name__ == "PolarMACE": + from openmmml.models._polarmace_external import ( + enable_polarmace_external_sources, + ) + + model = enable_polarmace_external_sources(model) + return model, device def addForces( self, @@ -131,7 +210,11 @@ def addForces( atoms: Optional[Iterable[int]], forceGroup: int, precision: Optional[str] = None, - returnEnergyType: str = "interaction_energy", + returnEnergyType: str = "energy", + linkRecords: LinkRecordsArg = None, + linkChargeScheme: str = "dz1", + embedding: str = "mechanical", + customNonbondedChargeParameter: Optional[str] = None, **args, ) -> None: """ @@ -151,8 +234,44 @@ def addForces( The precision of the model. Supported options are 'single' and 'double'. If ``None``, the default precision of the model is used. returnEnergyType : str, optional - Whether to return the interaction energy or the energy including the self-energy. - Default is 'interaction_energy'. Supported options are 'interaction_energy' and 'energy'. + Which scalar from the MACE model output is reported to OpenMM as + the potential energy. Default ``'energy'`` is the same quantity + the force vector is differentiated against, so OpenMM sees a + self-consistent (conservative) potential. ``'interaction_energy'`` + returns only the message-passing readout; for PolarMACE this is + **not** the gradient partner of ``forces`` and will produce + apparent NVE drift / a non-zero finite-difference plateau. + linkRecords : str / path / sequence of (q_global, m_global, target_dist) / None + Hydrogen link-atom cap records for QM/MM boundary bonds. + linkChargeScheme : str, optional + How to handle the partial charges on the MM-side boundary atoms + (M atoms) when ``linkRecords`` is provided. The M-atom partial + charge would otherwise sit ~1.5 Å from the nearest QM atom (through + the link H) and over-polarise the QM region's MACE-predicted + electronic structure. Supported in this iteration: + + - ``"none"``: leave MM charges untouched (legacy behaviour). + - ``"z1"``: set q_M = 0 for every M atom. Cheapest fix; breaks + total MM-charge neutrality by -q_M_orig. + - ``"dz1"`` (default): q_M = 0 plus q_M_orig is distributed + equally onto M's MM neighbours (M1 atoms). Preserves total + MM charge to round-off. Falls back to Z1 (with a warning) for + any M atom that has zero MM neighbours. + + Only modifies the MM charge array passed to the ML potential; + the OpenMM ``NonbondedForce`` is left untouched, so MM-MM Coulomb + is bit-exact with the original force field (standard QM/MM + practice). Z2 and RCD (which require per-step virtual charges) + are not in this iteration. + embedding : {"mechanical", "electrostatic"} + Which embedding method the caller is implementing. Set by + ``createMixedSystem``; there is normally no reason to pass it here + directly. ``mechanical`` (the default) does not pass MM positions + or charges into MACE. ``electrostatic`` passes them into PolarMACE + and scatters the returned ``mm_forces`` back onto the MM atoms; the + removal of the classical ML-MM Coulomb that this assumes is done by + ``createMixedSystem``, not here. It is an error to request it for a + model that cannot accept MM charges. """ import torch try: @@ -169,35 +288,35 @@ def addForces( # Load the model. - device = self._getTorchDevice(args) - if self.name in MACEPotentialImpl.KNOWN_MODELS: - functions = { - 'mace_off': mace_off, - 'mace_mp': mace_mp, - 'mace_omol': mace_omol, - 'mace_polar': mace_polar, - } - fnName, name, restrictiveLicense, _ = MACEPotentialImpl.KNOWN_MODELS[self.name] - model = functions[fnName](model=name, device=device, return_raw_model=True).to(device) - if restrictiveLicense is not None: - import logging - logging.warning(f'The model {self.name} is distributed under the restrictive {restrictiveLicense} license. Commercial use is not permitted.') - elif self.name == "mace": - if self.modelPath is not None: - model = torch.load(self.modelPath, map_location=device).to(device) - else: - raise ValueError("No modelPath provided for local MACE model.") - else: - raise ValueError(f"Unsupported MACE model: {self.name}") + model, device = self._loadModel(args) - # Get the atomic numbers of the ML region. + use_mm_embedding = _should_use_mm_embedding(model, atoms, embedding) includedAtoms = list(topology.atoms()) if atoms is not None: includedAtoms = [includedAtoms[i] for i in atoms] atomicNumbers = [atom.element.atomic_number for atom in includedAtoms] - # Set the precision that the model will be used with. + if returnEnergyType == "energy": + try: + e0_max = float(model.atomic_energies_fn.atomic_energies.detach().abs().max()) + except AttributeError: + e0_max = 0.0 + if e0_max > 100.0: # 1 eV per atom is conservative; foundation models far exceed this + import warnings as _w + _w.warn( + f"returnEnergyType='energy' includes per-atom reference " + f"energies (max |e0| = {e0_max:.2f} eV/atom over {int(model.atomic_energies_fn.atomic_energies.numel())} " + f"element entries). Use precision='double' if you need accurate absolute " + f"energies, or pass returnEnergyType='interaction_energy' for the " + f"e0-subtracted readout (note: only 'energy' is gradient-consistent " + f"with PolarMACE — see the docstring).", + stacklevel=2, + ) + + linkInfo = _prepareLinkRecords(linkRecords, atoms, topology, system) + if linkInfo is not None: + atomicNumbers = atomicNumbers + [1] * linkInfo["K"] modelDefaultDtype = next(model.parameters()).dtype if precision is None: @@ -210,34 +329,65 @@ def addForces( raise ValueError(f"Unsupported precision {precision} for the model. Supported values are 'single' and 'double'.") if dtype != modelDefaultDtype: print(f"Model dtype is {modelDefaultDtype} and requested dtype is {dtype}. The model will be converted to the requested dtype.") + # Actually do the conversion. The previous code only printed the + # warning and left the model untouched, which caused dtype + # mismatches inside e3nn's compiled TensorProduct submodules + # when inputs were passed at the requested dtype. model = model.to(dtype) - # One hot encoding of atomic numbers + model_device = device + try: + model_device = next(model.parameters()).device + except (AttributeError, StopIteration): + pass zTable = utils.AtomicNumberTable([int(z) for z in model.atomic_numbers]) nodeAttrs = to_one_hot( - torch.tensor(atomic_numbers_to_indices(atomicNumbers, z_table=zTable), dtype=torch.long, device=device).unsqueeze(-1), + torch.tensor(atomic_numbers_to_indices(atomicNumbers, z_table=zTable), dtype=torch.long, device=model_device).unsqueeze(-1), num_classes=len(zTable)) if atoms is None: indices = None else: indices = np.array(atoms) + mmInfo = None + if use_mm_embedding: + # ML-MM Coulomb is removed by MLPotential.createMixedSystem when + # embedding='electrostatic'. We only need MM positions/charges for + # the PolarMACE input here. + mmInfo = _prepareMMEmbedding(system, atoms, customNonbondedChargeParameter) + + # Optional Z1 / DZ1 link-atom charge redistribution. Standard QM/MM + # correction to stop the QM region from being over-polarised by + # the partial charge on the MM-side boundary atom. + if linkInfo is not None and linkChargeScheme not in (None, "none"): + from openmmml.models._links import ( + apply_link_charge_redistribution as _apply_link_q, + ) + mmInfo["mm_charges"] = _apply_link_q( + mm_atoms=mmInfo["mm_atoms"], + mm_charges=mmInfo["mm_charges"], + link_info=linkInfo, + topology=topology, + scheme=linkChargeScheme, + ) + print(f"[link-charge-redistribution] scheme={linkChargeScheme} " + f"M atoms touched={len(linkInfo['m_global'])}") periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions() - # Create the PythonForce and add it to the System. - compute = partial(_computeMACE, model=model, - ptr=torch.tensor([0, nodeAttrs.shape[0]], dtype=torch.long, device=device, requires_grad=False), + ptr=torch.tensor([0, nodeAttrs.shape[0]], dtype=torch.long, device=model_device, requires_grad=False), node_attrs=nodeAttrs.to(dtype), - batch=torch.zeros(nodeAttrs.shape[0], dtype=torch.long, device=device, requires_grad=False), - pbc=torch.tensor([periodic, periodic, periodic], dtype=torch.bool, device=device, requires_grad=False), + batch=torch.zeros(nodeAttrs.shape[0], dtype=torch.long, device=model_device, requires_grad=False), + pbc=torch.tensor([periodic, periodic, periodic], dtype=torch.bool, device=model_device, requires_grad=False), returnEnergyType=returnEnergyType, - charge=torch.tensor([float(args.get('charge', 0))], dtype=dtype, device=device, requires_grad=False), - multiplicity=torch.tensor([float(args.get('multiplicity', 1))], dtype=dtype, device=device, requires_grad=False), + charge=torch.tensor([float(args.get('charge', 0))], dtype=dtype, device=model_device, requires_grad=False), + multiplicity=torch.tensor([float(args.get('multiplicity', 1))], dtype=dtype, device=model_device, requires_grad=False), indices=indices, - periodic=periodic) + periodic=periodic, + linkInfo=linkInfo, + mmInfo=mmInfo) force = openmm.PythonForce(compute) force.setForceGroup(forceGroup) force.setUsesPeriodicBoundaryConditions(periodic) @@ -245,20 +395,332 @@ def addForces( def getMLLongRange(self) -> bool | None: if self.name in MACEPotentialImpl.KNOWN_MODELS: - _, _, _, longRange = MACEPotentialImpl.KNOWN_MODELS[self.name] + _, _, _, longRange, _ = MACEPotentialImpl.KNOWN_MODELS[self.name] return longRange return None + def getSupportedEmbeddings(self) -> list[str]: + + # Electrostatic embedding requires a model that accepts the charges and + # positions of the atoms outside the ML subset, which of the pretrained + # models only the PolarMACE family does. A custom checkpoint may be a + # PolarMACE model too, but that cannot be known without loading it, so + # the method is offered and createMixedSystem() rejects the checkpoint + # once loaded if it turns out not to be one. + + if self.name in MACEPotentialImpl.KNOWN_MODELS: + _, _, _, _, acceptsMMCharges = MACEPotentialImpl.KNOWN_MODELS[self.name] + return ["electrostatic"] if acceptsMMCharges else [] + return ["electrostatic"] + + def createMixedSystem(self, + topology: openmm.app.Topology, + system: openmm.System, + atoms: list[int], + forceGroup: int, + interpolate: bool, + embedding: str, + customNonbondedHasCharges: Optional[bool] = None, + customNonbondedChargeParameter: Optional[str] = None, + **args) -> openmm.System: + """Create a mixed system using electrostatic embedding. + + The model, rather than the conventional force field, is responsible for + the electrostatic interactions between the atoms within the ML subset + and those outside of it: it is passed the positions and conventional + force field charges of the atoms outside the ML subset, and returns + forces on them alongside the forces on the ML subset. The ML subset can + therefore polarize in response to its surroundings, which mechanical + embedding does not allow. + + This is implemented as the "global charge zero" variant: the + conventional force field charge of every atom in the ML subset is set to + zero. Every Coulomb term involving an ML atom is then zero by + construction, including the reciprocal space part of PME, without any + per-pair exceptions being added. Adding an exception for each ML-MM + pair would instead be incorrect under periodic boundary conditions, + since NonbondedForce evaluates exceptions using plain Cartesian + distances rather than the minimum image convention, so the ML-MM + Lennard-Jones interaction would silently vanish for any pair that is + only within the cutoff across a periodic boundary. Lennard-Jones is + left to the conventional force field and continues to use the ordinary, + periodicity-aware pair list. + + Interactions within the ML subset are excluded entirely, as the model + computes them. Bonded terms that cross the ML/MM boundary are retained. + + Only models that accept MM charges and positions, which at present means + the PolarMACE family, can be used with this embedding method. An error + is raised for any other model rather than falling back to mechanical + embedding, since by that point the ML-MM electrostatics have already + been removed from the conventional force field and a fallback would + simply lose them. + + Because this method has to account for every Coulomb term in the force + field, it requires the System to contain exactly one NonbondedForce: the + MM charges given to the model are read from one, so several would be + ambiguous. + + It also needs to be told about any CustomNonbondedForce, whose energy + expression is arbitrary and cannot be inspected here. Pass + customNonbondedHasCharges=False to declare that it holds no + electrostatics, or True together with customNonbondedChargeParameter + naming the per-particle parameter that holds the charge, which is then + zeroed on the ML atoms exactly as for the NonbondedForce. An error is + raised if the answer is needed and has not been given. + + Note that zeroing that parameter removes the ML terms only if the + expression is multiplicatively separable in the charge, as the usual + q1*q2/r is. That cannot be verified here, so it is the caller's + responsibility. + """ + + if embedding != "electrostatic": + raise ValueError(f"Unsupported embedding type: {embedding}") + + # Check that the model can actually accept MM charges and positions + # before touching the System, so that an unsuitable model is rejected + # with the System left alone rather than stripped of its ML-MM + # electrostatics. This is also the first point at which the check is + # possible, since it needs the loaded checkpoint; the model is handed to + # addForces() below so the checkpoint is only read once. + + model, device = self._loadModel(args) + if not _supports_mm_embedding(model): + raise ValueError( + f"embedding='{embedding}' requires a model that accepts MM charges " + f"and positions (PolarMACE); got {model.__class__.__name__}." + ) + + if interpolate: + # At lambda_interpolate=0 the conventional endpoint would be missing + # the ML-MM Coulomb energy, which is removed from the conventional + # force field outside of the interpolating CustomCVForce and cannot + # be restored from within it. + raise ValueError("Electrostatic embedding does not support interpolation.") + + periodic = system.usesPeriodicBoundaryConditions() + + # Electrostatic embedding has to account for every Coulomb term in the + # force field: the ones involving the ML subset are removed here on the + # understanding that the model supplies them. Anything it cannot see is + # either left in place and counted twice, or removed and never replaced, + # and in both cases the result is a wrong energy rather than an error. + # So refuse the cases where the electrostatics cannot be located rather + # than guessing. + + nonbondedForces = [f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)] + if len(nonbondedForces) > 1: + # The MM charges handed to the model are read from a single + # NonbondedForce, so several of them are ambiguous. + raise ValueError("Multiple NonbondedForce objects encountered; electrostatic embedding requires exactly one.") + + if any(isinstance(f, openmm.CustomNonbondedForce) for f in system.getForces()): + # A CustomNonbondedForce's energy expression is arbitrary, so + # whether it contains electrostatics cannot be determined here. + if customNonbondedHasCharges is None: + raise ValueError("The System contains a CustomNonbondedForce and it is unknown whether it includes electrostatic interactions; pass customNonbondedHasCharges to specify.") + if customNonbondedHasCharges and customNonbondedChargeParameter is None: + raise ValueError("A CustomNonbondedForce includes electrostatic interactions, so customNonbondedChargeParameter must name the per-particle parameter holding the charge.") + + # Create the new system with the ML-ML interactions that the model + # computes removed. + + newSystem = utilities.removeBonds(system, topology, atoms, True) + atomSet = set(atoms) + + for force in newSystem.getForces(): + if isinstance(force, openmm.NonbondedForce): + + # Zero the charge of every ML atom, which removes all Coulomb + # interactions involving the ML subset while leaving its + # Lennard-Jones parameters, and the MM-MM interactions, + # untouched. + + for atom in atoms: + charge, sigma, epsilon = force.getParticleParameters(atom) + force.setParticleParameters(atom, 0.0, sigma, epsilon) + + # setParticleParameters() does not update the charge products + # that were precomputed for existing exceptions, so the 1-4 + # Coulomb terms crossing the ML/MM boundary have to be zeroed + # separately. + + for index in range(force.getNumExceptions()): + p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(index) + if p1 in atomSet or p2 in atomSet: + force.setExceptionParameters(index, p1, p2, 0.0, sigma, epsilon) + + # Exclude the ML-ML interactions entirely. + + for i in range(len(atoms)): + for j in range(i): + force.addException(atoms[i], atoms[j], 0, 1, 0, True) + + force.setExceptionsUsePeriodicBoundaryConditions(periodic) + + elif isinstance(force, openmm.CustomNonbondedForce): + + # Zero the named charge parameter on the ML atoms, the same + # trick used for the NonbondedForce above. Unlike there it is + # not guaranteed to work: it removes the ML terms only if the + # energy expression is multiplicatively separable in the charge, + # as q1*q2/r is, and that cannot be checked here. + + if customNonbondedChargeParameter is not None: + names = [force.getPerParticleParameterName(i) + for i in range(force.getNumPerParticleParameters())] + if customNonbondedChargeParameter not in names: + raise ValueError(f"A CustomNonbondedForce has no per-particle parameter '{customNonbondedChargeParameter}'; it defines {names}.") + chargeIndex = names.index(customNonbondedChargeParameter) + for atom in atoms: + parameters = list(force.getParticleParameters(atom)) + parameters[chargeIndex] = 0.0 + force.setParticleParameters(atom, parameters) + + utilities.addCustomNonbondedExclusions(force, atoms) + + # Add the ML potential, telling it that it is responsible for the + # electrostatic interactions with the atoms outside the ML subset, and + # handing over the model already loaded above. + + self._preloadedModel = (model, device) + try: + self.addForces(topology, newSystem, atoms, forceGroup, embedding=embedding, + customNonbondedChargeParameter=customNonbondedChargeParameter, **args) + finally: + self._preloadedModel = None + + return newSystem + + +def _supports_mm_embedding(model) -> bool: + return bool(getattr(model, "supports_external_electrostatics", False)) + + +_SUPPORTED_EMBEDDINGS = ("mechanical", "electrostatic") +_MM_EMBEDDING_MODES = ("electrostatic",) + + +def _should_use_mm_embedding(model, atoms: Optional[Iterable[int]], embedding: str) -> bool: + if embedding not in _SUPPORTED_EMBEDDINGS: + raise ValueError( + f"Unsupported embedding mode '{embedding}'. Supported values are " + + ", ".join(repr(m) for m in _SUPPORTED_EMBEDDINGS) + + "." + ) + if embedding not in _MM_EMBEDDING_MODES: + return False + if not _supports_mm_embedding(model): + # The mixed system has had its ML-MM Coulomb removed on the assumption + # that the model will supply it, so falling back to mechanical + # embedding here would silently discard those interactions. + raise ValueError( + f"embedding='{embedding}' requires a model that accepts MM charges " + f"and positions (PolarMACE); got {model.__class__.__name__}." + ) + if atoms is None: + raise ValueError( + f"embedding='{embedding}' requires an ML subset; it cannot be used " + "with createSystem()." + ) + return True + + +def _prepareMMEmbedding(system: openmm.System, atoms: Optional[Iterable[int]], + customNonbondedChargeParameter: Optional[str] = None): + """Extract the MM complement and its charges from the system's NonbondedForce.""" + if atoms is None: + return None + + num_particles = int(system.getNumParticles()) + ml_atoms = np.asarray(list(atoms), dtype=np.int64) + ml_set = set(int(i) for i in ml_atoms.tolist()) + mm_atoms = np.asarray( + [i for i in range(num_particles) if i not in ml_set], dtype=np.int64 + ) + + # The charges given to the model have to come from wherever the force field + # actually keeps them, which is the same force createMixedSystem zeroed the + # ML charges in. When that is a CustomNonbondedForce the caller has named + # the parameter holding them. + mm_charges = np.empty(len(mm_atoms), dtype=np.float64) + + if customNonbondedChargeParameter is not None: + custom = None + for force in system.getForces(): + if isinstance(force, openmm.CustomNonbondedForce): + names = [force.getPerParticleParameterName(i) + for i in range(force.getNumPerParticleParameters())] + if customNonbondedChargeParameter in names: + custom = force + chargeIndex = names.index(customNonbondedChargeParameter) + break + if custom is None: + raise ValueError(f"No CustomNonbondedForce defines a per-particle parameter '{customNonbondedChargeParameter}'.") + for row, atom_index in enumerate(mm_atoms): + mm_charges[row] = custom.getParticleParameters(int(atom_index))[chargeIndex] + return { + "ml_atoms": ml_atoms, + "mm_atoms": mm_atoms, + "mm_charges": mm_charges, + } + + nonbonded = None + for force in system.getForces(): + if isinstance(force, openmm.NonbondedForce): + nonbonded = force + break + if nonbonded is None: + raise ValueError( + "PolarMACE MM embedding requires a NonbondedForce to source MM charges." + ) + + for row, atom_index in enumerate(mm_atoms): + charge, _, _ = nonbonded.getParticleParameters(int(atom_index)) + mm_charges[row] = charge.value_in_unit(unit.elementary_charge) + + return { + "ml_atoms": ml_atoms, + "mm_atoms": mm_atoms, + "mm_charges": mm_charges, + } -def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, indices, periodic): + +def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, indices, periodic, linkInfo=None, mmInfo=None): import torch from mace.data.neighborhood import get_neighborhood energyScale = 96.4853 lengthScale = 10.0 - positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom) - numAtoms = positions.shape[0] + positions_full = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom) + numAtoms = positions_full.shape[0] if indices is not None: - positions = positions[indices] + positions = positions_full[indices] + else: + positions = positions_full + + # Link atoms: append K fictitious H positions placed each step from current + # Q, M coordinates. These extend the model input only; they never enter the + # OpenMM system. See docs/plans/link-atom-inference.md. + if linkInfo is not None: + if indices is None: + raise ValueError("linkRecords requires an explicit `atoms` subset.") + from openmmml.models._links import compute_cap_positions, minimum_image_M + r_Q = positions_full[linkInfo["q_global"]] + r_M = positions_full[linkInfo["m_global"]] + if periodic: + # Minimum-image the Q->M bond so caps are placed correctly even if Q + # and M sit across a periodic boundary. The same imaged r_M is reused + # for force redistribution below, so the returned force stays the + # gradient of the reported energy. Triclinic-correct; raises on a + # genuinely wrapped pair (see minimum_image_M). + cell_A = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) + r_M = minimum_image_M(r_Q, r_M, cell_A) + pos_link, _C_L_unused = compute_cap_positions( + r_Q, r_M, linkInfo["target_dist"] + ) + positions = np.concatenate([positions, pos_link], axis=0) + if periodic: cell = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) else: @@ -266,7 +728,12 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch dtype = node_attrs.dtype cutoff = float(model.r_max.detach()) edgeIndex, shifts, _, _ = get_neighborhood(positions, cutoff, [periodic, periodic, periodic], cell) - cellTensor = torch.tensor(cell, dtype=dtype, device=ptr.device) + cell_tensor = torch.tensor(cell, dtype=dtype, device=ptr.device) + volume = torch.linalg.det(cell_tensor) + if torch.abs(volume) > 0: + rcell = 2 * torch.pi * torch.linalg.inv(cell_tensor.mT) + else: + rcell = torch.zeros((3, 3), dtype=dtype, device=ptr.device) inputDict = { "ptr": ptr, "node_attrs": node_attrs, @@ -275,19 +742,153 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch "positions": torch.tensor(positions, dtype=dtype, device=ptr.device), "edge_index": torch.tensor(edgeIndex, dtype=torch.int64, device=ptr.device), "shifts": torch.tensor(shifts, dtype=dtype, device=ptr.device), - "cell": cellTensor, - "rcell": 2 * torch.pi * torch.linalg.inv(cellTensor.mT), - "volume": torch.linalg.det(cellTensor), + "cell": cell_tensor, + "rcell": rcell, + "volume": volume.reshape(-1), "total_charge": charge, "total_spin": multiplicity, - "external_field": torch.zeros((1, 3), dtype=dtype, device=ptr.device), + "external_field": torch.zeros((charge.shape[0], 3), dtype=dtype, device=ptr.device), "fermi_level": torch.zeros((1,), dtype=dtype, device=ptr.device) } + # load mm position and charges + if mmInfo is not None: + mm_positions = positions_full[mmInfo["mm_atoms"]] + inputDict["mm_positions"] = torch.tensor( + mm_positions, dtype=dtype, device=ptr.device + ) + inputDict["mm_charges"] = torch.tensor( + mmInfo["mm_charges"], dtype=dtype, device=ptr.device + ) + inputDict["mm_source_batch"] = torch.zeros( + len(mmInfo["mm_atoms"]), dtype=torch.long, device=ptr.device + ) results = model(inputDict, compute_force=True) energy = float(results[returnEnergyType].detach())*energyScale forces = (results["forces"]*energyScale*lengthScale).detach().cpu().numpy() + mm_forces = results.get("mm_forces") + if mmInfo is not None and mm_forces is None: + # The mixed system has had its ML-MM electrostatics removed on the + # understanding that this model supplies them. A model that returns no + # forces on the MM atoms did not compute them, so continuing would leave + # those interactions missing entirely rather than merely approximated, + # and nothing downstream would report it. The usual cause is a + # PolarMACE checkpoint whose forward does not accept mm_charges. + raise ValueError("The model returned no 'mm_forces' although MM charges were supplied; it does not implement electrostatic embedding.") + if mm_forces is not None: + mm_forces = (mm_forces * energyScale * lengthScale).detach().cpu().numpy() + + # force redistribution and add mm_forces back to the full system if indices is not None: f = np.zeros((numAtoms, 3), dtype=(np.float64 if dtype == torch.float64 else np.float32)) - f[indices] = forces + if linkInfo is None: + f[indices] = forces + else: + from openmmml.models._links import ( + compute_cap_positions, + minimum_image_M, + redistribute_cap_force, + ) + N = len(indices) + f_ml = forces[:N] + f_link = forces[N:] + f[indices] = f_ml + + # Redistribute each link atom's force onto its (Q, M) partners. + r_Q = positions_full[linkInfo["q_global"]] + r_M = positions_full[linkInfo["m_global"]] + if periodic: + cell_A = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) + r_M = minimum_image_M(r_Q, r_M, cell_A) + _, C_L = compute_cap_positions(r_Q, r_M, linkInfo["target_dist"]) + F_Q_add, F_M_add = redistribute_cap_force(f_link, r_Q, r_M, C_L) + + f[linkInfo["q_global"]] += F_Q_add.astype(f.dtype, copy=False) + f[linkInfo["m_global"]] += F_M_add.astype(f.dtype, copy=False) + if mmInfo is not None and mm_forces is not None: + f[mmInfo["mm_atoms"]] += mm_forces.astype(f.dtype, copy=False) forces = f return energy, forces + + +def _prepareLinkRecords(linkRecords, atoms, topology, system): + """Normalize the ``linkRecords`` argument into a frozen bundle for the + per-step closure. Returns ``None`` if no link records were supplied. + + Assertions (failing loudly): + * non-periodic system (PBC deferred to a follow-up PR) + * ``atoms`` is not None + * every ``q_global`` is in ``atoms`` + * no ``m_global`` is in ``atoms`` + * every (Q, M) pair is unique, and no atom appears as Q or M in more + than one cap (simplifies per-step scatter; relax later if needed) + """ + if linkRecords is None: + return None + + # Periodic systems are allowed: cap positions are placed with minimum-image + # at runtime (_computeMACE), which is exact as long as each link bond is + # shorter than half the box -- always true for real frontier bonds. A + # genuinely wrapped (Q, M) pair (longer than half the box) raises there. + if atoms is None: + raise ValueError("linkRecords requires an explicit `atoms` subset.") + + if isinstance(linkRecords, (str, Path)): + # A capping-mapping CSV is 1-based: read (q_idx1, m_idx1, + # target_dist_ang) and convert to 0-based OpenMM indices. target_dist + # stays in Angstroms, canonical because it matches the MACE-side + # `positions_full`. Cap positions are not stored in the CSV; they are + # recomputed each step from (q, m, target_dist). + import csv as _csv + tuples = [] + with open(linkRecords, newline="") as _f: + for row in _csv.DictReader(_f): + tuples.append( + (int(row["q_idx1"]) - 1, int(row["m_idx1"]) - 1, float(row["target_dist_ang"])) + ) + else: + # Tuple path: target_dist is also in Å (matches the docstring + # parameter name `target_dist_ang` and the MACE convention used + # in every upstream test fixture, e.g. linkRecords=[(q, m, 1.09)]). + tuples = [(int(q), int(m), float(td)) for (q, m, td) in linkRecords] + + if not tuples: + return None + + num_particles = int(system.getNumParticles()) + atoms_set = set(int(a) for a in atoms) + seen_pairs: set = set() + seen_q: set = set() + seen_m: set = set() + q_global = np.empty(len(tuples), dtype=np.int64) + m_global = np.empty(len(tuples), dtype=np.int64) + target_dist = np.empty(len(tuples), dtype=np.float64) + for k, (q, m, td) in enumerate(tuples): + if not (0 <= q < num_particles): + raise ValueError(f"linkRecords[{k}]: q_global={q} out of range [0, {num_particles}).") + if not (0 <= m < num_particles): + raise ValueError(f"linkRecords[{k}]: m_global={m} out of range [0, {num_particles}).") + if q not in atoms_set: + raise ValueError(f"linkRecords[{k}]: q_global={q} is not in `atoms`.") + if m in atoms_set: + raise ValueError(f"linkRecords[{k}]: m_global={m} is in `atoms` (should be MM).") + if (q, m) in seen_pairs: + raise ValueError(f"linkRecords[{k}]: duplicate (Q, M) pair ({q}, {m}).") + if q in seen_q or m in seen_m: + raise ValueError( + f"linkRecords[{k}]: atom {q if q in seen_q else m} appears in more than " + "one cap; current implementation requires unique Q and M across caps." + ) + if td <= 0: + raise ValueError(f"linkRecords[{k}]: target_dist must be positive; got {td}.") + seen_pairs.add((q, m)) + seen_q.add(q) + seen_m.add(m) + q_global[k] = q + m_global[k] = m + target_dist[k] = td + return { + "K": len(tuples), + "q_global": q_global, + "m_global": m_global, + "target_dist": target_dist, + } diff --git a/test/TestElectrostaticEmbeddingPME.py b/test/TestElectrostaticEmbeddingPME.py new file mode 100644 index 0000000..7689f4b --- /dev/null +++ b/test/TestElectrostaticEmbeddingPME.py @@ -0,0 +1,506 @@ +"""Unit tests for the periodic NonbondedForce / PME path of +``MLPotential.createMixedSystem(embedding='electrostatic')``. + +These tests exercise the MM-side surgery only. A no-op MLPotentialImpl +is registered so the tests do not depend on MACE, torch, or GPU. + +Slice 4 of docs/codex-plans/pbc-electrostatic-embedding-small-plan.md. +""" +from __future__ import annotations + +import numpy as np +import openmm +import openmm.app as app +import openmm.unit as unit +import pytest + +from openmmml import MLPotential +from openmmml.mlpotential import MLPotentialImpl, MLPotentialImplFactory +from openmmml.models.macepotential import MACEPotentialImpl + + +# --------------------------------------------------------------------------- +# Register a no-op MLPotentialImpl for tests +# --------------------------------------------------------------------------- + +class _NoopImpl(MACEPotentialImpl): + """The real electrostatic embedding with the model evaluation stubbed out. + + Electrostatic embedding is a MACE-specific embedding method, so these tests + inherit MACEPotentialImpl.createMixedSystem() to exercise the actual + nonbonded surgery under test. Only addForces() is stubbed, which is what + would otherwise require a MACE checkpoint and the PolarMACE stack; the + surgery runs before it and is unaffected. + """ + + def __init__(self): + super().__init__("mace", None) + + def _loadModel(self, args): + # createMixedSystem() loads the model to check that it accepts MM + # charges. Stand in for a real checkpoint with an object of the class + # it looks for, so the surgery under test runs without the PolarMACE + # stack being installed. + class PolarMACE: + pass + + return PolarMACE(), "cpu" + + def addForces(self, topology, system, atoms, forceGroup, **args): + return + + +class _NoopFactory(MLPotentialImplFactory): + def createImpl(self, name, **args): + return _NoopImpl() + + +MLPotential.registerImplFactory("noop_test_impl", _NoopFactory()) + + +# --------------------------------------------------------------------------- +# Test system construction +# --------------------------------------------------------------------------- + +# 4 atoms: ML = {0, 1}, MM = {2, 3}. Periodic 2 nm cubic box. +_PARAMS = [ + # (mass, charge_e, sigma_nm, epsilon_kj) + (12.0, 0.6, 0.30, 0.20), # ML 0 + (12.0, -0.4, 0.32, 0.25), # ML 1 + (16.0, -0.8, 0.31, 0.65), # MM 2 + (1.0, 0.4, 0.10, 0.05), # MM 3 +] +_BOX_NM = 2.0 +_ML_ATOMS = [0, 1] +_MM_ATOMS = [2, 3] + + +def _build_periodic_system(use_pme=True, with_boundary_bonded=True, with_pre_existing_exception=False): + system = openmm.System() + box_vec = _BOX_NM * unit.nanometer + system.setDefaultPeriodicBoxVectors( + openmm.Vec3(_BOX_NM, 0, 0) * unit.nanometer, + openmm.Vec3(0, _BOX_NM, 0) * unit.nanometer, + openmm.Vec3(0, 0, _BOX_NM) * unit.nanometer, + ) + + nonbonded = openmm.NonbondedForce() + if use_pme: + nonbonded.setNonbondedMethod(openmm.NonbondedForce.PME) + nonbonded.setCutoffDistance(0.6 * unit.nanometer) + else: + nonbonded.setNonbondedMethod(openmm.NonbondedForce.NoCutoff) + for mass, charge, sigma, epsilon in _PARAMS: + system.addParticle(mass) + nonbonded.addParticle( + charge * unit.elementary_charge, + sigma * unit.nanometer, + epsilon * unit.kilojoule_per_mole, + ) + + # A pre-existing 1-3 ML-MM exception lets us check that sigma/epsilon + # are read from the exception rather than synthesized via Lorentz-Berthelot. + if with_pre_existing_exception: + nonbonded.addException( + 1, 2, + 0.123 * unit.elementary_charge * unit.elementary_charge, + 0.345 * unit.nanometer, + 0.678 * unit.kilojoule_per_mole, + ) + + system.addForce(nonbonded) + + if with_boundary_bonded: + # ML-ML bond + ML-MM bond (boundary) + MM-MM bond + bonds = openmm.HarmonicBondForce() + bonds.addBond(0, 1, 0.15 * unit.nanometer, 1000.0 * unit.kilojoule_per_mole / unit.nanometer ** 2) # ML-ML + bonds.addBond(1, 2, 0.15 * unit.nanometer, 800.0 * unit.kilojoule_per_mole / unit.nanometer ** 2) # ML-MM (boundary) + bonds.addBond(2, 3, 0.10 * unit.nanometer, 500.0 * unit.kilojoule_per_mole / unit.nanometer ** 2) # MM-MM + system.addForce(bonds) + + angles = openmm.HarmonicAngleForce() + angles.addAngle(0, 1, 2, 1.9, 100.0 * unit.kilojoule_per_mole / unit.radian ** 2) # ML-ML-MM (boundary) + angles.addAngle(1, 2, 3, 1.9, 100.0 * unit.kilojoule_per_mole / unit.radian ** 2) # ML-MM-MM (boundary) + system.addForce(angles) + + return system + + +def _build_topology(): + topology = app.Topology() + chain = topology.addChain() + res = topology.addResidue("MOL", chain) + elements = [app.element.carbon, app.element.carbon, app.element.oxygen, app.element.hydrogen] + for i, el in enumerate(elements): + topology.addAtom(f"A{i}", el, res) + topology.setPeriodicBoxVectors( + unit.Quantity(np.diag([_BOX_NM, _BOX_NM, _BOX_NM]), unit.nanometer) + ) + return topology + + +def _get_nonbonded(system): + for force in system.getForces(): + if isinstance(force, openmm.NonbondedForce): + return force + raise AssertionError("No NonbondedForce in system") + + +def _read_particle(force, i): + charge, sigma, epsilon = force.getParticleParameters(i) + return ( + charge.value_in_unit(unit.elementary_charge), + sigma.value_in_unit(unit.nanometer), + epsilon.value_in_unit(unit.kilojoule_per_mole), + ) + + +def _read_exception(force, i): + p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(i) + return ( + int(p1), int(p2), + chargeProd.value_in_unit(unit.elementary_charge * unit.elementary_charge), + sigma.value_in_unit(unit.nanometer), + epsilon.value_in_unit(unit.kilojoule_per_mole), + ) + + +def _all_exceptions(force): + return { + tuple(sorted((p1, p2))): (cp, s, e) + for (p1, p2, cp, s, e) in (_read_exception(force, i) for i in range(force.getNumExceptions())) + } + + +def _make_mixed_system(**system_kwargs): + system = _build_periodic_system(**system_kwargs) + topology = _build_topology() + potential = MLPotential("noop_test_impl") + return potential.createMixedSystem( + topology, system, _ML_ATOMS, embedding="electrostatic" + ), system + + +# --------------------------------------------------------------------------- +# Slice 4 unit tests +# --------------------------------------------------------------------------- + +def test_pme_method_preserved(): + """The PME setting on NonbondedForce must survive the surgery.""" + mixed, _ = _make_mixed_system() + nb = _get_nonbonded(mixed) + assert nb.getNonbondedMethod() == openmm.NonbondedForce.PME + + +def test_ml_charges_zeroed(): + """ML particle charges must be zero so reciprocal-space ML-* contribs vanish.""" + mixed, _ = _make_mixed_system() + nb = _get_nonbonded(mixed) + for i in _ML_ATOMS: + charge, _, _ = _read_particle(nb, i) + assert charge == pytest.approx(0.0, abs=1e-12) + + +def test_mm_charges_unchanged(): + """MM particle charges must be untouched by the surgery.""" + mixed, original = _make_mixed_system() + nb_mixed = _get_nonbonded(mixed) + nb_orig = _get_nonbonded(original) + for i in _MM_ATOMS: + new_charge, _, _ = _read_particle(nb_mixed, i) + old_charge, _, _ = _read_particle(nb_orig, i) + assert new_charge == pytest.approx(old_charge, abs=1e-12) + + +def test_lj_parameters_preserved_for_all_particles(): + """sigma/epsilon must be untouched (LJ stays in MM force field).""" + mixed, original = _make_mixed_system() + nb_mixed = _get_nonbonded(mixed) + nb_orig = _get_nonbonded(original) + for i in range(nb_orig.getNumParticles()): + _, s_new, e_new = _read_particle(nb_mixed, i) + _, s_old, e_old = _read_particle(nb_orig, i) + assert s_new == pytest.approx(s_old, abs=1e-12) + assert e_new == pytest.approx(e_old, abs=1e-12) + + +def test_ml_mm_pairs_get_no_exception(): + """ML-MM pairs must be left on the ordinary pair list. + + ML-MM Coulomb is removed by zeroing the ML particle charges, not by adding + an exception per ML-MM pair. An exception would also be wrong: OpenMM + evaluates exceptions at the plain Cartesian distance rather than the + minimum image one, so under PBC the ML-MM Lennard-Jones interaction would + silently disappear for any pair that is only within the cutoff across a + periodic boundary. + """ + mixed, _ = _make_mixed_system(with_pre_existing_exception=False) + nb = _get_nonbonded(mixed) + excs = _all_exceptions(nb) + for ml in _ML_ATOMS: + for mm in _MM_ATOMS: + key = tuple(sorted((ml, mm))) + assert key not in excs, f"Unexpected ML-MM exception for {key}" + + +def test_ml_charges_zeroed_and_lj_untouched(): + """The Coulomb removal is done by zeroing the ML particle charges, which + leaves their Lennard-Jones parameters, and so ML-MM LJ, intact.""" + mixed, _ = _make_mixed_system(with_pre_existing_exception=False) + nb = _get_nonbonded(mixed) + for ml in _ML_ATOMS: + charge, sigma, epsilon = _read_particle(nb, ml) + assert charge == pytest.approx(0.0, abs=1e-15) + assert sigma == pytest.approx(_PARAMS[ml][2], rel=1e-12) + assert epsilon == pytest.approx(_PARAMS[ml][3], rel=1e-12) + for mm in _MM_ATOMS: + charge, sigma, epsilon = _read_particle(nb, mm) + assert charge == pytest.approx(_PARAMS[mm][1], rel=1e-12) + assert sigma == pytest.approx(_PARAMS[mm][2], rel=1e-12) + assert epsilon == pytest.approx(_PARAMS[mm][3], rel=1e-12) + + +def test_ml_mm_pre_existing_exception_keeps_lj_zeroes_charge(): + """If a pre-existing ML-MM exception had nonzero chargeProd, surgery must + zero it but keep the explicit sigma/epsilon.""" + mixed, _ = _make_mixed_system(with_pre_existing_exception=True) + nb = _get_nonbonded(mixed) + excs = _all_exceptions(nb) + cp, sigma, epsilon = excs[(1, 2)] + assert cp == pytest.approx(0.0, abs=1e-15) + assert sigma == pytest.approx(0.345, rel=1e-12) + assert epsilon == pytest.approx(0.678, rel=1e-12) + + +def test_ml_ml_exceptions_zero_charge_and_zero_lj(): + """ML-ML pairs are fully internal to MACE — both Coulomb and LJ must be + zeroed in the MM force field.""" + mixed, _ = _make_mixed_system() + nb = _get_nonbonded(mixed) + excs = _all_exceptions(nb) + key = tuple(sorted(_ML_ATOMS)) + assert key in excs + cp, _, epsilon = excs[key] + assert cp == pytest.approx(0.0, abs=1e-15) + assert epsilon == pytest.approx(0.0, abs=1e-15) + + +def test_mm_mm_exceptions_untouched(): + """MM-MM pairs must keep whatever the original system specified + (here: nothing — no exceptions added by surgery).""" + mixed, original = _make_mixed_system() + nb_mixed = _get_nonbonded(mixed) + nb_orig = _get_nonbonded(original) + excs_orig = _all_exceptions(nb_orig) + excs_mixed = _all_exceptions(nb_mixed) + mm_key = tuple(sorted(_MM_ATOMS)) + assert excs_mixed.get(mm_key) == excs_orig.get(mm_key) + + +def test_boundary_bonded_terms_preserved(): + """Boundary bonded terms (HarmonicBond/Angle that connect ML-MM) must + survive — they remain classical in electrostatic embedding.""" + mixed, _ = _make_mixed_system(with_boundary_bonded=True) + bond_force = next(f for f in mixed.getForces() if isinstance(f, openmm.HarmonicBondForce)) + bond_pairs = set() + for i in range(bond_force.getNumBonds()): + p1, p2, _, _ = bond_force.getBondParameters(i) + bond_pairs.add(tuple(sorted((int(p1), int(p2))))) + # ML-MM boundary bond and MM-MM bond must remain. + assert (1, 2) in bond_pairs + assert (2, 3) in bond_pairs + # ML-internal bond must be removed. + assert (0, 1) not in bond_pairs + + angle_force = next(f for f in mixed.getForces() if isinstance(f, openmm.HarmonicAngleForce)) + angle_triples = set() + for i in range(angle_force.getNumAngles()): + p1, p2, p3, _, _ = angle_force.getAngleParameters(i) + angle_triples.add((int(p1), int(p2), int(p3))) + # Boundary angles (touching at least one MM atom) preserved. + assert (0, 1, 2) in angle_triples + assert (1, 2, 3) in angle_triples + + +def test_reciprocal_space_pme_matches_mm_only_reference(): + """PME reciprocal-space contribution from the mixed-system NonbondedForce + must equal that of an MM-only reference (where ML charges are explicitly + zeroed but the system is otherwise identical).""" + mixed, original = _make_mixed_system(with_boundary_bonded=False) + + # Strip the bonded forces so only NonbondedForce contributes (we are + # comparing the reciprocal-space + direct-space Coulomb/LJ summed term). + def _strip_to_nonbonded(system): + # Remove non-NonbondedForce forces by index, descending. + for idx in reversed(range(system.getNumForces())): + if not isinstance(system.getForce(idx), openmm.NonbondedForce): + system.removeForce(idx) + return system + + _strip_to_nonbonded(mixed) + + # Build MM-only reference by zeroing ML charges in a fresh original system. + mm_only = _build_periodic_system(with_boundary_bonded=False) + nb_ref = _get_nonbonded(mm_only) + for i in _ML_ATOMS: + _, sigma, epsilon = nb_ref.getParticleParameters(i) + nb_ref.setParticleParameters(i, 0.0 * unit.elementary_charge, sigma, epsilon) + # Add the same exception structure (LB ML-MM zero-charge + ML-ML LJ-zero) + # so that direct-space exclusion subtractions match between the two systems. + # That is what the surgery itself produced; copy it onto the reference. + nb_mixed = _get_nonbonded(mixed) + existing_pairs = set( + tuple(sorted((int(p1), int(p2)))) + for i in range(nb_ref.getNumExceptions()) + for (p1, p2, *_rest) in [nb_ref.getExceptionParameters(i)] + ) + for i in range(nb_mixed.getNumExceptions()): + p1, p2, cp, sigma, epsilon = nb_mixed.getExceptionParameters(i) + if tuple(sorted((int(p1), int(p2)))) in existing_pairs: + continue + nb_ref.addException(int(p1), int(p2), cp, sigma, epsilon) + _strip_to_nonbonded(mm_only) + + positions = [ + openmm.Vec3(0.20, 0.30, 0.40), + openmm.Vec3(0.45, 0.30, 0.40), + openmm.Vec3(0.70, 0.30, 0.40), + openmm.Vec3(0.80, 0.30, 0.40), + ] * unit.nanometer + + platform = openmm.Platform.getPlatformByName("Reference") + ctx_mixed = openmm.Context(mixed, openmm.VerletIntegrator(0.001), platform) + ctx_ref = openmm.Context(mm_only, openmm.VerletIntegrator(0.001), platform) + ctx_mixed.setPositions(positions) + ctx_ref.setPositions(positions) + + e_mixed = ctx_mixed.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) + e_ref = ctx_ref.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) + + assert e_mixed == pytest.approx(e_ref, rel=1e-6, abs=1e-6) + + +# --------------------------------------------------------------------------- +# Electrostatics the surgery cannot account for. +# +# The ML-MM Coulomb is removed from the force field on the understanding that +# the model supplies it. Any Coulomb term this method cannot find is either left +# in place and counted twice, or removed and never replaced. Both give a wrong +# energy with no error, so these cases are refused rather than guessed at. +# --------------------------------------------------------------------------- + + +def _minimal_topology(numParticles): + from openmm.app import element + + topology = app.Topology() + chain = topology.addChain() + residue = topology.addResidue("X", chain) + for i in range(numParticles): + topology.addAtom(f"H{i}", element.hydrogen, residue) + return topology + + +def _createMixedSystem(system, **args): + impl = _NoopImpl() + return impl.createMixedSystem( + _minimal_topology(system.getNumParticles()), system, _ML_ATOMS, 0, False, + "electrostatic", **args, + ) + + +def test_multiple_nonbonded_forces_rejected(): + """The MM charges given to the model come from one NonbondedForce, so more + than one is ambiguous: the surgery would zero ML charges in all of them + while the model saw the charges of only the first.""" + system = _build_periodic_system() + system.addForce(openmm.NonbondedForce()) + with pytest.raises(ValueError, match="Multiple NonbondedForce"): + _createMixedSystem(system) + + +def _addCoulombCustomNonbondedForce(system): + force = openmm.CustomNonbondedForce("138.935456*q1*q2/r") + force.addPerParticleParameter("q") + for index in range(system.getNumParticles()): + force.addParticle([_PARAMS[index][1]]) + system.addForce(force) + return system + + +def test_custom_nonbonded_force_requires_an_answer(): + """A CustomNonbondedForce's energy expression is arbitrary, so whether it + carries electrostatics cannot be determined here and must be declared.""" + system = _addCoulombCustomNonbondedForce(_build_periodic_system()) + with pytest.raises(ValueError, match="unknown whether it includes electrostatic"): + _createMixedSystem(system) + + +def test_custom_nonbonded_force_with_charges_needs_the_parameter_name(): + """Declaring that it does carry electrostatics is not enough on its own: + the charge cannot be zeroed without knowing which parameter holds it.""" + system = _addCoulombCustomNonbondedForce(_build_periodic_system()) + with pytest.raises(ValueError, match="must name the per-particle parameter"): + _createMixedSystem(system, customNonbondedHasCharges=True) + + +def test_custom_nonbonded_force_unknown_parameter_name_rejected(): + """Naming a parameter the force does not define is an error, not a no-op + that would leave the electrostatics in place.""" + system = _addCoulombCustomNonbondedForce(_build_periodic_system()) + with pytest.raises(ValueError, match="no per-particle parameter"): + _createMixedSystem(system, customNonbondedHasCharges=True, + customNonbondedChargeParameter="charge") + + +def test_custom_nonbonded_charge_parameter_zeroed_on_ml_atoms(): + """Naming the charge parameter zeroes it on the ML atoms, which removes the + ML-MM Coulomb the custom force would otherwise still contribute.""" + system = _addCoulombCustomNonbondedForce(_build_periodic_system()) + mixed = _createMixedSystem(system, customNonbondedHasCharges=True, + customNonbondedChargeParameter="q") + custom = next(f for f in mixed.getForces() if isinstance(f, openmm.CustomNonbondedForce)) + for atom in _ML_ATOMS: + assert custom.getParticleParameters(atom)[0] == pytest.approx(0.0, abs=1e-12) + for atom in _MM_ATOMS: + assert custom.getParticleParameters(atom)[0] == pytest.approx(_PARAMS[atom][1], rel=1e-12) + + +def test_custom_nonbonded_force_without_charges_accepted(): + """Declaring it carries none proceeds, with ML-ML excluded as usual.""" + system = _addCoulombCustomNonbondedForce(_build_periodic_system()) + mixed = _createMixedSystem(system, customNonbondedHasCharges=False) + custom = next(f for f in mixed.getForces() if isinstance(f, openmm.CustomNonbondedForce)) + exclusions = { + tuple(sorted(custom.getExclusionParticles(i))) + for i in range(custom.getNumExclusions()) + } + assert tuple(sorted(_ML_ATOMS)) in exclusions + + +def test_custom_nonbonded_charges_reach_the_model(): + """When the electrostatics live in the CustomNonbondedForce, the charges the + model is given must come from there too. + + Reading them from the NonbondedForce in that case hands the model zeros + while the surgery has already removed the real ML-MM Coulomb, so the + interaction disappears rather than being computed by the model. + """ + from openmmml.models.macepotential import _prepareMMEmbedding + + system = openmm.System() + for _ in range(4): + system.addParticle(1.0) + lj = openmm.NonbondedForce() # Lennard-Jones only, no charges + for _, _, sigma, epsilon in _PARAMS: + lj.addParticle(0.0, sigma, epsilon) + system.addForce(lj) + coulomb = openmm.CustomNonbondedForce("138.935456*q1*q2/r") + coulomb.addPerParticleParameter("q") + for _, charge, _, _ in _PARAMS: + coulomb.addParticle([charge]) + system.addForce(coulomb) + + expected = [_PARAMS[i][1] for i in _MM_ATOMS] + charges = _prepareMMEmbedding(system, _ML_ATOMS, "q")["mm_charges"] + np.testing.assert_allclose(charges, expected, atol=1e-12) diff --git a/test/TestLinkChargeRedistribution.py b/test/TestLinkChargeRedistribution.py new file mode 100644 index 0000000..18cc498 --- /dev/null +++ b/test/TestLinkChargeRedistribution.py @@ -0,0 +1,222 @@ +"""Unit tests for Z1 / DZ1 link-atom charge redistribution. + +Covers the new ``apply_link_charge_redistribution`` helper in +``openmmml.models._links``. Does not exercise the full MACE stack — +that's covered by the existing electrostatic-embedding smoke tests. +""" +import numpy as np +import openmm as mm +import openmm.app as app +from openmm.app import element as elem +import pytest + +from openmmml.models._links import apply_link_charge_redistribution + + +# ----------------------------------------------------------------------------- +# Helpers: build minimal Topology fixtures. +# ----------------------------------------------------------------------------- + +def _topology_with_bonds(bonds: list[tuple[int, int]], n_atoms: int = 4): + """Tiny topology: a single chain with the requested bonds. + + All atoms are CARBON (element doesn't matter for these tests). + """ + topo = app.Topology() + chain = topo.addChain() + residue = topo.addResidue("X", chain) + atoms = [topo.addAtom(f"C{i}", elem.carbon, residue) for i in range(n_atoms)] + for a, b in bonds: + topo.addBond(atoms[a], atoms[b]) + return topo + + +def _toy_linkinfo(q_globals, m_globals): + """Minimum link_info dict expected by apply_link_charge_redistribution.""" + return { + "q_global": np.asarray(q_globals, dtype=np.int64), + "m_global": np.asarray(m_globals, dtype=np.int64), + "target_dist": np.full(len(q_globals), 1.09, dtype=np.float64), + } + + +# ----------------------------------------------------------------------------- +# Basic API +# ----------------------------------------------------------------------------- + +def test_none_returns_copy_unchanged(): + mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) + mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10], [1]), + topology=_topology_with_bonds([]), scheme="none", + ) + np.testing.assert_array_equal(out, mm_q) + assert out is not mm_q # must be a copy + + +def test_unknown_scheme_raises(): + with pytest.raises(ValueError, match="Unsupported linkChargeScheme"): + apply_link_charge_redistribution( + np.array([0]), np.array([0.0]), + _toy_linkinfo([], []), + topology=_topology_with_bonds([]), scheme="z3", + ) + + +# ----------------------------------------------------------------------------- +# Z1 +# ----------------------------------------------------------------------------- + +def test_z1_zeros_M_only(): + mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) + mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10], [1]), + topology=_topology_with_bonds([]), scheme="z1", + ) + assert out[1] == 0.0 + np.testing.assert_array_equal(out[[0, 2, 3]], mm_q[[0, 2, 3]]) + + +def test_z1_multiple_M_atoms(): + mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) + mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10, 11], [1, 3]), + topology=_topology_with_bonds([]), scheme="z1", + ) + assert out[1] == 0.0 + assert out[3] == 0.0 + np.testing.assert_array_equal(out[[0, 2]], mm_q[[0, 2]]) + + +# ----------------------------------------------------------------------------- +# DZ1 +# ----------------------------------------------------------------------------- + +def test_dz1_distributes_to_MM_neighbors_and_preserves_total(): + # Bonds: M (atom 1) is bonded to MM atoms 2 and 3, plus Q atom 10. + # The Q-M bond should be ignored (not an MM-MM bond). + # q_M_orig = -0.2 -> each MM neighbour gets -0.1. + topology = _topology_with_bonds([(1, 2), (1, 3)], n_atoms=4) + mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) + mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10], [1]), + topology=topology, scheme="dz1", + ) + assert out[1] == 0.0 + assert out[2] == pytest.approx(0.3 + (-0.2) / 2) + assert out[3] == pytest.approx(-0.4 + (-0.2) / 2) + # Total preserved to round-off + assert out.sum() == pytest.approx(mm_q.sum(), abs=1e-12) + + +def test_dz1_skips_Q_neighbors(): + # M is bonded to Q (10) AND to MM (2). Only the MM bond counts. + # q_M_orig = -0.5 -> atom 2 gets +(-0.5)/1 = -0.5. + topology = _topology_with_bonds([(1, 2), (1, 10)], n_atoms=11) + mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) + mm_q = np.array([0.1, -0.5, 0.2, 0.2], dtype=np.float64) + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10], [1]), + topology=topology, scheme="dz1", + ) + assert out[1] == 0.0 + assert out[2] == pytest.approx(0.2 - 0.5) + # atom 3 unchanged + assert out[3] == pytest.approx(0.2) + + +def test_dz1_no_MM_neighbors_warns_and_falls_back_to_Z1(): + # M (atom 1) has no MM bonds — only bonded to Q (atom 10). + topology = _topology_with_bonds([(1, 10)], n_atoms=11) + mm_atoms = np.array([0, 1], dtype=np.int64) + mm_q = np.array([0.5, -0.5], dtype=np.float64) + with pytest.warns(UserWarning, match="no MM neighbours"): + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10], [1]), + topology=topology, scheme="dz1", + ) + assert out[1] == 0.0 + # Total not preserved here — warning told us. + assert out.sum() != pytest.approx(mm_q.sum(), abs=1e-12) + + +def test_dz1_multiple_M_atoms_each_get_own_neighbors(): + # M atoms = [1, 3]. Each is bonded to a different MM neighbour (0 and 2). + # q_orig: atom 1 = -0.4 -> atom 0 += -0.4 + # atom 3 = +0.6 -> atom 2 += +0.6 + topology = _topology_with_bonds([(1, 0), (3, 2)], n_atoms=4) + mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) + mm_q = np.array([0.0, -0.4, 0.0, 0.6], dtype=np.float64) + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10, 11], [1, 3]), + topology=topology, scheme="dz1", + ) + assert out[1] == 0.0 + assert out[3] == 0.0 + assert out[0] == pytest.approx(-0.4) + assert out[2] == pytest.approx(0.6) + assert out.sum() == pytest.approx(mm_q.sum(), abs=1e-12) + + +def test_dz1_split_across_three_neighbors(): + # M (atom 1) has three MM neighbours (0, 2, 3). q_M_orig=-0.3 -> each +(-0.1). + topology = _topology_with_bonds([(1, 0), (1, 2), (1, 3)], n_atoms=4) + mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) + mm_q = np.array([0.0, -0.3, 0.0, 0.0], dtype=np.float64) + out = apply_link_charge_redistribution( + mm_atoms, mm_q, _toy_linkinfo([10], [1]), + topology=topology, scheme="dz1", + ) + assert out[1] == 0.0 + for idx in (0, 2, 3): + assert out[idx] == pytest.approx(-0.1) + assert out.sum() == pytest.approx(mm_q.sum(), abs=1e-12) + + +def test_dz1_with_two_bonded_M_atoms(): + """Two M atoms bonded to each other must both end up at exactly zero. + + Redistributing onto a neighbour that is itself an M atom hands it charge + that its own turn then discards, so the outcome depended on the order the + link records were listed in and left one M atom holding charge -- the very + charge DZ1 exists to move away from the ML region. + """ + import numpy as np + import openmm.app as app + from openmm.app import element + + topology = app.Topology() + chain = topology.addChain() + residue = topology.addResidue("X", chain) + atoms = [topology.addAtom(f"H{i}", element.hydrogen, residue) for i in range(6)] + # ML = {0, 1}; caps Q0->M2 and Q1->M3; M2 and M3 are bonded to each other. + for i, j in [(2, 3), (2, 4), (3, 5)]: + topology.addBond(atoms[i], atoms[j]) + + mm_atoms = np.array([2, 3, 4, 5]) + mm_charges = np.array([-0.5, -0.3, 0.2, 0.1]) + + out = apply_link_charge_redistribution( + mm_atoms=mm_atoms, + mm_charges=mm_charges, + link_info={"q_global": np.array([0, 1]), "m_global": np.array([2, 3])}, + topology=topology, + scheme="dz1", + ) + + assert out[0] == 0.0 and out[1] == 0.0 + assert out.sum() == pytest.approx(mm_charges.sum(), abs=1e-12) + + # And the answer must not depend on the order of the link records. + reversed_order = apply_link_charge_redistribution( + mm_atoms=mm_atoms, + mm_charges=mm_charges, + link_info={"q_global": np.array([1, 0]), "m_global": np.array([3, 2])}, + topology=topology, + scheme="dz1", + ) + np.testing.assert_allclose(out, reversed_order, atol=1e-12) diff --git a/test/TestLinkRecordsTargetDistUnits.py b/test/TestLinkRecordsTargetDistUnits.py new file mode 100644 index 0000000..42143ac --- /dev/null +++ b/test/TestLinkRecordsTargetDistUnits.py @@ -0,0 +1,147 @@ +"""Regression test for the `target_dist` unit canonicalisation in +`_prepareLinkRecords` and its downstream consumers. + +The historical bug: when `linkRecords` was loaded from a +`capping_mapping.csv`, the loader +multiplied the CSV's `target_dist_ang` column by 0.1 — turning a 1.09 Å +record into a 0.109 nm record stored in `linkInfo["target_dist"]`. But +`_computeMACE` reads `linkInfo["target_dist"]` together with `r_Q`, +`r_M` already in Å (via `positions_full = state.getPositions(asNumpy=True +).value_in_unit(unit.angstrom)`). The unit mismatch placed CSV-path +caps at ~0.1 Å from the Q atom — about 10× closer than the intended +~1.09 Å (typical C-H bond). + +The fix: `linkInfo["target_dist"]` is canonically in **Å**. The CSV path +no longer multiplies by 0.1; the tuple path was already in Å. + +Tests in this file pin the contract at three levels so a regression on +*any* of them surfaces clearly: + +1. Tuple-path loader stores the value verbatim (Å). +2. CSV-path loader stores the value verbatim (Å, no `* 0.1` conversion). +3. Downstream end-to-end: `compute_cap_positions` placed against + linkInfo + Å-scale Q/M positions gives a Q-L distance equal to the + target value. Any future "fix" that re-introduces `* 0.1` at the + consumer side (instead of the loader side) trips this end-to-end + case even if it leaves the loader unchanged. +""" +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + + +@pytest.fixture +def tiny_system_topology(): + """Minimal 5-atom system + topology so _prepareLinkRecords' validation + (q in atoms, m not in atoms, etc.) is satisfied. Atoms 0,1,2 are ML; + atoms 3,4 are MM.""" + import openmm as mm + import openmm.app as app + from openmm.app import element as elem + + system = mm.System() + for _ in range(5): + system.addParticle(1.0) + top = app.Topology() + chain = top.addChain() + res = top.addResidue("X", chain) + atoms = [top.addAtom(f"H{i}", elem.hydrogen, res) for i in range(5)] + for a, b in [(0, 1), (1, 2), (2, 3), (3, 4)]: + top.addBond(atoms[a], atoms[b]) + return top, system + + +def test_tuple_path_target_dist_is_angstrom(tiny_system_topology): + """Tuple path: linkInfo['target_dist'] equals the value the caller passed. + + Locks the canonical interpretation as Å. + """ + from openmmml.models.macepotential import _prepareLinkRecords + + top, system = tiny_system_topology + link_info = _prepareLinkRecords( + linkRecords=[(2, 3, 1.09)], + atoms=[0, 1, 2], + topology=top, + system=system, + ) + assert link_info is not None + assert link_info["K"] == 1 + np.testing.assert_allclose(link_info["target_dist"], np.array([1.09])) + + +# The CSV-path test requires the MLMM-side `load_link_records` parser. +# openmm-ml proper does not depend on MLMM (it's an optional consumer), +# so the test skips cleanly when `mlmm` isn't on the path — keeping +# openmm-ml's own CI (which doesn't install MLMM) reproducible. +_mlmm_available = pytest.importorskip.__doc__ is not None # always True; placeholder +try: + import mlmm.link_atoms # noqa: F401 + _HAS_MLMM = True +except Exception: + _HAS_MLMM = False + + +@pytest.mark.skipif(not _HAS_MLMM, reason="MLMM not installed; CSV path needs mlmm.link_atoms.load_link_records") +def test_csv_path_target_dist_in_angstrom(tmp_path, tiny_system_topology): + """CSV path: linkInfo['target_dist'] matches the CSV's target_dist_ang + column verbatim (no spurious 0.1 conversion). + + Historical bug had this multiplied by 0.1, producing + linkInfo['target_dist'] = 0.109 for a CSV target_dist_ang of 1.09, + which then placed MACE caps 10x too close to Q. + """ + from openmmml.models.macepotential import _prepareLinkRecords + + csv = tmp_path / "capping_mapping.csv" + csv.write_text( + "q_idx1,m_idx1,q_element,m_element,bond_order,qm_mm_distance_ang," + "target_dist_ang,r_Q_x,r_Q_y,r_Q_z,r_M_x,r_M_y,r_M_z,r_L_x,r_L_y,r_L_z\n" + "3,4,H,H,1.000,1.500000,1.090000," + "0.0,0.0,0.0,1.5,0.0,0.0,1.09,0.0,0.0\n" + ) + + top, system = tiny_system_topology + link_info = _prepareLinkRecords( + linkRecords=str(csv), + atoms=[0, 1, 2], + topology=top, + system=system, + ) + assert link_info is not None + assert link_info["K"] == 1 + np.testing.assert_allclose( + link_info["target_dist"], + np.array([1.09]), + err_msg=( + "CSV target_dist_ang=1.09 must land in linkInfo as 1.09 (Å). " + "If you see 0.109 here, the historical *0.1 nm conversion has " + "regressed; see openmm-ml fix/target-dist-unit-canonical-angstrom." + ), + ) + + +def test_compute_cap_positions_q_to_l_distance_matches_target(): + """End-to-end contract: with linkInfo['target_dist'] in Å and Q,M + positions in Å, the cap placement helper produces a Q-L distance + equal to target_dist Å. + + This is the consumer-side test codex review of PR #20 called out: a + future regression that re-introduces ``* 0.1`` at the consumer side + (e.g. ``compute_cap_positions(r_Q, r_M, linkInfo['target_dist'] * 0.1)``) + would slip past the loader-only test but trip this one. + """ + from openmmml.models._links import compute_cap_positions + + # Q at origin, M at (1.5, 0, 0) — a typical C-C single bond, 1.5 Å. + r_Q = np.array([[0.0, 0.0, 0.0]]) + r_M = np.array([[1.5, 0.0, 0.0]]) + target_dist_ang = np.array([1.09]) # typical C-H, in Å + + cap_pos, C_L = compute_cap_positions(r_Q, r_M, target_dist_ang) + q_to_l = np.linalg.norm(cap_pos - r_Q, axis=-1) + np.testing.assert_allclose(q_to_l, target_dist_ang, atol=1e-12) + np.testing.assert_allclose(C_L, target_dist_ang / 1.5, atol=1e-12) diff --git a/test/TestMACEPotential.py b/test/TestMACEPotential.py index 4470278..8b8e729 100644 --- a/test/TestMACEPotential.py +++ b/test/TestMACEPotential.py @@ -5,14 +5,159 @@ import openmm.app as app import openmm.unit as unit import pytest +import torch from openmmml import MLPotential +from openmmml.models.macepotential import ( + MACEPotentialImpl, + _computeMACE, + _prepareMMEmbedding, + _should_use_mm_embedding, +) mace = pytest.importorskip("mace", reason="mace is not installed") platform_ints = range(mm.Platform.getNumPlatforms()) # Get the path to the test data test_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") + +def _simple_nonbonded_system(): + system = mm.System() + nonbonded = mm.NonbondedForce() + params = [ + (1.0, 0.30, 0.20), + (-0.5, 0.40, 0.50), + (0.25, 0.50, 0.80), + ] + for charge_e, sigma_nm, epsilon_kj in params: + system.addParticle(12.0) + nonbonded.addParticle( + charge_e * unit.elementary_charge, + sigma_nm * unit.nanometer, + epsilon_kj * unit.kilojoule_per_mole, + ) + system.addForce(nonbonded) + return system, nonbonded + + +def testPrepareMMEmbedding(): + """The MM complement and its charges are read off the NonbondedForce.""" + system, nonbonded = _simple_nonbonded_system() + info = _prepareMMEmbedding(system, [0, 1]) + assert info is not None + np.testing.assert_array_equal(info["ml_atoms"], [0, 1]) + np.testing.assert_array_equal(info["mm_atoms"], [2]) + np.testing.assert_allclose(info["mm_charges"], [0.25], atol=1e-12) + + +class _FakeState: + def __init__(self, positions_angstrom): + self._positions = np.asarray(positions_angstrom, dtype=np.float64) * unit.angstrom + + def getPositions(self, asNumpy=False): + return self._positions + + +class _FakeModel: + def __init__(self, dtype=torch.float32): + self.r_max = torch.tensor(3.0, dtype=dtype) + self.dtype = dtype + + def __call__(self, input_dict, compute_force=True): + del compute_force + n_ml = input_dict["positions"].shape[0] + forces = torch.tensor( + [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]][:n_ml], + dtype=self.dtype, + device=input_dict["positions"].device, + ) + out = { + "interaction_energy": torch.tensor([2.5], dtype=self.dtype, device=forces.device), + "forces": forces, + } + if "mm_positions" in input_dict: + out["mm_forces"] = torch.tensor( + [[0.5, 0.25, -0.75]], + dtype=self.dtype, + device=forces.device, + ) + return out + + +class PolarMACE: + supports_external_electrostatics = True + + +class MACE: + pass + + +def testComputeMACEScattersMMForces(): + ptr = torch.tensor([0, 2], dtype=torch.long) + node_attrs = torch.ones((2, 1), dtype=torch.float32) + batch = torch.zeros(2, dtype=torch.long) + pbc = torch.tensor([False, False, False], dtype=torch.bool) + mm_info = {"mm_atoms": np.array([2], dtype=np.int64), "mm_charges": np.array([0.25], dtype=np.float64)} + + energy, forces = _computeMACE( + state=_FakeState([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]]), + model=_FakeModel(dtype=torch.float32), + ptr=ptr, + node_attrs=node_attrs, + batch=batch, + pbc=pbc, + returnEnergyType="interaction_energy", + charge=torch.tensor([0.0], dtype=torch.float32), + multiplicity=torch.tensor([1.0], dtype=torch.float32), + indices=np.array([0, 1], dtype=np.int64), + periodic=False, + linkInfo=None, + mmInfo=mm_info, + ) + + assert energy == pytest.approx(2.5 * 96.4853) + expected = np.array( + [ + [1.0, 2.0, 3.0], + [-1.0, -2.0, -3.0], + [0.5, 0.25, -0.75], + ], + dtype=np.float32, + ) * (96.4853 * 10.0) + np.testing.assert_allclose(forces, expected, rtol=1e-6, atol=1e-6) + + +def testShouldUseMMEmbedding(): + assert _should_use_mm_embedding(PolarMACE(), [0, 1], "electrostatic") + assert not _should_use_mm_embedding(PolarMACE(), [0, 1], "mechanical") + assert not _should_use_mm_embedding(MACE(), [0, 1], "mechanical") + with pytest.raises(ValueError, match="Unsupported embedding mode"): + _should_use_mm_embedding(PolarMACE(), [0, 1], "bad-mode") + + +def testElectrostaticEmbeddingRejectsUnsupportedCases(): + """Electrostatic embedding must fail loudly rather than falling back to + mechanical embedding: by the time the potential is asked to add its forces, + the ML-MM electrostatics have already been removed from the MM force field, + so a silent fallback would simply lose them.""" + with pytest.raises(ValueError, match="requires a model that accepts MM charges"): + _should_use_mm_embedding(MACE(), [0, 1], "electrostatic") + with pytest.raises(ValueError, match="requires an ML subset"): + _should_use_mm_embedding(PolarMACE(), None, "electrostatic") + + +def testOnlyPolarModelsOfferElectrostaticEmbedding(): + """Electrostatic embedding needs a model that accepts MM charges, which of + the pretrained models only the PolarMACE family does. A custom checkpoint + may be one, so it is offered there too and checked once the model loads.""" + for name, (_, _, _, _, acceptsMMCharges) in MACEPotentialImpl.KNOWN_MODELS.items(): + embeddings = MLPotential(name).getSupportedEmbeddings() + assert ('electrostatic' in embeddings) == acceptsMMCharges, name + assert acceptsMMCharges == name.startswith('mace-polar'), name + custom = MLPotential('mace', modelPath='unused-until-forces-are-added.model') + assert 'electrostatic' in custom.getSupportedEmbeddings() + + @pytest.mark.parametrize("platform_int", list(platform_ints)) class TestMACE: @@ -105,3 +250,35 @@ def testPrecisionApplied(self, platform_int, precision): "Energy is not finite under precision {}".format(precision) assert np.isclose(energyML, self.refEnergy['mace-off23-small'], rtol=1e-6),\ "Energy is not close to reference under precision {}".format(precision) + + +def testMissingMMForcesIsAnError(): + """A model given MM charges must return forces on the MM atoms.""" + class _NoMMForces(_FakeModel): + def __call__(self, inputDict, compute_force=True): + out = super().__call__(inputDict, compute_force) + out.pop("mm_forces", None) + return out + + mmInfo = { + "mm_atoms": np.array([2], dtype=np.int64), + "mm_charges": np.array([0.25], dtype=np.float64), + } + with pytest.raises(ValueError, match="no 'mm_forces'"): + _computeMACE( + state=_FakeState( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]] + ), + model=_NoMMForces(dtype=torch.float32), + ptr=torch.tensor([0, 2], dtype=torch.long), + node_attrs=torch.ones((2, 1), dtype=torch.float32), + batch=torch.zeros(2, dtype=torch.long), + pbc=torch.tensor([False, False, False], dtype=torch.bool), + returnEnergyType="interaction_energy", + charge=torch.tensor([0.0], dtype=torch.float32), + multiplicity=torch.tensor([1.0], dtype=torch.float32), + indices=np.array([0, 1], dtype=np.int64), + periodic=False, + linkInfo=None, + mmInfo=mmInfo, + ) diff --git a/test/TestMechanicalEmbedding.py b/test/TestMechanicalEmbedding.py index ac01f50..0257437 100644 --- a/test/TestMechanicalEmbedding.py +++ b/test/TestMechanicalEmbedding.py @@ -489,3 +489,40 @@ def testLinkAtomInfo(self, platform_int): assert mixed_info["oldToNew"] == list(range(original_count)) for i in range(mixed_system.getNumParticles()): assert mixed_system.isVirtualSite(i) == (i >= original_count) + + +def testCustomNonbondedForce(): + """Mechanical embedding excludes ML-ML CustomNonbondedForce pairs.""" + numParticles = 4 + mlAtoms = [0, 1] + system = openmm.System() + nonbonded = openmm.NonbondedForce() + custom = openmm.CustomNonbondedForce("4*epsilon*((sigma/r)^12-(sigma/r)^6)") + custom.addPerParticleParameter("sigma") + custom.addPerParticleParameter("epsilon") + for _ in range(numParticles): + system.addParticle(1.0) + nonbonded.addParticle(0.0, 0.3, 0.2) + custom.addParticle([0.3, 0.2]) + system.addForce(nonbonded) + system.addForce(custom) + + topology = openmm.app.Topology() + chain = topology.addChain() + residue = topology.addResidue("X", chain) + for i in range(numParticles): + topology.addAtom(f"H{i}", openmm.app.element.hydrogen, residue) + + mixed = MLPotential("mace-off23-small").createMixedSystem( + topology, system, mlAtoms + ) + mixedCustom = next( + force + for force in mixed.getForces() + if isinstance(force, openmm.CustomNonbondedForce) + ) + exclusions = { + tuple(sorted(mixedCustom.getExclusionParticles(i))) + for i in range(mixedCustom.getNumExclusions()) + } + assert tuple(sorted(mlAtoms)) in exclusions diff --git a/test/TestPolarMACEElectrostaticEmbedding.py b/test/TestPolarMACEElectrostaticEmbedding.py new file mode 100644 index 0000000..0bfa01e --- /dev/null +++ b/test/TestPolarMACEElectrostaticEmbedding.py @@ -0,0 +1,658 @@ +"""End-to-end tests for PolarMACE electrostatic embedding via openmmml. + +Covers both: + +* Non-periodic plumbing: MM positions/charges flow into PolarMACE, + electrostatic embedding shifts the ML energy and creates a back-reaction + force on the MM atoms (Slice 4 cross-check, fills the gap that the only + existing electrostatic-embedding tests in TestMACEPotential.py exercise + the helpers in isolation rather than the full plumbing). + +* PBC: ``MLPotential('mace', modelPath=...).createMixedSystem( + embedding='electrostatic')`` is translation-invariant under PBC -- + shifting all positions by one full box vector must leave the potential + energy and per-atom forces unchanged. (Slice 5 of + docs/codex-plans/pbc-electrostatic-embedding-small-plan.md.) + +The tests train nothing -- a tiny random PolarMACE is built and serialized +in a module-scoped fixture so the suite runs in seconds on CPU. +""" +from __future__ import annotations + +import os +import tempfile + +import numpy as np +import openmm +import openmm.app as app +import openmm.unit as unit +import pytest +import torch + +torch.serialization.add_safe_globals([slice]) + +mace = pytest.importorskip("mace", reason="mace is not installed") + +# OpenMM-ML installs its external-source adapter when this checkpoint is loaded. +# The checkpoint itself stays a stock PolarMACE and does not need the mlmm +# monkey-patch that these tests previously imported. + +from e3nn import o3 # noqa: E402 + +from openmmml import MLPotential # noqa: E402 +from openmmml.models import macepotential # noqa: E402 +from mace.modules import interaction_classes # noqa: E402 +from mace.modules.extensions import PolarMACE # noqa: E402 + + +_BOX_NM = 1.4 # 14 Angstrom box (large enough for r_max=4 A and PME) +_DTYPE = torch.float64 + + +# --------------------------------------------------------------------------- +# PolarMACE model + save to tempfile +# --------------------------------------------------------------------------- + +def _build_polar_mace(device: torch.device, dtype: torch.dtype) -> PolarMACE: + fixedpoint_update_config = { + "type": "AgnosticEmbeddedOneBodyVariableUpdate", + "potential_embedding_cls": "AgnosticChargeBiasedLinearPotentialEmbedding", + "nonlinearity_cls": "MLPNonLinearity", + } + field_readout_config = {"type": "OneBodyMLPFieldReadout"} + return PolarMACE( + r_max=4.0, + num_bessel=4, + num_polynomial_cutoff=3, + max_ell=1, + interaction_cls=interaction_classes[ + "RealAgnosticResidualNonLinearInteractionBlock" + ], + interaction_cls_first=interaction_classes[ + "RealAgnosticResidualNonLinearInteractionBlock" + ], + num_interactions=2, + num_elements=2, + hidden_irreps=o3.Irreps("4x0e + 4x1o"), + MLP_irreps=o3.Irreps("8x0e"), + atomic_energies=torch.zeros(2, dtype=dtype, device=device), + avg_num_neighbors=3.0, + atomic_numbers=[1, 8], + correlation=1, + gate=torch.nn.functional.silu, + radial_MLP=[16, 16], + radial_type="bessel", + kspace_cutoff_factor=1.0, + atomic_multipoles_max_l=1, + atomic_multipoles_smearing_width=1.0, + field_feature_max_l=1, + field_feature_widths=[1.0], + field_feature_norms=[1.0, 1.0], + num_recursion_steps=1, + field_si=False, + include_electrostatic_self_interaction=False, + add_local_electron_energy=True, + field_dependence_type="AgnosticEmbeddedOneBodyVariableUpdate", + final_field_readout_type="OneBodyMLPFieldReadout", + return_electrostatic_potentials=False, + heads=["Default"], + field_norm_factor=1.0, + fixedpoint_update_config=fixedpoint_update_config, + field_readout_config=field_readout_config, + ).to(device=device, dtype=dtype) + + +@pytest.fixture(scope="module") +def polar_mace_model_path(): + torch.manual_seed(7) + model = _build_polar_mace(torch.device("cpu"), _DTYPE) + model.eval() + tmpdir = tempfile.mkdtemp(prefix="polar_mace_pbc_test_") + path = os.path.join(tmpdir, "polar_mace_test.pt") + torch.save(model, path) + yield path + try: + os.remove(path) + os.rmdir(tmpdir) + except OSError: + pass + + +# --------------------------------------------------------------------------- +# Periodic OpenMM topology + system: 1 ML water + N MM waters +# --------------------------------------------------------------------------- + +# TIP3P-ish charges and LJ for the test (units consistent with OpenMM defaults). +_O_CHARGE = -0.834 +_H_CHARGE = 0.417 +_O_SIGMA_NM = 0.31507 +_O_EPS_KJ = 0.6364 +_H_SIGMA_NM = 1.0e-3 +_H_EPS_KJ = 0.0 + + +def _add_water(system, nonbonded, bonds, angles, masses_charges_lj): + """Append a 3-atom water to system+forces; returns particle indices.""" + indices = [] + for mass, charge, sigma, epsilon in masses_charges_lj: + idx = system.addParticle(mass) + nonbonded.addParticle( + charge * unit.elementary_charge, + sigma * unit.nanometer, + epsilon * unit.kilojoule_per_mole, + ) + indices.append(idx) + o, h1, h2 = indices + bonds.addBond(o, h1, 0.09572 * unit.nanometer, 4.5e5 * unit.kilojoule_per_mole / unit.nanometer ** 2) + bonds.addBond(o, h2, 0.09572 * unit.nanometer, 4.5e5 * unit.kilojoule_per_mole / unit.nanometer ** 2) + angles.addAngle(h1, o, h2, 1.824, 460.0 * unit.kilojoule_per_mole / unit.radian ** 2) + nonbonded.addException( + o, h1, + 0.0 * unit.elementary_charge * unit.elementary_charge, + 0.5 * (_O_SIGMA_NM + _H_SIGMA_NM) * unit.nanometer, + 0.0 * unit.kilojoule_per_mole, + ) + nonbonded.addException( + o, h2, + 0.0 * unit.elementary_charge * unit.elementary_charge, + 0.5 * (_O_SIGMA_NM + _H_SIGMA_NM) * unit.nanometer, + 0.0 * unit.kilojoule_per_mole, + ) + nonbonded.addException( + h1, h2, + 0.0 * unit.elementary_charge * unit.elementary_charge, + _H_SIGMA_NM * unit.nanometer, + 0.0 * unit.kilojoule_per_mole, + ) + return indices + + +def _build_topology_and_system(num_mm_waters=3, periodic=True, mm_oxygen_charge=_O_CHARGE): + topology = app.Topology() + chain = topology.addChain() + + # ML water (residue 0). + res = topology.addResidue("HOH", chain) + a_o = topology.addAtom("O", app.element.oxygen, res) + a_h1 = topology.addAtom("H1", app.element.hydrogen, res) + a_h2 = topology.addAtom("H2", app.element.hydrogen, res) + topology.addBond(a_o, a_h1) + topology.addBond(a_o, a_h2) + + for _ in range(num_mm_waters): + res = topology.addResidue("HOH", chain) + a_o = topology.addAtom("O", app.element.oxygen, res) + a_h1 = topology.addAtom("H1", app.element.hydrogen, res) + a_h2 = topology.addAtom("H2", app.element.hydrogen, res) + topology.addBond(a_o, a_h1) + topology.addBond(a_o, a_h2) + + if periodic: + topology.setPeriodicBoxVectors( + unit.Quantity(np.diag([_BOX_NM, _BOX_NM, _BOX_NM]), unit.nanometer) + ) + + system = openmm.System() + if periodic: + system.setDefaultPeriodicBoxVectors( + openmm.Vec3(_BOX_NM, 0, 0) * unit.nanometer, + openmm.Vec3(0, _BOX_NM, 0) * unit.nanometer, + openmm.Vec3(0, 0, _BOX_NM) * unit.nanometer, + ) + nonbonded = openmm.NonbondedForce() + if periodic: + nonbonded.setNonbondedMethod(openmm.NonbondedForce.PME) + nonbonded.setCutoffDistance(0.5 * unit.nanometer) + else: + nonbonded.setNonbondedMethod(openmm.NonbondedForce.NoCutoff) + bonds = openmm.HarmonicBondForce() + angles = openmm.HarmonicAngleForce() + + ml_o = (15.999, _O_CHARGE, _O_SIGMA_NM, _O_EPS_KJ) + ml_h = (1.008, _H_CHARGE, _H_SIGMA_NM, _H_EPS_KJ) + mm_o = (15.999, mm_oxygen_charge, _O_SIGMA_NM, _O_EPS_KJ) + mm_h_charge = -mm_oxygen_charge / 2.0 + mm_h = (1.008, mm_h_charge, _H_SIGMA_NM, _H_EPS_KJ) + + _add_water(system, nonbonded, bonds, angles, [ml_o, ml_h, ml_h]) + for _ in range(num_mm_waters): + _add_water(system, nonbonded, bonds, angles, [mm_o, mm_h, mm_h]) + + system.addForce(nonbonded) + system.addForce(bonds) + system.addForce(angles) + + return topology, system + + +def _initial_positions(num_mm_waters): + """Return positions in nanometers. ML water at ~box centre, MM waters scattered.""" + base_o = np.array([ + [0.55, 0.55, 0.55], # ML O + [0.20, 0.20, 0.20], + [0.95, 0.20, 0.30], + [0.30, 0.95, 0.95], + [0.95, 0.95, 0.95], + [0.10, 0.50, 0.10], + ])[: 1 + num_mm_waters] + H1_OFFSET = np.array([0.09572, 0.0, 0.0]) + H2_OFFSET = np.array([-0.0240, 0.0927, 0.0]) # 104.5 deg, 0.09572 nm bond + coords = [] + for o in base_o: + coords.append(o) + coords.append(o + H1_OFFSET) + coords.append(o + H2_OFFSET) + return np.array(coords) * unit.nanometer + + +# --------------------------------------------------------------------------- +# Integration tests +# --------------------------------------------------------------------------- + +def _energy_and_forces(context): + state = context.getState(getEnergy=True, getForces=True) + e = state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) + f = state.getForces(asNumpy=True).value_in_unit(unit.kilojoule_per_mole / unit.nanometer) + return float(e), np.asarray(f) + + +def test_polar_mace_pbc_runs_and_is_finite(polar_mace_model_path): + """Smoke test: createMixedSystem(embedding='electrostatic') with PolarMACE + under PBC returns a finite energy.""" + topology, mm_system = _build_topology_and_system(num_mm_waters=2) + potential = MLPotential("mace", modelPath=polar_mace_model_path) + mixed_system = potential.createMixedSystem( + topology, mm_system, [0, 1, 2], embedding="electrostatic" + ) + platform = openmm.Platform.getPlatformByName("Reference") + context = openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) + context.setPositions(_initial_positions(num_mm_waters=2)) + e, f = _energy_and_forces(context) + assert np.isfinite(e) + assert np.all(np.isfinite(f)) + + +def test_polar_mace_pbc_translation_invariance(polar_mace_model_path): + """Energy and forces must be invariant under whole-system translation by + one full box vector — the fundamental PBC sanity check.""" + topology, mm_system = _build_topology_and_system(num_mm_waters=2) + potential = MLPotential("mace", modelPath=polar_mace_model_path) + mixed_system = potential.createMixedSystem( + topology, mm_system, [0, 1, 2], embedding="electrostatic" + ) + platform = openmm.Platform.getPlatformByName("Reference") + context = openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) + + pos = _initial_positions(num_mm_waters=2) + context.setPositions(pos) + e0, f0 = _energy_and_forces(context) + + # Translate all atoms by +x box length. Forces are reported per-atom in the + # same order, so direct comparison is valid. + pos_shifted = ( + pos.value_in_unit(unit.nanometer) + np.array([_BOX_NM, 0.0, 0.0]) + ) * unit.nanometer + context.setPositions(pos_shifted) + e1, f1 = _energy_and_forces(context) + + assert e1 == pytest.approx(e0, rel=1e-6, abs=1e-4) + np.testing.assert_allclose(f1, f0, rtol=1e-5, atol=1e-3) + + +def test_polar_mace_pbc_ml_atom_near_boundary(polar_mace_model_path): + """If an ML atom sits near a periodic boundary, its energy must match the + energy from the unwrapped image (translation by one full box must give + the same result).""" + topology, mm_system = _build_topology_and_system(num_mm_waters=2) + potential = MLPotential("mace", modelPath=polar_mace_model_path) + mixed_system = potential.createMixedSystem( + topology, mm_system, [0, 1, 2], embedding="electrostatic" + ) + platform = openmm.Platform.getPlatformByName("Reference") + context = openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) + + # Put the ML water O near x=0 (boundary) with H atoms straddling the box. + pos_np = _initial_positions(num_mm_waters=2).value_in_unit(unit.nanometer) + pos_np[0] = np.array([0.05, 0.55, 0.55]) # O near +x boundary edge + pos_np[1] = pos_np[0] + np.array([0.09572, 0.0, 0.0]) # H still inside box + pos_np[2] = pos_np[0] + np.array([-0.024, 0.0927, 0.0]) # H near boundary + context.setPositions(pos_np * unit.nanometer) + e_near, f_near = _energy_and_forces(context) + assert np.isfinite(e_near) + assert np.all(np.isfinite(f_near)) + + # Translate everything by one full box: identical energy and forces. + context.setPositions((pos_np + np.array([_BOX_NM, 0.0, 0.0])) * unit.nanometer) + e_shift, f_shift = _energy_and_forces(context) + assert e_shift == pytest.approx(e_near, rel=1e-6, abs=1e-4) + np.testing.assert_allclose(f_shift, f_near, rtol=1e-5, atol=1e-3) + + +# --------------------------------------------------------------------------- +# Non-PBC tests: MACE electrostatic plumbing through openmmml +# --------------------------------------------------------------------------- + +def _build_nonpbc_mixed(potential, num_mm_waters, embedding, mm_oxygen_charge=_O_CHARGE): + topology, mm_system = _build_topology_and_system( + num_mm_waters=num_mm_waters, periodic=False, mm_oxygen_charge=mm_oxygen_charge + ) + mixed_system = potential.createMixedSystem( + topology, mm_system, [0, 1, 2], embedding=embedding + ) + platform = openmm.Platform.getPlatformByName("Reference") + return mixed_system, openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) + + +def _nonpbc_initial_positions(num_mm_waters): + base_o = np.array([ + [0.0, 0.0, 0.0], # ML O at origin + [0.50, 0.0, 0.0], # MM water on +x + [0.0, 0.50, 0.0], # MM water on +y + [0.0, 0.0, 0.50], # MM water on +z + ])[: 1 + num_mm_waters] + H1_OFFSET = np.array([0.09572, 0.0, 0.0]) + H2_OFFSET = np.array([-0.0240, 0.0927, 0.0]) + coords = [] + for o in base_o: + coords.append(o) + coords.append(o + H1_OFFSET) + coords.append(o + H2_OFFSET) + return np.array(coords) * unit.nanometer + + +def test_polar_mace_electrostatic_nonpbc_smoke(polar_mace_model_path): + """Non-PBC: createMixedSystem(embedding='electrostatic') with PolarMACE + runs end-to-end and produces finite energy/forces.""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + _, ctx = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="electrostatic") + ctx.setPositions(_nonpbc_initial_positions(num_mm_waters=2)) + e, f = _energy_and_forces(ctx) + assert np.isfinite(e) + assert np.all(np.isfinite(f)) + + +def test_polar_mace_electrostatic_changes_energy_vs_mechanical(polar_mace_model_path): + """Mechanical embedding does NOT pass MM charges into MACE; electrostatic + does. With nonzero MM charges, the two modes must give different ML energies + -- the only path that exercises the openmmml -> mm_charges plumbing.""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + positions = _nonpbc_initial_positions(num_mm_waters=2) + + _, ctx_mech = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="mechanical") + ctx_mech.setPositions(positions) + e_mech, _ = _energy_and_forces(ctx_mech) + + _, ctx_elec = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="electrostatic") + ctx_elec.setPositions(positions) + e_elec, _ = _energy_and_forces(ctx_elec) + + # Energies must not coincide -- with TIP3P-ish MM charges there is a real + # ML-MM Coulomb piece that mechanical embedding doesn't see. + assert abs(e_elec - e_mech) > 1e-3 + + +def test_polar_mace_electrostatic_zero_mm_charges_matches_mechanical(polar_mace_model_path): + """If MM charges are zero, the ML-MM Coulomb term vanishes and electrostatic + embedding should reproduce mechanical embedding (within float64 round-off).""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + positions = _nonpbc_initial_positions(num_mm_waters=2) + + _, ctx_mech = _build_nonpbc_mixed( + potential, num_mm_waters=2, embedding="mechanical", mm_oxygen_charge=0.0 + ) + ctx_mech.setPositions(positions) + e_mech, f_mech = _energy_and_forces(ctx_mech) + + _, ctx_elec = _build_nonpbc_mixed( + potential, num_mm_waters=2, embedding="electrostatic", mm_oxygen_charge=0.0 + ) + ctx_elec.setPositions(positions) + e_elec, f_elec = _energy_and_forces(ctx_elec) + + assert e_elec == pytest.approx(e_mech, rel=1e-5, abs=1e-3) + np.testing.assert_allclose(f_elec, f_mech, rtol=1e-4, atol=1e-2) + + +def test_polar_mace_electrostatic_mm_atoms_receive_back_reaction(polar_mace_model_path): + """The MACE mm_forces path must scatter forces onto MM atoms. With nonzero + MM charges, MM atoms must feel a non-trivial force; with zero MM charges + they must not (modulo MM-internal forces in the unmodified pieces of the + system, which we factor out by subtraction).""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + positions = _nonpbc_initial_positions(num_mm_waters=2) + + _, ctx_charged = _build_nonpbc_mixed( + potential, num_mm_waters=2, embedding="electrostatic", mm_oxygen_charge=_O_CHARGE + ) + ctx_charged.setPositions(positions) + _, f_charged = _energy_and_forces(ctx_charged) + + _, ctx_neutral = _build_nonpbc_mixed( + potential, num_mm_waters=2, embedding="electrostatic", mm_oxygen_charge=0.0 + ) + ctx_neutral.setPositions(positions) + _, f_neutral = _energy_and_forces(ctx_neutral) + + # MM atoms are indices 3..8 (2 waters); their force differs between charged + # and neutral runs, and that difference is exactly the ML-MM back reaction + # routed through openmmml. + delta_mm = f_charged[3:] - f_neutral[3:] + assert np.linalg.norm(delta_mm) > 1e-2 + + +def test_polar_mace_electrostatic_mm_charge_displacement_changes_ml_force( + polar_mace_model_path, +): + """Moving an MM charge changes the field on the ML region and therefore + the force on the ML atoms. This is the cleanest end-to-end check that the + field embedding is wired through openmmml into MACE.""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + + pos_a = _nonpbc_initial_positions(num_mm_waters=2) + # Move the second MM water (atoms 6, 7, 8) by +0.05 nm in x. + pos_b_np = pos_a.value_in_unit(unit.nanometer).copy() + pos_b_np[6:9, 0] += 0.05 + pos_b = pos_b_np * unit.nanometer + + _, ctx = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="electrostatic") + ctx.setPositions(pos_a) + _, f_a = _energy_and_forces(ctx) + ctx.setPositions(pos_b) + _, f_b = _energy_and_forces(ctx) + + # ML forces (atoms 0..2) must have changed: the ML region polarization + # responds to the shifted MM field. + delta_ml = f_b[:3] - f_a[:3] + assert np.linalg.norm(delta_ml) > 1e-3 + + +# --------------------------------------------------------------------------- +# Link-atom CHARGE redistribution (Z1 / DZ1) — end-to-end via MACE +# --------------------------------------------------------------------------- +# +# The unit tests in TestLinkChargeRedistribution.py cover the pure helper +# (apply_link_charge_redistribution). The tests below exercise the actual +# integration into MACEPotentialImpl.addForces and verify that: +# (a) the scheme parameter propagates all the way to MACE's mm_charges +# input (different schemes give different energies/forces); +# (b) forces remain conservative under each scheme (F = -dE/dx, FD-tested). + + +def _build_capped_nonpbc_chain(): + """Tiny 6-atom chain: ML region {0,1,2} bonded to MM atom 3 (the M atom) + which is bonded to MM atoms {4,5} (M's neighbours, for DZ1 to populate). + All MM atoms carry non-zero partial charges so the linkChargeScheme has + something to redistribute. + """ + system = openmm.System() + # Make all six atoms hydrogen (element 1) so MACE's atomic-number map + # exercised by the polar_mace_model_path fixture works (it knows H+O only). + for _ in range(6): + system.addParticle(1.0) + nb = openmm.NonbondedForce() + # ML atoms get zero charge (they'll be re-zeroed by createMixedSystem anyway, + # but make it explicit). MM atoms get something the scheme can act on. + nb.addParticle(0.0, 0.30, 0.20) # 0 ML + nb.addParticle(0.0, 0.30, 0.20) # 1 ML + nb.addParticle(0.0, 0.30, 0.20) # 2 ML (Q atom in the cut) + nb.addParticle(-0.40, 0.30, 0.20) # 3 MM (the M atom) + nb.addParticle(0.15, 0.30, 0.20) # 4 MM (M1 neighbour) + nb.addParticle(0.25, 0.30, 0.20) # 5 MM (M1 neighbour) + nb.setNonbondedMethod(openmm.NonbondedForce.NoCutoff) + system.addForce(nb) + + topology = app.Topology() + chain = topology.addChain() + res = topology.addResidue("X", chain) + atoms = [topology.addAtom(f"H{i}", elem.hydrogen, res) for i in range(6)] + # Bonds: 0-1, 1-2, 2-3 (Q-M cut), 3-4 (M-M1), 3-5 (M-M1) + for a, b in [(0, 1), (1, 2), (2, 3), (3, 4), (3, 5)]: + topology.addBond(atoms[a], atoms[b]) + return topology, system + + +def _capped_chain_positions(): + return np.array([ + [0.00, 0.0, 0.0], + [0.11, 0.0, 0.0], + [0.22, 0.0, 0.0], + [0.33, 0.0, 0.0], # M atom + [0.40, 0.07, 0.0], # M1 + [0.40, -0.07, 0.0], # M1 + ]) * unit.nanometer + + +def _build_capped_mixed(potential, scheme): + topology, mm_system = _build_capped_nonpbc_chain() + mixed = potential.createMixedSystem( + topology, mm_system, atoms=[0, 1, 2], + embedding="electrostatic", + linkRecords=[(2, 3, 1.09)], # Q=2 (in ML), M=3 (MM) + linkChargeScheme=scheme, + ) + plat = openmm.Platform.getPlatformByName("Reference") + ctx = openmm.Context(mixed, openmm.VerletIntegrator(0.001), plat) + ctx.setPositions(_capped_chain_positions()) + return mixed, ctx + + +# elem import for the topology builder; sits with the openmm.app namespace. +from openmm.app import element as elem # noqa: E402 (kept near point of use) + + +def _mmChargesReachingModel(potential, scheme): + """The mm_charges array that createMixedSystem bakes into the ML force. + + Asserting on the array rather than on the energy it produces is deliberate. + The energy route would only distinguish the schemes if the model in use + responded to the MM charge *values*, and the synthetic checkpoint built by + the polar_mace_model_path fixture does not for this fixture's geometry: its + energy is bit-identical whether mm_charges are passed unchanged, zeroed, or + scaled tenfold. What the redistribution schemes are responsible for is the + contents of that array, so that is what these tests pin, exactly rather + than as an inequality. + """ + captured = {} + original = macepotential._computeMACE + + def capture(state, **args): + if args.get("mmInfo") is not None and "mm_charges" not in captured: + captured["mm_atoms"] = np.asarray(args["mmInfo"]["mm_atoms"]) + captured["mm_charges"] = np.asarray(args["mmInfo"]["mm_charges"]) + return original(state, **args) + + macepotential._computeMACE = capture + try: + _, context = _build_capped_mixed(potential, scheme=scheme) + context.getState(getEnergy=True) + finally: + macepotential._computeMACE = original + assert "mm_charges" in captured, "the ML force never received MM charges" + return captured["mm_atoms"], captured["mm_charges"] + + +def test_link_charge_scheme_none_passes_forcefield_charges(polar_mace_model_path): + """With scheme='none' the model sees the MM force field charges verbatim.""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + mmAtoms, mmCharges = _mmChargesReachingModel(potential, "none") + # Atoms 3 (the M atom), 4 and 5 (its MM neighbours) from + # _build_capped_nonpbc_chain. + np.testing.assert_array_equal(mmAtoms, [3, 4, 5]) + np.testing.assert_allclose(mmCharges, [-0.40, 0.15, 0.25], atol=1e-12) + + +def test_link_charge_scheme_z1_zeros_the_m_atom(polar_mace_model_path): + """Z1 zeros the M-atom charge and leaves every other MM charge alone, + which changes the total MM charge by -q_M.""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + _, mmCharges = _mmChargesReachingModel(potential, "z1") + np.testing.assert_allclose(mmCharges, [0.0, 0.15, 0.25], atol=1e-12) + assert mmCharges.sum() == pytest.approx(0.40, abs=1e-12) + + +def test_link_charge_scheme_dz1_spreads_the_m_charge_and_conserves(polar_mace_model_path): + """DZ1 zeros the M-atom charge and spreads it over M's two MM neighbours, + -0.40/2 = -0.20 each, leaving the total MM charge unchanged.""" + potential = MLPotential("mace", modelPath=polar_mace_model_path) + _, mmCharges = _mmChargesReachingModel(potential, "dz1") + np.testing.assert_allclose(mmCharges, [0.0, 0.15 - 0.20, 0.25 - 0.20], atol=1e-12) + assert mmCharges.sum() == pytest.approx(0.0, abs=1e-12) + + +@pytest.mark.parametrize("scheme", ["none", "z1", "dz1"]) +def test_link_charge_scheme_fd_force_consistency(polar_mace_model_path, scheme): + """For each scheme, the MACE-side forces must equal the negative gradient + of the MACE-side energy. Central finite difference in f64. + + Atoms 2 (Q) and 3 (M) are included deliberately: the cap sits between them + and is repositioned from their coordinates every step, so they are the only + atoms whose forces exercise the cap chain rule in redistribute_cap_force. + Perturbing only the interior QM atoms would leave that untested. + + The criterion is relative to the size of the force being checked, not an + absolute number of kJ/mol/nm. This runs against a synthetic *untrained* + checkpoint at an arbitrary geometry, where forces reach ~1e7 kJ/mol/nm on + the boundary atoms and only ~1e6 in the interior; a fixed absolute + threshold is meaninglessly loose for the former and tight for the latter. + The central-difference truncation error scales as delta^2 while a genuine + force/energy inconsistency does not shrink with delta at all, which is what + separates the two and what this bound is sized against. + """ + potential = MLPotential("mace", modelPath=polar_mace_model_path) + _, ctx = _build_capped_mixed(potential, scheme=scheme) + + pos0 = _capped_chain_positions().value_in_unit(unit.nanometer).copy() + delta = 1e-5 # nm; the truncation term goes as delta^2 + + ctx.setPositions(pos0 * unit.nanometer) + e_ref = ctx.getState(getEnergy=True).getPotentialEnergy().value_in_unit( + unit.kilojoules_per_mole) + f_ref = ctx.getState(getForces=True).getForces(asNumpy=True).value_in_unit( + unit.kilojoules_per_mole / unit.nanometer) + + max_abs_err = 0.0 + # Atoms 2 (Q) and 3 (M) are the ones the cap is placed between, so they are + # the only ones whose forces exercise the cap chain rule. Perturbing just + # the interior QM atoms 0 and 1 would leave the redistribution untested. + for ia in (0, 1, 2, 3): + for ax in range(3): + saved = pos0[ia, ax] + pos0[ia, ax] = saved + delta + ctx.setPositions(pos0 * unit.nanometer) + e_p = ctx.getState(getEnergy=True).getPotentialEnergy().value_in_unit( + unit.kilojoules_per_mole) + pos0[ia, ax] = saved - delta + ctx.setPositions(pos0 * unit.nanometer) + e_m = ctx.getState(getEnergy=True).getPotentialEnergy().value_in_unit( + unit.kilojoules_per_mole) + pos0[ia, ax] = saved + f_num = -(e_p - e_m) / (2.0 * delta) + err = abs(f_num - f_ref[ia, ax]) + max_abs_err = max(max_abs_err, err) + forceScale = float(np.abs(f_ref).max()) + max_rel_err = max_abs_err / forceScale + assert max_rel_err < 5e-5, ( + f"FD vs analytical force mismatch for scheme='{scheme}': " + f"max |dF| = {max_abs_err:.3e} kJ/mol/nm against a force scale of " + f"{forceScale:.3e}, i.e. relative {max_rel_err:.3e} (expected < 5e-5)" + ) From 95c57fec6bb614d8df8821c984c3e5f9c1c6fd49 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 14:09:29 +0000 Subject: [PATCH 03/20] Use graph-longrange external-source wrappers --- doc/userguide.md | 9 +- openmmml/models/_polarmace_external.py | 156 +++++-------------------- 2 files changed, 35 insertions(+), 130 deletions(-) diff --git a/doc/userguide.md b/doc/userguide.md index 29efd6f..f8de749 100644 --- a/doc/userguide.md +++ b/doc/userguide.md @@ -403,11 +403,10 @@ with the model name `mace` and a `modelPath`. For a custom checkpoint the metho known until it is loaded; a non-PolarMACE checkpoint is rejected at that point rather than silently falling back to mechanical embedding, since the mixed system has already had its ML-MM electrostatics removed. -The damped real-space and reciprocal-space ML-MM cross energy, including slab and molecule-in-box corrections, is -implemented by `GTOElectrostaticCrossEnergy` in `graph_longrange`. OpenMM-ML only adapts the external positions and -charges to the stock PolarMACE electrostatic feature and energy blocks, then returns the corresponding MM forces. The -normalization is therefore inherited from the checkpoint's own GTO density and feature bases rather than duplicated in -the OpenMM interface. +The damped real-space and reciprocal-space ML-MM cross energy, including slab and molecule-in-box corrections, and the +external-source feature and energy wrappers are implemented in `graph_longrange`. OpenMM-ML only supplies the external +positions and charges to those blocks, then returns the corresponding MM forces. The normalization is therefore +inherited from the checkpoint's own GTO density and feature bases rather than duplicated in the OpenMM interface. Interpolation is not supported, because at `lambda_interpolate=0` the ML-MM electrostatics would be missing from the MM endpoint. diff --git a/openmmml/models/_polarmace_external.py b/openmmml/models/_polarmace_external.py index f690743..ab518d8 100644 --- a/openmmml/models/_polarmace_external.py +++ b/openmmml/models/_polarmace_external.py @@ -65,125 +65,6 @@ def _rebuild_energy_block(block, pbc_handling: str = "auto"): return rebuilt -class _ExternalFeatureBlock: - """Mixin-like implementation installed around a graph feature module.""" - - def __init__(self, base): - import torch - - class FeatureBlock(torch.nn.Module): - def __init__(inner_self, wrapped): - super().__init__() - inner_self.base = wrapped - inner_self._external = None - - def __getattr__(inner_self, name): - try: - return super().__getattr__(name) - except AttributeError: - return getattr(inner_self.base, name) - - def set_external_sources(inner_self, external): - inner_self._external = external - - def precompute_geometry(inner_self, **kwargs): - base_kwargs = dict(kwargs) - base_kwargs.pop("force_pbc_evaluator", None) - cache = inner_self.base.precompute_geometry(**base_kwargs) - external = inner_self._external - if external is None: - return cache - external_cache = inner_self.base.precompute_geometry_source_target( - k_vectors=base_kwargs["k_vectors"], - k_norm2=base_kwargs["k_norm2"], - k_vector_batch=base_kwargs["k_vector_batch"], - k0_mask=base_kwargs["k0_mask"], - src_positions=external["positions"], - src_batch=external["batch"], - tgt_positions=base_kwargs["node_positions"], - tgt_batch=base_kwargs["batch"], - volume=base_kwargs["volume"], - pbc=base_kwargs["pbc"], - ) - external_field = inner_self.base.forward_dynamic_source_target( - cache=external_cache, - source_feats=external["features"], - ) - result = dict(cache) - result["_openmmml_external_field"] = external_field - return result - - def forward_dynamic(inner_self, cache, source_feats, pbc=None): - if source_feats.dim() == 3 and source_feats.shape[-2] == 1: - source_feats = source_feats.squeeze(-2) - value = inner_self.base.forward_dynamic( - cache=cache, source_feats=source_feats - ) - external_field = cache.get("_openmmml_external_field") - if external_field is not None: - # PolarMACE has two spin channels. Each channel receives half - # of the physical external potential. - value = value + 0.5 * external_field - return value - - self.module = FeatureBlock(base) - - -class _ExternalEnergyBlock: - def __init__(self, base): - import torch - try: - from graph_longrange.external_source_energy import ( - GTOElectrostaticCrossEnergy, - ) - except ImportError as exc: - raise ImportError( - "PolarMACE electrostatic embedding requires a graph_longrange " - "release that provides GTOElectrostaticCrossEnergy." - ) from exc - - class EnergyBlock(torch.nn.Module): - def __init__(inner_self, wrapped): - super().__init__() - inner_self.base = wrapped - inner_self.cross = GTOElectrostaticCrossEnergy.from_energy(wrapped) - inner_self._external = None - - def __getattr__(inner_self, name): - try: - return super().__getattr__(name) - except AttributeError: - return getattr(inner_self.base, name) - - def set_external_sources(inner_self, external): - inner_self._external = external - - def forward(inner_self, **kwargs): - base_kwargs = dict(kwargs) - base_kwargs.pop("force_pbc_evaluator", None) - energy = inner_self.base(**base_kwargs) - external = inner_self._external - if external is None: - return energy - cross = inner_self.cross( - k_vectors=base_kwargs["k_vectors"], - k_norm2=base_kwargs["k_norm2"], - k_vector_batch=base_kwargs["k_vector_batch"], - k0_mask=base_kwargs["k0_mask"], - source_feats=base_kwargs["source_feats"], - source_positions=base_kwargs["node_positions"], - source_batch=base_kwargs["batch"], - target_feats=external["features"], - target_positions=external["positions"], - target_batch=external["batch"], - volume=base_kwargs["volume"], - pbc=base_kwargs["pbc"], - ) - return energy + cross - - self.module = EnergyBlock(base) - - def _prepare_external_sources(model, data, compute_force: bool): import torch @@ -247,10 +128,28 @@ def enable_polarmace_external_sources(model): f"{model.__class__.__name__}." ) + try: + from graph_longrange.external_source_energy import ( + GTOElectrostaticExternalSourceEnergy, + ) + from graph_longrange.external_source_features import ( + GTOElectrostaticExternalSourceFeatures, + ) + except ImportError as exc: + raise ImportError( + "PolarMACE electrostatic embedding requires a graph_longrange " + "release that provides the external-source energy and feature blocks." + ) from exc + feature_base = _rebuild_feature_block(model.electric_potential_descriptor) energy_base = _rebuild_energy_block(model.coulomb_energy) - feature_block = _ExternalFeatureBlock(feature_base).module - energy_block = _ExternalEnergyBlock(energy_base).module + feature_block = GTOElectrostaticExternalSourceFeatures.from_features( + feature_base, + # PolarMACE has two spin channels. Each receives half of the physical + # external potential. + external_scale=0.5, + ) + energy_block = GTOElectrostaticExternalSourceEnergy.from_energy(energy_base) model.electric_potential_descriptor = feature_block model.coulomb_energy = energy_block @@ -303,8 +202,15 @@ def forward( "Cartesian forces only." ) - self.model.electric_potential_descriptor.set_external_sources(external) - self.model.coulomb_energy.set_external_sources(external) + external_kwargs = { + "external_feats": external["features"], + "external_positions": external["positions"], + "external_batch": external["batch"], + } + self.model.electric_potential_descriptor.set_external_sources( + **external_kwargs + ) + self.model.coulomb_energy.set_external_sources(**external_kwargs) try: result = self.model( data, @@ -339,8 +245,8 @@ def forward( result["mm_forces"] = None return result finally: - self.model.electric_potential_descriptor.set_external_sources(None) - self.model.coulomb_energy.set_external_sources(None) + self.model.electric_potential_descriptor.clear_external_sources() + self.model.coulomb_energy.clear_external_sources() return PolarMACEExternalSources(model) From 7b1004f0c8a5d69d09230ea3fbb13c8c74785870 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 14:26:28 +0000 Subject: [PATCH 04/20] Remove link-atom scope from electrostatic embedding --- doc/userguide.md | 23 - openmmml/models/_links.py | 331 -------------- openmmml/models/_polarmace_external.py | 254 ----------- openmmml/models/macepotential.py | 458 ++++++++++++-------- test/TestLinkChargeRedistribution.py | 222 ---------- test/TestLinkRecordsTargetDistUnits.py | 147 ------- test/TestMACEPotential.py | 2 - test/TestPolarMACEElectrostaticEmbedding.py | 192 -------- 8 files changed, 270 insertions(+), 1359 deletions(-) delete mode 100644 openmmml/models/_links.py delete mode 100644 openmmml/models/_polarmace_external.py delete mode 100644 test/TestLinkChargeRedistribution.py delete mode 100644 test/TestLinkRecordsTargetDistUnits.py diff --git a/doc/userguide.md b/doc/userguide.md index f8de749..97b0727 100644 --- a/doc/userguide.md +++ b/doc/userguide.md @@ -105,8 +105,6 @@ When using MACE models, the following extra keyword arguments to `createSystem() | --- | --- | | `precision` | The numerical precision of the model. Supported options are `'single'` and `'double'`. If `None`, the default precision of the model is used. | | `returnEnergyType` | Whether to return the interaction energy or the energy including the self-energy. The default is `'energy'`, which is the scalar the reported forces are the gradient of, so the potential is conservative. Supported options are `'interaction_energy'` and `'energy'`. | -| `linkRecords` | Hydrogen link atoms to cap covalent bonds crossing the ML/MM boundary; see *Link atoms* below. Only meaningful for `createMixedSystem()`. | -| `linkChargeScheme` | How to treat the MM charges on the MM-side boundary atoms when `linkRecords` is given; see *Link atoms* below. The default is `'dz1'`. | | `embedding` | Passed by the embedding method rather than by you; see *Embeddings* below. | | `device` | The PyTorch device to perform calculations on, either a `torch.device` object or a string (such as `'cuda'` or `'cpu'`.) If omitted, a device is chosen automatically. | | `charge` | The total charge of the system. If omitted, it is assumed to be 0. This is only used by MACE-OMOL-0. For other models it is ignored. | @@ -414,27 +412,6 @@ endpoint. system = potential.createMixedSystem(topology, mm_system, ml_atoms, embedding='electrostatic') ``` -#### Link atoms - -When the ML subset is carved out of the middle of a molecule, the covalent bonds crossing the boundary leave the ML -atoms with dangling valences. The MACE interface can cap each such bond with a hydrogen link atom, placed along the -bond vector at every force evaluation and with its forces redistributed onto the two real atoms it sits between, so the -capped system stays conservative. Caps are requested with the `linkRecords` argument to `createMixedSystem()`, which -takes a sequence of `(q, m, target_dist)` tuples — `q` is the ML-side atom index, `m` the MM-side atom index, and -`target_dist` the desired cap distance from `q` in Ångströms — or the path to a CSV file with `q_idx1`, `m_idx1`, and -`target_dist_ang` columns. Note that the tuples are 0-based, matching OpenMM's atom indices, while the CSV columns are -1-based, hence their names. - -The MM partial charge on each MM-side boundary atom would otherwise sit about 1.5 Å from the nearest ML atom, through -the cap, and over-polarise it. The `linkChargeScheme` argument controls how that charge is handled, affecting only the -charges passed to the MLIP; the MM `NonbondedForce` is left alone, so MM-MM electrostatics remain exact. - -| `linkChargeScheme` | Description | -| --- | --- | -| `'dz1'` | The default. Sets the MM-side boundary atom charge to zero and spreads it evenly over that atom's MM neighbours, conserving total charge. Falls back to `'z1'`, with a warning, for a boundary atom with no MM neighbours. | -| `'z1'` | Sets the MM-side boundary atom charge to zero without redistributing it, which changes the total charge. | -| `'none'` | Leaves the MM charges untouched. | - ### Molecules Spanning the ML-MM Region OpenMM-ML's mechanical embedding implementation supports the link-atom method for molecules having bonds crossing the diff --git a/openmmml/models/_links.py b/openmmml/models/_links.py deleted file mode 100644 index 25a93d1..0000000 --- a/openmmml/models/_links.py +++ /dev/null @@ -1,331 +0,0 @@ -"""Hydrogen link-atom (cap) primitives shared across the codebase. - -Caps are virtual hydrogen atoms placed on the QM/ML side of a Q-M -boundary bond. Their positions are deterministic functions of the -Q (ML-side) and M (MM-side) atom positions; their forces from the -ML evaluator must be redistributed onto Q and M before being handed -back to OpenMM. - -The caller is `openmmml/models/macepotential.py::_computeMACE`, which -places caps in MACE's input each step and redistributes forces from -MACE's output. The math lives here rather than there so that any -further evaluator needing caps uses the *same* placement formula and -the *same* redistribution Jacobian: - - r_cap = (1 - C_L) * r_Q + C_L * r_M where C_L = target_dist / |r_M - r_Q| - - F_Q ← F_Q + (1 - C_L) F_cap + C_L (F_cap · ê_b) ê_b - F_M ← F_M + C_L F_cap - C_L (F_cap · ê_b) ê_b - where ê_b = (r_M - r_Q) / |r_M - r_Q| - -The redistribution preserves total force and total torque (the cap is -constrained to lie along the Q-M axis at fixed fraction `C_L`, so -its motion is entirely determined by Q and M motion; the chain rule -of energy w.r.t. Q,M gives the formulas above). -""" -from __future__ import annotations - -import numpy as np - - -# All 27 neighbouring lattice-cell offsets (coefficients in {-1, 0, 1}^3). -_OFFSETS_27 = np.array( - [[i, j, k] for i in (-1, 0, 1) for j in (-1, 0, 1) for k in (-1, 0, 1)], - dtype=np.float64, -) - - -def minimum_image_M( - r_Q: np.ndarray, - r_M: np.ndarray, - cell: np.ndarray, - max_bond: float = 2.5, -) -> np.ndarray: - """Return each M atom shifted to its minimum image relative to its Q partner. - - When a Q-M boundary bond straddles a periodic boundary the raw M position is - in a different image than Q, and both cap *placement* and cap-force - *redistribution* must use the same imaged M (otherwise the redistributed - force is not the gradient of the energy the cap produced). Both call sites in - ``macepotential._computeMACE`` therefore route ``r_M`` through this helper. - - The nearest image is found by a 27-cell search: round the fractional - displacement to the nearest lattice cell, then pick the shortest candidate - over that cell and its 26 neighbours. This is exact for a real (weakly skewed) - simulation box; it is *not* a certified closest-vector solver for - pathologically skewed triclinic cells, which do not occur for physical MD - systems. As a sanity check, the imaged Q-M distance must stay below - ``max_bond`` (Angstrom); a longer bond means a genuinely wrapped / broken pair - (or a box thinner than the bond) and raises rather than placing a bad cap. - - Parameters - ---------- - r_Q, r_M : arrays of shape (K, 3) - Q (ML-side) and M (MM-side) positions, one row per cap, in Angstrom - (the unit ``cell`` is given in). - cell : array of shape (3, 3) - Periodic box vectors as rows (OpenMM convention), in Angstrom. - max_bond : float, optional - Chemical sanity ceiling on the imaged Q-M bond length (Angstrom). Frontier - bonds are ~1.0-1.6 A; the default 2.5 A leaves margin while still catching - a wrongly-imaged / wrapped pair. - - Returns - ------- - r_M_imaged : array of shape (K, 3) - M positions shifted to the image nearest their Q partner. - """ - r_Q = np.atleast_2d(np.asarray(r_Q, dtype=np.float64)) - r_M = np.atleast_2d(np.asarray(r_M, dtype=np.float64)) - cell = np.asarray(cell, dtype=np.float64) - - if cell.shape != (3, 3): - raise ValueError("cell must have shape (3, 3)") - if abs(np.linalg.det(cell)) < 1e-12: - raise ValueError("cell is singular or nearly singular") - if r_Q.shape != r_M.shape: - raise ValueError("r_Q and r_M must have the same shape") - - dr = r_M - r_Q - base = np.round(dr @ np.linalg.inv(cell)) # nearest lattice cell (orthorhombic guess) - shifts = base[:, None, :] + _OFFSETS_27[None, :, :] # (K, 27, 3) integer coefficients - cand = dr[:, None, :] - shifts @ cell # (K, 27, 3) candidate displacements - best = np.argmin(np.einsum("kij,kij->ki", cand, cand), axis=1) - dr_mi = cand[np.arange(cand.shape[0]), best] - - bond = np.linalg.norm(dr_mi, axis=-1) - if np.any(bond >= max_bond): - raise ValueError( - f"imaged Q-M link bond length {bond.max():.2f} A exceeds max_bond " - f"{max_bond:.2f} A; the pair is genuinely wrapped/broken (or the box is " - "thinner than the bond). Increase max_bond only if this bond is real." - ) - return r_Q + dr_mi - - -def compute_cap_positions( - r_Q: np.ndarray, - r_M: np.ndarray, - target_dist: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - """Compute cap positions and the per-cap C_L ratio. - - Parameters - ---------- - r_Q : array of shape (K, 3) - Position of each Q atom (ML side of the boundary bond), one - row per cap, in the same length unit as `target_dist`. - r_M : array of shape (K, 3) - Position of each M atom (MM side), aligned 1:1 with `r_Q`. - target_dist : array of shape (K,) - Target distance from Q to the cap. Typically a hydrogen bond - length (e.g., 1.09 Å for C-H). - - Returns - ------- - r_cap : array of shape (K, 3) - Cap positions. - C_L : array of shape (K,) - The placement ratio used for each cap. The caller usually - also needs `C_L` for the corresponding force redistribution. - - Raises - ------ - ValueError - If any C_L falls outside (0, 1) — meaning the requested - target distance exceeds (or equals) the actual Q-M distance, - or Q and M are coincident. Both cases are user-error. - """ - r_Q = np.asarray(r_Q, dtype=np.float64) - r_M = np.asarray(r_M, dtype=np.float64) - target_dist = np.asarray(target_dist, dtype=np.float64) - v = r_M - r_Q - s = np.linalg.norm(v, axis=-1) - if np.any(s == 0.0): - raise ValueError("Cap Q and M atoms are coincident.") - C_L = target_dist / s - if np.any(~np.isfinite(C_L)) or np.any(C_L <= 0.0) or np.any(C_L >= 1.0): - raise ValueError( - f"Link-atom C_L out of (0, 1): {C_L}. " - "Check that target_dist < |r_M - r_Q| for every cap." - ) - r_cap = (1.0 - C_L)[:, None] * r_Q + C_L[:, None] * r_M - return r_cap, C_L - - -def redistribute_cap_force( - F_cap: np.ndarray, - r_Q: np.ndarray, - r_M: np.ndarray, - C_L: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - """Map per-cap forces onto the corresponding Q and M atoms. - - Parameters - ---------- - F_cap : array of shape (K, 3) - Force on each cap from the ML evaluator. - r_Q, r_M : arrays of shape (K, 3) - Positions of the Q (ML side) and M (MM side) atoms. - C_L : array of shape (K,) - Placement ratio for each cap (same as returned by - `compute_cap_positions`). - - Returns - ------- - F_Q_add : array of shape (K, 3) - Force contribution to add onto each Q atom. - F_M_add : array of shape (K, 3) - Force contribution to add onto each M atom. - - Notes - ----- - The caller is responsible for accumulating the returned - contributions onto the right global atom indices (e.g., via - `f[q_global] += F_Q_add` if multiple caps share a Q or M atom - is forbidden by `_prepareLinkRecords`). - """ - F_cap = np.asarray(F_cap, dtype=np.float64) - r_Q = np.asarray(r_Q, dtype=np.float64) - r_M = np.asarray(r_M, dtype=np.float64) - C_L = np.asarray(C_L, dtype=np.float64) - v = r_M - r_Q - s = np.linalg.norm(v, axis=-1) - e_b = v / s[:, None] - proj = np.einsum("ki,ki->k", F_cap, e_b) - F_Q_add = (1.0 - C_L)[:, None] * F_cap + (C_L * proj)[:, None] * e_b - F_M_add = C_L[:, None] * F_cap - (C_L * proj)[:, None] * e_b - return F_Q_add, F_M_add - - -# ----------------------------------------------------------------------------- -# Link-atom CHARGE redistribution (Z1 / DZ1) -# -# Standard QM/MM correction. When a Q–M bond is cut and the M atom is replaced -# (on the QM side) by an H link atom, the M atom's full partial charge sits -# ~1.5 Å from the nearest QM atom. With electrostatic embedding the QM region -# sees that close-in partial charge through the link H and gets -# over-polarised. The "charge shift" / "redistributed charge" schemes move -# the M partial charge away from the boundary. -# -# Z1 — set q_M ← 0. Cheapest; breaks total MM-charge neutrality by -q_M_orig. -# DZ1 — Z1 + redistribute q_M_orig / N(M1) onto each MM neighbour (M1 atom). -# Preserves total MM charge to round-off. -# -# Both schemes only modify the *constant* mm_charges array fed into the ML -# potential's electrostatic embedding — they do not touch the OpenMM MM -# scaffold's NonbondedForce, which keeps the MM-MM Coulomb sum bit-exact -# with the original force field (matching standard QM/MM practice). -# ----------------------------------------------------------------------------- - -def apply_link_charge_redistribution( - mm_atoms: np.ndarray, - mm_charges: np.ndarray, - link_info: dict, - topology, - scheme: str, -) -> np.ndarray: - """Apply Z1 / DZ1 link-atom charge redistribution to an MM charges array. - - Parameters - ---------- - mm_atoms : np.ndarray, shape (n_mm,) - Global atom indices of the MM atoms, in the order their charges appear - in ``mm_charges``. - mm_charges : np.ndarray, shape (n_mm,) - Original MM partial charges aligned with ``mm_atoms``. - link_info : dict - Output of ``MACEPotentialImpl._prepareLinkRecords``; must contain - ``q_global`` and ``m_global`` arrays. - topology : openmm.app.Topology - The full system topology — needed to look up M's MM neighbours (M1) for - DZ1. - scheme : str - One of ``"none"``, ``"z1"``, ``"dz1"``. - - Returns - ------- - np.ndarray, shape (n_mm,) - New MM charges array (always a copy of the input; never an alias). - - Notes - ----- - For ``"dz1"``, if any M atom has zero MM neighbours the call falls back to - Z1 for that atom (charge zeroed, nothing to redistribute) and emits a - ``UserWarning``. Total MM charge will then drift by the orphan q_M_orig. - """ - scheme = scheme.lower() - if scheme not in {"none", "z1", "dz1"}: - raise ValueError( - f"Unsupported linkChargeScheme {scheme!r}. " - "Supported in this iteration: 'none', 'z1', 'dz1'." - ) - new_charges = np.asarray(mm_charges, dtype=np.float64).copy() - if scheme == "none": - return new_charges - - mm_idx_to_row = {int(g): k for k, g in enumerate(mm_atoms)} - qm_set = set(int(g) for g in link_info["q_global"]) - m_globals = [int(g) for g in link_info["m_global"]] - - # Snapshot original q_M before any zeroing so DZ1 redistributes the - # *original* value even if two M atoms happened to be the same row (they - # are validated unique upstream, but the snapshot is also free insurance). - orig_qM = {m: float(new_charges[mm_idx_to_row[m]]) for m in m_globals} - - if scheme == "z1": - for m in m_globals: - new_charges[mm_idx_to_row[m]] = 0.0 - return new_charges - - # ---- scheme == "dz1" ---- - # Build the MM-neighbour list for each M atom by walking topology bonds. - # Filter: the neighbour must (a) be in mm_atoms, (b) not be a Q atom. - # The QM filter is redundant against (a) when q_global ⊂ atoms ⊂ mm complement, - # but cheap and explicit. - # A neighbour that is itself an M atom is excluded: it is about to be zeroed - # too, so charge handed to it would simply be discarded. Without that filter - # the result depends on the order the link records happen to be listed in, - # and two M atoms bonded to each other leave one of them holding charge that - # DZ1 exists to remove. - m_set = set(m_globals) - mm_neighbors_of_M: dict[int, list[int]] = {m: [] for m in m_globals} - for bond in topology.bonds(): - a = bond.atom1.index - b = bond.atom2.index - if a in m_set and b in mm_idx_to_row and b not in qm_set and b not in m_set: - mm_neighbors_of_M[a].append(b) - if b in m_set and a in mm_idx_to_row and a not in qm_set and a not in m_set: - mm_neighbors_of_M[b].append(a) - - # Zero every M atom before redistributing any charge, so a share can never - # be written to a row that a later iteration then clears. - import warnings - for m in m_globals: - new_charges[mm_idx_to_row[m]] = 0.0 - - for m in m_globals: - m1s = mm_neighbors_of_M[m] - if not m1s: - warnings.warn( - f"DZ1: M atom {m} has no MM neighbours to redistribute " - f"q={orig_qM[m]:+.4f} e onto. Falling back to Z1 for this atom.", - stacklevel=2, - ) - continue - share = orig_qM[m] / len(m1s) - for m1 in m1s: - new_charges[mm_idx_to_row[m1]] += share - - # Conservation check: warn if numerically non-trivial drift (DZ1 should - # preserve total to f64 round-off). - delta = float(new_charges.sum() - - np.asarray(mm_charges, dtype=np.float64).sum()) - if abs(delta) > 1e-9: - import warnings as _w - _w.warn( - f"DZ1: post-redistribution total MM charge drift = {delta:+.3e} e " - "(likely due to M atoms with zero MM neighbours).", - stacklevel=2, - ) - return new_charges diff --git a/openmmml/models/_polarmace_external.py b/openmmml/models/_polarmace_external.py deleted file mode 100644 index ab518d8..0000000 --- a/openmmml/models/_polarmace_external.py +++ /dev/null @@ -1,254 +0,0 @@ -"""Small eager-mode adapter for PolarMACE external electrostatic sources. - -The electrostatic mathematics lives in graph_longrange. This module only -bridges dynamic MM sources into an otherwise unchanged PolarMACE forward pass. -It is intended for OpenMM's PythonForce execution path, not TorchScript export. -""" - -from __future__ import annotations - - -def _floating_reference(module): - for tensor in module.buffers(): - if tensor.is_floating_point(): - return tensor - for tensor in module.parameters(): - if tensor.is_floating_point(): - return tensor - return None - - -def _rebuild_feature_block(block, pbc_handling: str = "auto"): - """Rebuild a deterministic graph block saved by an older graph release.""" - from graph_longrange.features import GTOElectrostaticFeatures - - realspace = block.realspace_features - quadrupoles = bool( - getattr( - getattr(block.non_periodic_correction_terms, "self_field", None), - "include_quadrupole_corrections", - False, - ) - ) - rebuilt = GTOElectrostaticFeatures( - density_max_l=int(realspace.density_max_l), - density_smearing_width=float(realspace.density_smearing_width), - feature_max_l=int(realspace.projection_max_l), - feature_smearing_widths=[ - float(x) for x in realspace.projection_smearing_widths - ], - include_self_interaction=bool(block.include_self_interaction), - kspace_cutoff=float(block.kspace_cutoff), - quadrupole_feature_corrections=quadrupoles, - integral_normalization=str(block.feature_basis.normalize), - pbc_handling=pbc_handling, - ) - reference = _floating_reference(block) - if reference is not None: - rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) - return rebuilt - - -def _rebuild_energy_block(block, pbc_handling: str = "auto"): - from graph_longrange.energy import GTOElectrostaticEnergy - - rebuilt = GTOElectrostaticEnergy( - density_max_l=int(block.density_max_l), - density_smearing_width=float(block.density_smearing_width), - kspace_cutoff=float(block.kspace_cutoff), - include_self_interaction=bool(block.include_self_interaction), - pbc_handling=pbc_handling, - ) - reference = _floating_reference(block) - if reference is not None: - rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) - return rebuilt - - -def _prepare_external_sources(model, data, compute_force: bool): - import torch - - positions = data.get("mm_positions") - charges = data.get("mm_charges") - multipoles = data.get("mm_multipoles") - if charges is not None and multipoles is not None: - raise ValueError("mm_charges and mm_multipoles are mutually exclusive.") - values = multipoles if multipoles is not None else charges - if positions is None or values is None or positions.numel() == 0 or values.numel() == 0: - return None - - ml_positions = data["positions"] - positions = positions.to(device=ml_positions.device, dtype=ml_positions.dtype) - positions = positions.clone().requires_grad_(compute_force) - width = (int(model.atomic_multipoles_max_l) + 1) ** 2 - if multipoles is None: - features = torch.zeros( - (charges.numel(), width), dtype=ml_positions.dtype, device=ml_positions.device - ) - features[:, 0] = charges.to(features).reshape(-1) - else: - features = multipoles.to(device=ml_positions.device, dtype=ml_positions.dtype).clone() - if features.dim() != 2 or features.shape != (positions.shape[0], width): - raise ValueError(f"mm_multipoles must have shape [N_mm, {width}].") - if width >= 4: - # Public Cartesian (q, px, py, pz) -> graph/e3nn (q, py, pz, px). - features[:, 1:4] = features[:, [2, 3, 1]] - if positions.shape[0] != features.shape[0]: - raise ValueError("MM positions and electrostatic sources must have the same length.") - - transform = getattr(model, "_charges_to_mul_ir", None) - if transform is not None: - features = transform(features) - - batch = data.get("mm_source_batch") - if batch is None: - if int(data["pbc"].reshape(-1, 3).shape[0]) != 1: - raise ValueError("mm_source_batch is required for batched PolarMACE inputs.") - batch = torch.zeros(positions.shape[0], dtype=torch.long, device=positions.device) - else: - batch = batch.to(device=positions.device, dtype=torch.long).reshape(-1) - if batch.shape[0] != positions.shape[0]: - raise ValueError("mm_source_batch and mm_positions must have the same length.") - return {"positions": positions, "features": features, "batch": batch} - - -def enable_polarmace_external_sources(model): - """Return an eager wrapper that adds dynamic MM electrostatic sources. - - Models that have already been wrapped are returned unchanged. Non-PolarMACE - models are rejected rather than silently dropping the external interaction. - """ - import torch - - if getattr(model, "supports_external_electrostatics", False): - return model - if model.__class__.__name__ != "PolarMACE": - raise TypeError( - "External electrostatic sources require a PolarMACE model; got " - f"{model.__class__.__name__}." - ) - - try: - from graph_longrange.external_source_energy import ( - GTOElectrostaticExternalSourceEnergy, - ) - from graph_longrange.external_source_features import ( - GTOElectrostaticExternalSourceFeatures, - ) - except ImportError as exc: - raise ImportError( - "PolarMACE electrostatic embedding requires a graph_longrange " - "release that provides the external-source energy and feature blocks." - ) from exc - - feature_base = _rebuild_feature_block(model.electric_potential_descriptor) - energy_base = _rebuild_energy_block(model.coulomb_energy) - feature_block = GTOElectrostaticExternalSourceFeatures.from_features( - feature_base, - # PolarMACE has two spin channels. Each receives half of the physical - # external potential. - external_scale=0.5, - ) - energy_block = GTOElectrostaticExternalSourceEnergy.from_energy(energy_base) - model.electric_potential_descriptor = feature_block - model.coulomb_energy = energy_block - - class PolarMACEExternalSources(torch.nn.Module): - supports_external_electrostatics = True - - def __init__(self, wrapped): - super().__init__() - self.model = wrapped - - def __getattr__(self, name): - try: - return super().__getattr__(name) - except AttributeError: - return getattr(self.model, name) - - def forward( - self, - data, - training: bool = False, - compute_force: bool = True, - compute_virials: bool = False, - compute_stress: bool = False, - compute_displacement: bool = False, - compute_hessian: bool = False, - compute_edge_forces: bool = False, - compute_atomic_stresses: bool = False, - **kwargs, - ): - external = _prepare_external_sources(self.model, data, compute_force) - if external is None: - return self.model( - data, - training=training, - compute_force=compute_force, - compute_virials=compute_virials, - compute_stress=compute_stress, - compute_displacement=compute_displacement, - compute_hessian=compute_hessian, - compute_edge_forces=compute_edge_forces, - compute_atomic_stresses=compute_atomic_stresses, - **kwargs, - ) - if any( - (compute_virials, compute_stress, compute_displacement, compute_hessian, - compute_edge_forces, compute_atomic_stresses) - ): - raise NotImplementedError( - "The OpenMM external-source adapter currently supports energies and " - "Cartesian forces only." - ) - - external_kwargs = { - "external_feats": external["features"], - "external_positions": external["positions"], - "external_batch": external["batch"], - } - self.model.electric_potential_descriptor.set_external_sources( - **external_kwargs - ) - self.model.coulomb_energy.set_external_sources(**external_kwargs) - try: - result = self.model( - data, - training=training, - compute_force=False, - compute_virials=False, - compute_stress=False, - compute_displacement=False, - compute_hessian=False, - compute_edge_forces=False, - compute_atomic_stresses=False, - **kwargs, - ) - if compute_force: - ml_gradient, mm_gradient = torch.autograd.grad( - outputs=[result["energy"]], - inputs=[data["positions"], external["positions"]], - grad_outputs=[torch.ones_like(result["energy"])], - create_graph=training, - retain_graph=training, - allow_unused=True, - ) - result["forces"] = ( - torch.zeros_like(data["positions"]) - if ml_gradient is None else -ml_gradient - ) - result["mm_forces"] = ( - torch.zeros_like(external["positions"]) - if mm_gradient is None else -mm_gradient - ) - else: - result["mm_forces"] = None - return result - finally: - self.model.electric_potential_descriptor.clear_external_sources() - self.model.coulomb_energy.clear_external_sources() - - return PolarMACEExternalSources(model) - - -__all__ = ["enable_polarmace_external_sources"] diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 6eeab59..8d51048 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -28,18 +28,279 @@ OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -import os import openmm from openmm import unit from openmmml.mlpotential import MLPotential, MLPotentialImpl, MLPotentialImplFactory from openmmml.embeddings import utilities -from typing import Iterable, Optional, Sequence, Tuple, Union +from typing import Iterable, Optional from functools import partial -from pathlib import Path import numpy as np -LinkRecordTuple = Tuple[int, int, float] # (q_global, m_global, target_dist) -LinkRecordsArg = Union[str, "os.PathLike[str]", Sequence[LinkRecordTuple], None] + +def _floating_reference(module): + for tensor in module.buffers(): + if tensor.is_floating_point(): + return tensor + for tensor in module.parameters(): + if tensor.is_floating_point(): + return tensor + return None + + +def _rebuild_feature_block(block, pbc_handling: str = "auto"): + """Rebuild a deterministic graph block saved by an older graph release.""" + from graph_longrange.features import GTOElectrostaticFeatures + + realspace = block.realspace_features + quadrupoles = bool( + getattr( + getattr(block.non_periodic_correction_terms, "self_field", None), + "include_quadrupole_corrections", + False, + ) + ) + rebuilt = GTOElectrostaticFeatures( + density_max_l=int(realspace.density_max_l), + density_smearing_width=float(realspace.density_smearing_width), + feature_max_l=int(realspace.projection_max_l), + feature_smearing_widths=[ + float(x) for x in realspace.projection_smearing_widths + ], + include_self_interaction=bool(block.include_self_interaction), + kspace_cutoff=float(block.kspace_cutoff), + quadrupole_feature_corrections=quadrupoles, + integral_normalization=str(block.feature_basis.normalize), + pbc_handling=pbc_handling, + ) + reference = _floating_reference(block) + if reference is not None: + rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) + return rebuilt + + +def _rebuild_energy_block(block, pbc_handling: str = "auto"): + from graph_longrange.energy import GTOElectrostaticEnergy + + rebuilt = GTOElectrostaticEnergy( + density_max_l=int(block.density_max_l), + density_smearing_width=float(block.density_smearing_width), + kspace_cutoff=float(block.kspace_cutoff), + include_self_interaction=bool(block.include_self_interaction), + pbc_handling=pbc_handling, + ) + reference = _floating_reference(block) + if reference is not None: + rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) + return rebuilt + + +def _prepare_external_sources(model, data, compute_force: bool): + import torch + + positions = data.get("mm_positions") + charges = data.get("mm_charges") + multipoles = data.get("mm_multipoles") + if charges is not None and multipoles is not None: + raise ValueError("mm_charges and mm_multipoles are mutually exclusive.") + values = multipoles if multipoles is not None else charges + if ( + positions is None + or values is None + or positions.numel() == 0 + or values.numel() == 0 + ): + return None + + ml_positions = data["positions"] + positions = positions.to(device=ml_positions.device, dtype=ml_positions.dtype) + positions = positions.clone().requires_grad_(compute_force) + width = (int(model.atomic_multipoles_max_l) + 1) ** 2 + if multipoles is None: + features = torch.zeros( + (charges.numel(), width), + dtype=ml_positions.dtype, + device=ml_positions.device, + ) + features[:, 0] = charges.to(features).reshape(-1) + else: + features = multipoles.to( + device=ml_positions.device, dtype=ml_positions.dtype + ).clone() + if features.dim() != 2 or features.shape != (positions.shape[0], width): + raise ValueError(f"mm_multipoles must have shape [N_mm, {width}].") + if width >= 4: + # Public Cartesian (q, px, py, pz) -> graph/e3nn (q, py, pz, px). + features[:, 1:4] = features[:, [2, 3, 1]] + if positions.shape[0] != features.shape[0]: + raise ValueError( + "MM positions and electrostatic sources must have the same length." + ) + + transform = getattr(model, "_charges_to_mul_ir", None) + if transform is not None: + features = transform(features) + + batch = data.get("mm_source_batch") + if batch is None: + if int(data["pbc"].reshape(-1, 3).shape[0]) != 1: + raise ValueError( + "mm_source_batch is required for batched PolarMACE inputs." + ) + batch = torch.zeros( + positions.shape[0], dtype=torch.long, device=positions.device + ) + else: + batch = batch.to(device=positions.device, dtype=torch.long).reshape(-1) + if batch.shape[0] != positions.shape[0]: + raise ValueError( + "mm_source_batch and mm_positions must have the same length." + ) + return {"positions": positions, "features": features, "batch": batch} + + +def _enable_polarmace_external_sources(model): + """Wrap PolarMACE with dynamic MM electrostatic sources in eager mode.""" + import torch + + if getattr(model, "supports_external_electrostatics", False): + return model + if model.__class__.__name__ != "PolarMACE": + raise TypeError( + "External electrostatic sources require a PolarMACE model; got " + f"{model.__class__.__name__}." + ) + + try: + from graph_longrange.external_source_energy import ( + GTOElectrostaticExternalSourceEnergy, + ) + from graph_longrange.external_source_features import ( + GTOElectrostaticExternalSourceFeatures, + ) + except ImportError as exc: + raise ImportError( + "PolarMACE electrostatic embedding requires a graph_longrange " + "release that provides the external-source energy and feature blocks." + ) from exc + + feature_base = _rebuild_feature_block(model.electric_potential_descriptor) + energy_base = _rebuild_energy_block(model.coulomb_energy) + model.electric_potential_descriptor = ( + GTOElectrostaticExternalSourceFeatures.from_features( + feature_base, + # PolarMACE has two spin channels. Each receives half of the + # physical external potential. + external_scale=0.5, + ) + ) + model.coulomb_energy = GTOElectrostaticExternalSourceEnergy.from_energy( + energy_base + ) + + class PolarMACEExternalSources(torch.nn.Module): + supports_external_electrostatics = True + + def __init__(self, wrapped): + super().__init__() + self.model = wrapped + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.model, name) + + def forward( + self, + data, + training: bool = False, + compute_force: bool = True, + compute_virials: bool = False, + compute_stress: bool = False, + compute_displacement: bool = False, + compute_hessian: bool = False, + compute_edge_forces: bool = False, + compute_atomic_stresses: bool = False, + **kwargs, + ): + external = _prepare_external_sources(self.model, data, compute_force) + if external is None: + return self.model( + data, + training=training, + compute_force=compute_force, + compute_virials=compute_virials, + compute_stress=compute_stress, + compute_displacement=compute_displacement, + compute_hessian=compute_hessian, + compute_edge_forces=compute_edge_forces, + compute_atomic_stresses=compute_atomic_stresses, + **kwargs, + ) + if any( + ( + compute_virials, + compute_stress, + compute_displacement, + compute_hessian, + compute_edge_forces, + compute_atomic_stresses, + ) + ): + raise NotImplementedError( + "The OpenMM external-source adapter currently supports " + "energies and Cartesian forces only." + ) + + external_kwargs = { + "external_feats": external["features"], + "external_positions": external["positions"], + "external_batch": external["batch"], + } + self.model.electric_potential_descriptor.set_external_sources( + **external_kwargs + ) + self.model.coulomb_energy.set_external_sources(**external_kwargs) + try: + result = self.model( + data, + training=training, + compute_force=False, + compute_virials=False, + compute_stress=False, + compute_displacement=False, + compute_hessian=False, + compute_edge_forces=False, + compute_atomic_stresses=False, + **kwargs, + ) + if compute_force: + ml_gradient, mm_gradient = torch.autograd.grad( + outputs=[result["energy"]], + inputs=[data["positions"], external["positions"]], + grad_outputs=[torch.ones_like(result["energy"])], + create_graph=training, + retain_graph=training, + allow_unused=True, + ) + result["forces"] = ( + torch.zeros_like(data["positions"]) + if ml_gradient is None + else -ml_gradient + ) + result["mm_forces"] = ( + torch.zeros_like(external["positions"]) + if mm_gradient is None + else -mm_gradient + ) + else: + result["mm_forces"] = None + return result + finally: + self.model.electric_potential_descriptor.clear_external_sources() + self.model.coulomb_energy.clear_external_sources() + + return PolarMACEExternalSources(model) class MACEPotentialImplFactory(MLPotentialImplFactory): @@ -196,11 +457,7 @@ def _loadModel(self, args): else: raise ValueError(f"Unsupported MACE model: {self.name}") if model.__class__.__name__ == "PolarMACE": - from openmmml.models._polarmace_external import ( - enable_polarmace_external_sources, - ) - - model = enable_polarmace_external_sources(model) + model = _enable_polarmace_external_sources(model) return model, device def addForces( @@ -211,8 +468,6 @@ def addForces( forceGroup: int, precision: Optional[str] = None, returnEnergyType: str = "energy", - linkRecords: LinkRecordsArg = None, - linkChargeScheme: str = "dz1", embedding: str = "mechanical", customNonbondedChargeParameter: Optional[str] = None, **args, @@ -241,28 +496,6 @@ def addForces( returns only the message-passing readout; for PolarMACE this is **not** the gradient partner of ``forces`` and will produce apparent NVE drift / a non-zero finite-difference plateau. - linkRecords : str / path / sequence of (q_global, m_global, target_dist) / None - Hydrogen link-atom cap records for QM/MM boundary bonds. - linkChargeScheme : str, optional - How to handle the partial charges on the MM-side boundary atoms - (M atoms) when ``linkRecords`` is provided. The M-atom partial - charge would otherwise sit ~1.5 Å from the nearest QM atom (through - the link H) and over-polarise the QM region's MACE-predicted - electronic structure. Supported in this iteration: - - - ``"none"``: leave MM charges untouched (legacy behaviour). - - ``"z1"``: set q_M = 0 for every M atom. Cheapest fix; breaks - total MM-charge neutrality by -q_M_orig. - - ``"dz1"`` (default): q_M = 0 plus q_M_orig is distributed - equally onto M's MM neighbours (M1 atoms). Preserves total - MM charge to round-off. Falls back to Z1 (with a warning) for - any M atom that has zero MM neighbours. - - Only modifies the MM charge array passed to the ML potential; - the OpenMM ``NonbondedForce`` is left untouched, so MM-MM Coulomb - is bit-exact with the original force field (standard QM/MM - practice). Z2 and RCD (which require per-step virtual charges) - are not in this iteration. embedding : {"mechanical", "electrostatic"} Which embedding method the caller is implementing. Set by ``createMixedSystem``; there is normally no reason to pass it here @@ -314,10 +547,6 @@ def addForces( stacklevel=2, ) - linkInfo = _prepareLinkRecords(linkRecords, atoms, topology, system) - if linkInfo is not None: - atomicNumbers = atomicNumbers + [1] * linkInfo["K"] - modelDefaultDtype = next(model.parameters()).dtype if precision is None: dtype = modelDefaultDtype @@ -356,23 +585,6 @@ def addForces( # embedding='electrostatic'. We only need MM positions/charges for # the PolarMACE input here. mmInfo = _prepareMMEmbedding(system, atoms, customNonbondedChargeParameter) - - # Optional Z1 / DZ1 link-atom charge redistribution. Standard QM/MM - # correction to stop the QM region from being over-polarised by - # the partial charge on the MM-side boundary atom. - if linkInfo is not None and linkChargeScheme not in (None, "none"): - from openmmml.models._links import ( - apply_link_charge_redistribution as _apply_link_q, - ) - mmInfo["mm_charges"] = _apply_link_q( - mm_atoms=mmInfo["mm_atoms"], - mm_charges=mmInfo["mm_charges"], - link_info=linkInfo, - topology=topology, - scheme=linkChargeScheme, - ) - print(f"[link-charge-redistribution] scheme={linkChargeScheme} " - f"M atoms touched={len(linkInfo['m_global'])}") periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions() compute = partial(_computeMACE, @@ -386,7 +598,6 @@ def addForces( multiplicity=torch.tensor([float(args.get('multiplicity', 1))], dtype=dtype, device=model_device, requires_grad=False), indices=indices, periodic=periodic, - linkInfo=linkInfo, mmInfo=mmInfo) force = openmm.PythonForce(compute) force.setForceGroup(forceGroup) @@ -687,7 +898,7 @@ def _prepareMMEmbedding(system: openmm.System, atoms: Optional[Iterable[int]], } -def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, indices, periodic, linkInfo=None, mmInfo=None): +def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, indices, periodic, mmInfo=None): import torch from mace.data.neighborhood import get_neighborhood energyScale = 96.4853 @@ -699,28 +910,6 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch else: positions = positions_full - # Link atoms: append K fictitious H positions placed each step from current - # Q, M coordinates. These extend the model input only; they never enter the - # OpenMM system. See docs/plans/link-atom-inference.md. - if linkInfo is not None: - if indices is None: - raise ValueError("linkRecords requires an explicit `atoms` subset.") - from openmmml.models._links import compute_cap_positions, minimum_image_M - r_Q = positions_full[linkInfo["q_global"]] - r_M = positions_full[linkInfo["m_global"]] - if periodic: - # Minimum-image the Q->M bond so caps are placed correctly even if Q - # and M sit across a periodic boundary. The same imaged r_M is reused - # for force redistribution below, so the returned force stays the - # gradient of the reported energy. Triclinic-correct; raises on a - # genuinely wrapped pair (see minimum_image_M). - cell_A = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) - r_M = minimum_image_M(r_Q, r_M, cell_A) - pos_link, _C_L_unused = compute_cap_positions( - r_Q, r_M, linkInfo["target_dist"] - ) - positions = np.concatenate([positions, pos_link], axis=0) - if periodic: cell = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) else: @@ -777,118 +966,11 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch if mm_forces is not None: mm_forces = (mm_forces * energyScale * lengthScale).detach().cpu().numpy() - # force redistribution and add mm_forces back to the full system + # Scatter ML and MM forces back to the full system. if indices is not None: f = np.zeros((numAtoms, 3), dtype=(np.float64 if dtype == torch.float64 else np.float32)) - if linkInfo is None: - f[indices] = forces - else: - from openmmml.models._links import ( - compute_cap_positions, - minimum_image_M, - redistribute_cap_force, - ) - N = len(indices) - f_ml = forces[:N] - f_link = forces[N:] - f[indices] = f_ml - - # Redistribute each link atom's force onto its (Q, M) partners. - r_Q = positions_full[linkInfo["q_global"]] - r_M = positions_full[linkInfo["m_global"]] - if periodic: - cell_A = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) - r_M = minimum_image_M(r_Q, r_M, cell_A) - _, C_L = compute_cap_positions(r_Q, r_M, linkInfo["target_dist"]) - F_Q_add, F_M_add = redistribute_cap_force(f_link, r_Q, r_M, C_L) - - f[linkInfo["q_global"]] += F_Q_add.astype(f.dtype, copy=False) - f[linkInfo["m_global"]] += F_M_add.astype(f.dtype, copy=False) + f[indices] = forces if mmInfo is not None and mm_forces is not None: f[mmInfo["mm_atoms"]] += mm_forces.astype(f.dtype, copy=False) forces = f return energy, forces - - -def _prepareLinkRecords(linkRecords, atoms, topology, system): - """Normalize the ``linkRecords`` argument into a frozen bundle for the - per-step closure. Returns ``None`` if no link records were supplied. - - Assertions (failing loudly): - * non-periodic system (PBC deferred to a follow-up PR) - * ``atoms`` is not None - * every ``q_global`` is in ``atoms`` - * no ``m_global`` is in ``atoms`` - * every (Q, M) pair is unique, and no atom appears as Q or M in more - than one cap (simplifies per-step scatter; relax later if needed) - """ - if linkRecords is None: - return None - - # Periodic systems are allowed: cap positions are placed with minimum-image - # at runtime (_computeMACE), which is exact as long as each link bond is - # shorter than half the box -- always true for real frontier bonds. A - # genuinely wrapped (Q, M) pair (longer than half the box) raises there. - if atoms is None: - raise ValueError("linkRecords requires an explicit `atoms` subset.") - - if isinstance(linkRecords, (str, Path)): - # A capping-mapping CSV is 1-based: read (q_idx1, m_idx1, - # target_dist_ang) and convert to 0-based OpenMM indices. target_dist - # stays in Angstroms, canonical because it matches the MACE-side - # `positions_full`. Cap positions are not stored in the CSV; they are - # recomputed each step from (q, m, target_dist). - import csv as _csv - tuples = [] - with open(linkRecords, newline="") as _f: - for row in _csv.DictReader(_f): - tuples.append( - (int(row["q_idx1"]) - 1, int(row["m_idx1"]) - 1, float(row["target_dist_ang"])) - ) - else: - # Tuple path: target_dist is also in Å (matches the docstring - # parameter name `target_dist_ang` and the MACE convention used - # in every upstream test fixture, e.g. linkRecords=[(q, m, 1.09)]). - tuples = [(int(q), int(m), float(td)) for (q, m, td) in linkRecords] - - if not tuples: - return None - - num_particles = int(system.getNumParticles()) - atoms_set = set(int(a) for a in atoms) - seen_pairs: set = set() - seen_q: set = set() - seen_m: set = set() - q_global = np.empty(len(tuples), dtype=np.int64) - m_global = np.empty(len(tuples), dtype=np.int64) - target_dist = np.empty(len(tuples), dtype=np.float64) - for k, (q, m, td) in enumerate(tuples): - if not (0 <= q < num_particles): - raise ValueError(f"linkRecords[{k}]: q_global={q} out of range [0, {num_particles}).") - if not (0 <= m < num_particles): - raise ValueError(f"linkRecords[{k}]: m_global={m} out of range [0, {num_particles}).") - if q not in atoms_set: - raise ValueError(f"linkRecords[{k}]: q_global={q} is not in `atoms`.") - if m in atoms_set: - raise ValueError(f"linkRecords[{k}]: m_global={m} is in `atoms` (should be MM).") - if (q, m) in seen_pairs: - raise ValueError(f"linkRecords[{k}]: duplicate (Q, M) pair ({q}, {m}).") - if q in seen_q or m in seen_m: - raise ValueError( - f"linkRecords[{k}]: atom {q if q in seen_q else m} appears in more than " - "one cap; current implementation requires unique Q and M across caps." - ) - if td <= 0: - raise ValueError(f"linkRecords[{k}]: target_dist must be positive; got {td}.") - seen_pairs.add((q, m)) - seen_q.add(q) - seen_m.add(m) - q_global[k] = q - m_global[k] = m - target_dist[k] = td - return { - "K": len(tuples), - "q_global": q_global, - "m_global": m_global, - "target_dist": target_dist, - } diff --git a/test/TestLinkChargeRedistribution.py b/test/TestLinkChargeRedistribution.py deleted file mode 100644 index 18cc498..0000000 --- a/test/TestLinkChargeRedistribution.py +++ /dev/null @@ -1,222 +0,0 @@ -"""Unit tests for Z1 / DZ1 link-atom charge redistribution. - -Covers the new ``apply_link_charge_redistribution`` helper in -``openmmml.models._links``. Does not exercise the full MACE stack — -that's covered by the existing electrostatic-embedding smoke tests. -""" -import numpy as np -import openmm as mm -import openmm.app as app -from openmm.app import element as elem -import pytest - -from openmmml.models._links import apply_link_charge_redistribution - - -# ----------------------------------------------------------------------------- -# Helpers: build minimal Topology fixtures. -# ----------------------------------------------------------------------------- - -def _topology_with_bonds(bonds: list[tuple[int, int]], n_atoms: int = 4): - """Tiny topology: a single chain with the requested bonds. - - All atoms are CARBON (element doesn't matter for these tests). - """ - topo = app.Topology() - chain = topo.addChain() - residue = topo.addResidue("X", chain) - atoms = [topo.addAtom(f"C{i}", elem.carbon, residue) for i in range(n_atoms)] - for a, b in bonds: - topo.addBond(atoms[a], atoms[b]) - return topo - - -def _toy_linkinfo(q_globals, m_globals): - """Minimum link_info dict expected by apply_link_charge_redistribution.""" - return { - "q_global": np.asarray(q_globals, dtype=np.int64), - "m_global": np.asarray(m_globals, dtype=np.int64), - "target_dist": np.full(len(q_globals), 1.09, dtype=np.float64), - } - - -# ----------------------------------------------------------------------------- -# Basic API -# ----------------------------------------------------------------------------- - -def test_none_returns_copy_unchanged(): - mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) - mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10], [1]), - topology=_topology_with_bonds([]), scheme="none", - ) - np.testing.assert_array_equal(out, mm_q) - assert out is not mm_q # must be a copy - - -def test_unknown_scheme_raises(): - with pytest.raises(ValueError, match="Unsupported linkChargeScheme"): - apply_link_charge_redistribution( - np.array([0]), np.array([0.0]), - _toy_linkinfo([], []), - topology=_topology_with_bonds([]), scheme="z3", - ) - - -# ----------------------------------------------------------------------------- -# Z1 -# ----------------------------------------------------------------------------- - -def test_z1_zeros_M_only(): - mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) - mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10], [1]), - topology=_topology_with_bonds([]), scheme="z1", - ) - assert out[1] == 0.0 - np.testing.assert_array_equal(out[[0, 2, 3]], mm_q[[0, 2, 3]]) - - -def test_z1_multiple_M_atoms(): - mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) - mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10, 11], [1, 3]), - topology=_topology_with_bonds([]), scheme="z1", - ) - assert out[1] == 0.0 - assert out[3] == 0.0 - np.testing.assert_array_equal(out[[0, 2]], mm_q[[0, 2]]) - - -# ----------------------------------------------------------------------------- -# DZ1 -# ----------------------------------------------------------------------------- - -def test_dz1_distributes_to_MM_neighbors_and_preserves_total(): - # Bonds: M (atom 1) is bonded to MM atoms 2 and 3, plus Q atom 10. - # The Q-M bond should be ignored (not an MM-MM bond). - # q_M_orig = -0.2 -> each MM neighbour gets -0.1. - topology = _topology_with_bonds([(1, 2), (1, 3)], n_atoms=4) - mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) - mm_q = np.array([0.1, -0.2, 0.3, -0.4], dtype=np.float64) - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10], [1]), - topology=topology, scheme="dz1", - ) - assert out[1] == 0.0 - assert out[2] == pytest.approx(0.3 + (-0.2) / 2) - assert out[3] == pytest.approx(-0.4 + (-0.2) / 2) - # Total preserved to round-off - assert out.sum() == pytest.approx(mm_q.sum(), abs=1e-12) - - -def test_dz1_skips_Q_neighbors(): - # M is bonded to Q (10) AND to MM (2). Only the MM bond counts. - # q_M_orig = -0.5 -> atom 2 gets +(-0.5)/1 = -0.5. - topology = _topology_with_bonds([(1, 2), (1, 10)], n_atoms=11) - mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) - mm_q = np.array([0.1, -0.5, 0.2, 0.2], dtype=np.float64) - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10], [1]), - topology=topology, scheme="dz1", - ) - assert out[1] == 0.0 - assert out[2] == pytest.approx(0.2 - 0.5) - # atom 3 unchanged - assert out[3] == pytest.approx(0.2) - - -def test_dz1_no_MM_neighbors_warns_and_falls_back_to_Z1(): - # M (atom 1) has no MM bonds — only bonded to Q (atom 10). - topology = _topology_with_bonds([(1, 10)], n_atoms=11) - mm_atoms = np.array([0, 1], dtype=np.int64) - mm_q = np.array([0.5, -0.5], dtype=np.float64) - with pytest.warns(UserWarning, match="no MM neighbours"): - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10], [1]), - topology=topology, scheme="dz1", - ) - assert out[1] == 0.0 - # Total not preserved here — warning told us. - assert out.sum() != pytest.approx(mm_q.sum(), abs=1e-12) - - -def test_dz1_multiple_M_atoms_each_get_own_neighbors(): - # M atoms = [1, 3]. Each is bonded to a different MM neighbour (0 and 2). - # q_orig: atom 1 = -0.4 -> atom 0 += -0.4 - # atom 3 = +0.6 -> atom 2 += +0.6 - topology = _topology_with_bonds([(1, 0), (3, 2)], n_atoms=4) - mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) - mm_q = np.array([0.0, -0.4, 0.0, 0.6], dtype=np.float64) - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10, 11], [1, 3]), - topology=topology, scheme="dz1", - ) - assert out[1] == 0.0 - assert out[3] == 0.0 - assert out[0] == pytest.approx(-0.4) - assert out[2] == pytest.approx(0.6) - assert out.sum() == pytest.approx(mm_q.sum(), abs=1e-12) - - -def test_dz1_split_across_three_neighbors(): - # M (atom 1) has three MM neighbours (0, 2, 3). q_M_orig=-0.3 -> each +(-0.1). - topology = _topology_with_bonds([(1, 0), (1, 2), (1, 3)], n_atoms=4) - mm_atoms = np.array([0, 1, 2, 3], dtype=np.int64) - mm_q = np.array([0.0, -0.3, 0.0, 0.0], dtype=np.float64) - out = apply_link_charge_redistribution( - mm_atoms, mm_q, _toy_linkinfo([10], [1]), - topology=topology, scheme="dz1", - ) - assert out[1] == 0.0 - for idx in (0, 2, 3): - assert out[idx] == pytest.approx(-0.1) - assert out.sum() == pytest.approx(mm_q.sum(), abs=1e-12) - - -def test_dz1_with_two_bonded_M_atoms(): - """Two M atoms bonded to each other must both end up at exactly zero. - - Redistributing onto a neighbour that is itself an M atom hands it charge - that its own turn then discards, so the outcome depended on the order the - link records were listed in and left one M atom holding charge -- the very - charge DZ1 exists to move away from the ML region. - """ - import numpy as np - import openmm.app as app - from openmm.app import element - - topology = app.Topology() - chain = topology.addChain() - residue = topology.addResidue("X", chain) - atoms = [topology.addAtom(f"H{i}", element.hydrogen, residue) for i in range(6)] - # ML = {0, 1}; caps Q0->M2 and Q1->M3; M2 and M3 are bonded to each other. - for i, j in [(2, 3), (2, 4), (3, 5)]: - topology.addBond(atoms[i], atoms[j]) - - mm_atoms = np.array([2, 3, 4, 5]) - mm_charges = np.array([-0.5, -0.3, 0.2, 0.1]) - - out = apply_link_charge_redistribution( - mm_atoms=mm_atoms, - mm_charges=mm_charges, - link_info={"q_global": np.array([0, 1]), "m_global": np.array([2, 3])}, - topology=topology, - scheme="dz1", - ) - - assert out[0] == 0.0 and out[1] == 0.0 - assert out.sum() == pytest.approx(mm_charges.sum(), abs=1e-12) - - # And the answer must not depend on the order of the link records. - reversed_order = apply_link_charge_redistribution( - mm_atoms=mm_atoms, - mm_charges=mm_charges, - link_info={"q_global": np.array([1, 0]), "m_global": np.array([3, 2])}, - topology=topology, - scheme="dz1", - ) - np.testing.assert_allclose(out, reversed_order, atol=1e-12) diff --git a/test/TestLinkRecordsTargetDistUnits.py b/test/TestLinkRecordsTargetDistUnits.py deleted file mode 100644 index 42143ac..0000000 --- a/test/TestLinkRecordsTargetDistUnits.py +++ /dev/null @@ -1,147 +0,0 @@ -"""Regression test for the `target_dist` unit canonicalisation in -`_prepareLinkRecords` and its downstream consumers. - -The historical bug: when `linkRecords` was loaded from a -`capping_mapping.csv`, the loader -multiplied the CSV's `target_dist_ang` column by 0.1 — turning a 1.09 Å -record into a 0.109 nm record stored in `linkInfo["target_dist"]`. But -`_computeMACE` reads `linkInfo["target_dist"]` together with `r_Q`, -`r_M` already in Å (via `positions_full = state.getPositions(asNumpy=True -).value_in_unit(unit.angstrom)`). The unit mismatch placed CSV-path -caps at ~0.1 Å from the Q atom — about 10× closer than the intended -~1.09 Å (typical C-H bond). - -The fix: `linkInfo["target_dist"]` is canonically in **Å**. The CSV path -no longer multiplies by 0.1; the tuple path was already in Å. - -Tests in this file pin the contract at three levels so a regression on -*any* of them surfaces clearly: - -1. Tuple-path loader stores the value verbatim (Å). -2. CSV-path loader stores the value verbatim (Å, no `* 0.1` conversion). -3. Downstream end-to-end: `compute_cap_positions` placed against - linkInfo + Å-scale Q/M positions gives a Q-L distance equal to the - target value. Any future "fix" that re-introduces `* 0.1` at the - consumer side (instead of the loader side) trips this end-to-end - case even if it leaves the loader unchanged. -""" -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest - - -@pytest.fixture -def tiny_system_topology(): - """Minimal 5-atom system + topology so _prepareLinkRecords' validation - (q in atoms, m not in atoms, etc.) is satisfied. Atoms 0,1,2 are ML; - atoms 3,4 are MM.""" - import openmm as mm - import openmm.app as app - from openmm.app import element as elem - - system = mm.System() - for _ in range(5): - system.addParticle(1.0) - top = app.Topology() - chain = top.addChain() - res = top.addResidue("X", chain) - atoms = [top.addAtom(f"H{i}", elem.hydrogen, res) for i in range(5)] - for a, b in [(0, 1), (1, 2), (2, 3), (3, 4)]: - top.addBond(atoms[a], atoms[b]) - return top, system - - -def test_tuple_path_target_dist_is_angstrom(tiny_system_topology): - """Tuple path: linkInfo['target_dist'] equals the value the caller passed. - - Locks the canonical interpretation as Å. - """ - from openmmml.models.macepotential import _prepareLinkRecords - - top, system = tiny_system_topology - link_info = _prepareLinkRecords( - linkRecords=[(2, 3, 1.09)], - atoms=[0, 1, 2], - topology=top, - system=system, - ) - assert link_info is not None - assert link_info["K"] == 1 - np.testing.assert_allclose(link_info["target_dist"], np.array([1.09])) - - -# The CSV-path test requires the MLMM-side `load_link_records` parser. -# openmm-ml proper does not depend on MLMM (it's an optional consumer), -# so the test skips cleanly when `mlmm` isn't on the path — keeping -# openmm-ml's own CI (which doesn't install MLMM) reproducible. -_mlmm_available = pytest.importorskip.__doc__ is not None # always True; placeholder -try: - import mlmm.link_atoms # noqa: F401 - _HAS_MLMM = True -except Exception: - _HAS_MLMM = False - - -@pytest.mark.skipif(not _HAS_MLMM, reason="MLMM not installed; CSV path needs mlmm.link_atoms.load_link_records") -def test_csv_path_target_dist_in_angstrom(tmp_path, tiny_system_topology): - """CSV path: linkInfo['target_dist'] matches the CSV's target_dist_ang - column verbatim (no spurious 0.1 conversion). - - Historical bug had this multiplied by 0.1, producing - linkInfo['target_dist'] = 0.109 for a CSV target_dist_ang of 1.09, - which then placed MACE caps 10x too close to Q. - """ - from openmmml.models.macepotential import _prepareLinkRecords - - csv = tmp_path / "capping_mapping.csv" - csv.write_text( - "q_idx1,m_idx1,q_element,m_element,bond_order,qm_mm_distance_ang," - "target_dist_ang,r_Q_x,r_Q_y,r_Q_z,r_M_x,r_M_y,r_M_z,r_L_x,r_L_y,r_L_z\n" - "3,4,H,H,1.000,1.500000,1.090000," - "0.0,0.0,0.0,1.5,0.0,0.0,1.09,0.0,0.0\n" - ) - - top, system = tiny_system_topology - link_info = _prepareLinkRecords( - linkRecords=str(csv), - atoms=[0, 1, 2], - topology=top, - system=system, - ) - assert link_info is not None - assert link_info["K"] == 1 - np.testing.assert_allclose( - link_info["target_dist"], - np.array([1.09]), - err_msg=( - "CSV target_dist_ang=1.09 must land in linkInfo as 1.09 (Å). " - "If you see 0.109 here, the historical *0.1 nm conversion has " - "regressed; see openmm-ml fix/target-dist-unit-canonical-angstrom." - ), - ) - - -def test_compute_cap_positions_q_to_l_distance_matches_target(): - """End-to-end contract: with linkInfo['target_dist'] in Å and Q,M - positions in Å, the cap placement helper produces a Q-L distance - equal to target_dist Å. - - This is the consumer-side test codex review of PR #20 called out: a - future regression that re-introduces ``* 0.1`` at the consumer side - (e.g. ``compute_cap_positions(r_Q, r_M, linkInfo['target_dist'] * 0.1)``) - would slip past the loader-only test but trip this one. - """ - from openmmml.models._links import compute_cap_positions - - # Q at origin, M at (1.5, 0, 0) — a typical C-C single bond, 1.5 Å. - r_Q = np.array([[0.0, 0.0, 0.0]]) - r_M = np.array([[1.5, 0.0, 0.0]]) - target_dist_ang = np.array([1.09]) # typical C-H, in Å - - cap_pos, C_L = compute_cap_positions(r_Q, r_M, target_dist_ang) - q_to_l = np.linalg.norm(cap_pos - r_Q, axis=-1) - np.testing.assert_allclose(q_to_l, target_dist_ang, atol=1e-12) - np.testing.assert_allclose(C_L, target_dist_ang / 1.5, atol=1e-12) diff --git a/test/TestMACEPotential.py b/test/TestMACEPotential.py index 8b8e729..af46d17 100644 --- a/test/TestMACEPotential.py +++ b/test/TestMACEPotential.py @@ -111,7 +111,6 @@ def testComputeMACEScattersMMForces(): multiplicity=torch.tensor([1.0], dtype=torch.float32), indices=np.array([0, 1], dtype=np.int64), periodic=False, - linkInfo=None, mmInfo=mm_info, ) @@ -279,6 +278,5 @@ def __call__(self, inputDict, compute_force=True): multiplicity=torch.tensor([1.0], dtype=torch.float32), indices=np.array([0, 1], dtype=np.int64), periodic=False, - linkInfo=None, mmInfo=mmInfo, ) diff --git a/test/TestPolarMACEElectrostaticEmbedding.py b/test/TestPolarMACEElectrostaticEmbedding.py index 0bfa01e..67ddfc6 100644 --- a/test/TestPolarMACEElectrostaticEmbedding.py +++ b/test/TestPolarMACEElectrostaticEmbedding.py @@ -464,195 +464,3 @@ def test_polar_mace_electrostatic_mm_charge_displacement_changes_ml_force( # responds to the shifted MM field. delta_ml = f_b[:3] - f_a[:3] assert np.linalg.norm(delta_ml) > 1e-3 - - -# --------------------------------------------------------------------------- -# Link-atom CHARGE redistribution (Z1 / DZ1) — end-to-end via MACE -# --------------------------------------------------------------------------- -# -# The unit tests in TestLinkChargeRedistribution.py cover the pure helper -# (apply_link_charge_redistribution). The tests below exercise the actual -# integration into MACEPotentialImpl.addForces and verify that: -# (a) the scheme parameter propagates all the way to MACE's mm_charges -# input (different schemes give different energies/forces); -# (b) forces remain conservative under each scheme (F = -dE/dx, FD-tested). - - -def _build_capped_nonpbc_chain(): - """Tiny 6-atom chain: ML region {0,1,2} bonded to MM atom 3 (the M atom) - which is bonded to MM atoms {4,5} (M's neighbours, for DZ1 to populate). - All MM atoms carry non-zero partial charges so the linkChargeScheme has - something to redistribute. - """ - system = openmm.System() - # Make all six atoms hydrogen (element 1) so MACE's atomic-number map - # exercised by the polar_mace_model_path fixture works (it knows H+O only). - for _ in range(6): - system.addParticle(1.0) - nb = openmm.NonbondedForce() - # ML atoms get zero charge (they'll be re-zeroed by createMixedSystem anyway, - # but make it explicit). MM atoms get something the scheme can act on. - nb.addParticle(0.0, 0.30, 0.20) # 0 ML - nb.addParticle(0.0, 0.30, 0.20) # 1 ML - nb.addParticle(0.0, 0.30, 0.20) # 2 ML (Q atom in the cut) - nb.addParticle(-0.40, 0.30, 0.20) # 3 MM (the M atom) - nb.addParticle(0.15, 0.30, 0.20) # 4 MM (M1 neighbour) - nb.addParticle(0.25, 0.30, 0.20) # 5 MM (M1 neighbour) - nb.setNonbondedMethod(openmm.NonbondedForce.NoCutoff) - system.addForce(nb) - - topology = app.Topology() - chain = topology.addChain() - res = topology.addResidue("X", chain) - atoms = [topology.addAtom(f"H{i}", elem.hydrogen, res) for i in range(6)] - # Bonds: 0-1, 1-2, 2-3 (Q-M cut), 3-4 (M-M1), 3-5 (M-M1) - for a, b in [(0, 1), (1, 2), (2, 3), (3, 4), (3, 5)]: - topology.addBond(atoms[a], atoms[b]) - return topology, system - - -def _capped_chain_positions(): - return np.array([ - [0.00, 0.0, 0.0], - [0.11, 0.0, 0.0], - [0.22, 0.0, 0.0], - [0.33, 0.0, 0.0], # M atom - [0.40, 0.07, 0.0], # M1 - [0.40, -0.07, 0.0], # M1 - ]) * unit.nanometer - - -def _build_capped_mixed(potential, scheme): - topology, mm_system = _build_capped_nonpbc_chain() - mixed = potential.createMixedSystem( - topology, mm_system, atoms=[0, 1, 2], - embedding="electrostatic", - linkRecords=[(2, 3, 1.09)], # Q=2 (in ML), M=3 (MM) - linkChargeScheme=scheme, - ) - plat = openmm.Platform.getPlatformByName("Reference") - ctx = openmm.Context(mixed, openmm.VerletIntegrator(0.001), plat) - ctx.setPositions(_capped_chain_positions()) - return mixed, ctx - - -# elem import for the topology builder; sits with the openmm.app namespace. -from openmm.app import element as elem # noqa: E402 (kept near point of use) - - -def _mmChargesReachingModel(potential, scheme): - """The mm_charges array that createMixedSystem bakes into the ML force. - - Asserting on the array rather than on the energy it produces is deliberate. - The energy route would only distinguish the schemes if the model in use - responded to the MM charge *values*, and the synthetic checkpoint built by - the polar_mace_model_path fixture does not for this fixture's geometry: its - energy is bit-identical whether mm_charges are passed unchanged, zeroed, or - scaled tenfold. What the redistribution schemes are responsible for is the - contents of that array, so that is what these tests pin, exactly rather - than as an inequality. - """ - captured = {} - original = macepotential._computeMACE - - def capture(state, **args): - if args.get("mmInfo") is not None and "mm_charges" not in captured: - captured["mm_atoms"] = np.asarray(args["mmInfo"]["mm_atoms"]) - captured["mm_charges"] = np.asarray(args["mmInfo"]["mm_charges"]) - return original(state, **args) - - macepotential._computeMACE = capture - try: - _, context = _build_capped_mixed(potential, scheme=scheme) - context.getState(getEnergy=True) - finally: - macepotential._computeMACE = original - assert "mm_charges" in captured, "the ML force never received MM charges" - return captured["mm_atoms"], captured["mm_charges"] - - -def test_link_charge_scheme_none_passes_forcefield_charges(polar_mace_model_path): - """With scheme='none' the model sees the MM force field charges verbatim.""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - mmAtoms, mmCharges = _mmChargesReachingModel(potential, "none") - # Atoms 3 (the M atom), 4 and 5 (its MM neighbours) from - # _build_capped_nonpbc_chain. - np.testing.assert_array_equal(mmAtoms, [3, 4, 5]) - np.testing.assert_allclose(mmCharges, [-0.40, 0.15, 0.25], atol=1e-12) - - -def test_link_charge_scheme_z1_zeros_the_m_atom(polar_mace_model_path): - """Z1 zeros the M-atom charge and leaves every other MM charge alone, - which changes the total MM charge by -q_M.""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - _, mmCharges = _mmChargesReachingModel(potential, "z1") - np.testing.assert_allclose(mmCharges, [0.0, 0.15, 0.25], atol=1e-12) - assert mmCharges.sum() == pytest.approx(0.40, abs=1e-12) - - -def test_link_charge_scheme_dz1_spreads_the_m_charge_and_conserves(polar_mace_model_path): - """DZ1 zeros the M-atom charge and spreads it over M's two MM neighbours, - -0.40/2 = -0.20 each, leaving the total MM charge unchanged.""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - _, mmCharges = _mmChargesReachingModel(potential, "dz1") - np.testing.assert_allclose(mmCharges, [0.0, 0.15 - 0.20, 0.25 - 0.20], atol=1e-12) - assert mmCharges.sum() == pytest.approx(0.0, abs=1e-12) - - -@pytest.mark.parametrize("scheme", ["none", "z1", "dz1"]) -def test_link_charge_scheme_fd_force_consistency(polar_mace_model_path, scheme): - """For each scheme, the MACE-side forces must equal the negative gradient - of the MACE-side energy. Central finite difference in f64. - - Atoms 2 (Q) and 3 (M) are included deliberately: the cap sits between them - and is repositioned from their coordinates every step, so they are the only - atoms whose forces exercise the cap chain rule in redistribute_cap_force. - Perturbing only the interior QM atoms would leave that untested. - - The criterion is relative to the size of the force being checked, not an - absolute number of kJ/mol/nm. This runs against a synthetic *untrained* - checkpoint at an arbitrary geometry, where forces reach ~1e7 kJ/mol/nm on - the boundary atoms and only ~1e6 in the interior; a fixed absolute - threshold is meaninglessly loose for the former and tight for the latter. - The central-difference truncation error scales as delta^2 while a genuine - force/energy inconsistency does not shrink with delta at all, which is what - separates the two and what this bound is sized against. - """ - potential = MLPotential("mace", modelPath=polar_mace_model_path) - _, ctx = _build_capped_mixed(potential, scheme=scheme) - - pos0 = _capped_chain_positions().value_in_unit(unit.nanometer).copy() - delta = 1e-5 # nm; the truncation term goes as delta^2 - - ctx.setPositions(pos0 * unit.nanometer) - e_ref = ctx.getState(getEnergy=True).getPotentialEnergy().value_in_unit( - unit.kilojoules_per_mole) - f_ref = ctx.getState(getForces=True).getForces(asNumpy=True).value_in_unit( - unit.kilojoules_per_mole / unit.nanometer) - - max_abs_err = 0.0 - # Atoms 2 (Q) and 3 (M) are the ones the cap is placed between, so they are - # the only ones whose forces exercise the cap chain rule. Perturbing just - # the interior QM atoms 0 and 1 would leave the redistribution untested. - for ia in (0, 1, 2, 3): - for ax in range(3): - saved = pos0[ia, ax] - pos0[ia, ax] = saved + delta - ctx.setPositions(pos0 * unit.nanometer) - e_p = ctx.getState(getEnergy=True).getPotentialEnergy().value_in_unit( - unit.kilojoules_per_mole) - pos0[ia, ax] = saved - delta - ctx.setPositions(pos0 * unit.nanometer) - e_m = ctx.getState(getEnergy=True).getPotentialEnergy().value_in_unit( - unit.kilojoules_per_mole) - pos0[ia, ax] = saved - f_num = -(e_p - e_m) / (2.0 * delta) - err = abs(f_num - f_ref[ia, ax]) - max_abs_err = max(max_abs_err, err) - forceScale = float(np.abs(f_ref).max()) - max_rel_err = max_abs_err / forceScale - assert max_rel_err < 5e-5, ( - f"FD vs analytical force mismatch for scheme='{scheme}': " - f"max |dF| = {max_abs_err:.3e} kJ/mol/nm against a force scale of " - f"{forceScale:.3e}, i.e. relative {max_rel_err:.3e} (expected < 5e-5)" - ) From 76427cca7379ce755ceb23fdbae7631d22e9d250 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 14:37:36 +0000 Subject: [PATCH 05/20] Restrict external MM sources to point charges --- openmmml/models/macepotential.py | 30 ++++++++---------------------- 1 file changed, 8 insertions(+), 22 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 8d51048..aa7b4fe 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -99,15 +99,11 @@ def _prepare_external_sources(model, data, compute_force: bool): positions = data.get("mm_positions") charges = data.get("mm_charges") - multipoles = data.get("mm_multipoles") - if charges is not None and multipoles is not None: - raise ValueError("mm_charges and mm_multipoles are mutually exclusive.") - values = multipoles if multipoles is not None else charges if ( positions is None - or values is None + or charges is None or positions.numel() == 0 - or values.numel() == 0 + or charges.numel() == 0 ): return None @@ -115,22 +111,12 @@ def _prepare_external_sources(model, data, compute_force: bool): positions = positions.to(device=ml_positions.device, dtype=ml_positions.dtype) positions = positions.clone().requires_grad_(compute_force) width = (int(model.atomic_multipoles_max_l) + 1) ** 2 - if multipoles is None: - features = torch.zeros( - (charges.numel(), width), - dtype=ml_positions.dtype, - device=ml_positions.device, - ) - features[:, 0] = charges.to(features).reshape(-1) - else: - features = multipoles.to( - device=ml_positions.device, dtype=ml_positions.dtype - ).clone() - if features.dim() != 2 or features.shape != (positions.shape[0], width): - raise ValueError(f"mm_multipoles must have shape [N_mm, {width}].") - if width >= 4: - # Public Cartesian (q, px, py, pz) -> graph/e3nn (q, py, pz, px). - features[:, 1:4] = features[:, [2, 3, 1]] + features = torch.zeros( + (charges.numel(), width), + dtype=ml_positions.dtype, + device=ml_positions.device, + ) + features[:, 0] = charges.to(features).reshape(-1) if positions.shape[0] != features.shape[0]: raise ValueError( "MM positions and electrostatic sources must have the same length." From ec8510bd2b09f461f4e17fbb5fd2912440039be3 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 17:30:28 +0000 Subject: [PATCH 06/20] Rely on MACE Polar compatibility repair --- openmmml/models/macepotential.py | 63 +------------------------------- 1 file changed, 2 insertions(+), 61 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index aa7b4fe..5998447 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -37,63 +37,6 @@ import numpy as np -def _floating_reference(module): - for tensor in module.buffers(): - if tensor.is_floating_point(): - return tensor - for tensor in module.parameters(): - if tensor.is_floating_point(): - return tensor - return None - - -def _rebuild_feature_block(block, pbc_handling: str = "auto"): - """Rebuild a deterministic graph block saved by an older graph release.""" - from graph_longrange.features import GTOElectrostaticFeatures - - realspace = block.realspace_features - quadrupoles = bool( - getattr( - getattr(block.non_periodic_correction_terms, "self_field", None), - "include_quadrupole_corrections", - False, - ) - ) - rebuilt = GTOElectrostaticFeatures( - density_max_l=int(realspace.density_max_l), - density_smearing_width=float(realspace.density_smearing_width), - feature_max_l=int(realspace.projection_max_l), - feature_smearing_widths=[ - float(x) for x in realspace.projection_smearing_widths - ], - include_self_interaction=bool(block.include_self_interaction), - kspace_cutoff=float(block.kspace_cutoff), - quadrupole_feature_corrections=quadrupoles, - integral_normalization=str(block.feature_basis.normalize), - pbc_handling=pbc_handling, - ) - reference = _floating_reference(block) - if reference is not None: - rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) - return rebuilt - - -def _rebuild_energy_block(block, pbc_handling: str = "auto"): - from graph_longrange.energy import GTOElectrostaticEnergy - - rebuilt = GTOElectrostaticEnergy( - density_max_l=int(block.density_max_l), - density_smearing_width=float(block.density_smearing_width), - kspace_cutoff=float(block.kspace_cutoff), - include_self_interaction=bool(block.include_self_interaction), - pbc_handling=pbc_handling, - ) - reference = _floating_reference(block) - if reference is not None: - rebuilt = rebuilt.to(device=reference.device, dtype=reference.dtype) - return rebuilt - - def _prepare_external_sources(model, data, compute_force: bool): import torch @@ -169,18 +112,16 @@ def _enable_polarmace_external_sources(model): "release that provides the external-source energy and feature blocks." ) from exc - feature_base = _rebuild_feature_block(model.electric_potential_descriptor) - energy_base = _rebuild_energy_block(model.coulomb_energy) model.electric_potential_descriptor = ( GTOElectrostaticExternalSourceFeatures.from_features( - feature_base, + model.electric_potential_descriptor, # PolarMACE has two spin channels. Each receives half of the # physical external potential. external_scale=0.5, ) ) model.coulomb_energy = GTOElectrostaticExternalSourceEnergy.from_energy( - energy_base + model.coulomb_energy ) class PolarMACEExternalSources(torch.nn.Module): From b56181c2fc4f22f30dc6d2fc92e39021e7a3f085 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 17:43:06 +0000 Subject: [PATCH 07/20] Keep user guide unchanged --- doc/userguide.md | 35 ++--------------------------------- 1 file changed, 2 insertions(+), 33 deletions(-) diff --git a/doc/userguide.md b/doc/userguide.md index 97b0727..2c579b5 100644 --- a/doc/userguide.md +++ b/doc/userguide.md @@ -104,8 +104,7 @@ When using MACE models, the following extra keyword arguments to `createSystem() | Argument | Description | | --- | --- | | `precision` | The numerical precision of the model. Supported options are `'single'` and `'double'`. If `None`, the default precision of the model is used. | -| `returnEnergyType` | Whether to return the interaction energy or the energy including the self-energy. The default is `'energy'`, which is the scalar the reported forces are the gradient of, so the potential is conservative. Supported options are `'interaction_energy'` and `'energy'`. | -| `embedding` | Passed by the embedding method rather than by you; see *Embeddings* below. | +| `returnEnergyType` | Whether to return the interaction energy or the energy including the self-energy. The default is `'interaction_energy'`. Supported options are `'interaction_energy'` and `'energy'`. | | `device` | The PyTorch device to perform calculations on, either a `torch.device` object or a string (such as `'cuda'` or `'cpu'`.) If omitted, a device is chosen automatically. | | `charge` | The total charge of the system. If omitted, it is assumed to be 0. This is only used by MACE-OMOL-0. For other models it is ignored. | | `multiplicity` | The spin multiplicity of the system. If omitted, it is assumed to be 1. This is only used by MACE-OMOL-0. For other models it is ignored. | @@ -382,37 +381,7 @@ 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. -### Electrostatic Embedding - -This is a potential-specific embedding method provided by the MACE interface rather than a generic one, and it is -selected with the embedding name `electrostatic`. The MLIP, rather than the MM force field, computes the electrostatic -interactions between the ML and MM atoms: it is given the positions and MM force field charges of the MM atoms, and -returns forces on them alongside the forces on the ML atoms. The ML subset can therefore polarise in response to its -surroundings, which mechanical embedding does not allow. Lennard-Jones interactions between the ML and MM atoms are -still computed by the MM force field, as are all bonded terms that cross the ML/MM boundary. - -This is implemented by setting the MM force field charge of every ML atom to zero, so that every Coulomb term involving -an ML atom vanishes, including the reciprocal space part of PME. As a result, the MM force field's own charges are -untouched and MM-MM electrostatics are unchanged. - -Only PolarMACE models can be used with this embedding method. This includes the pretrained -`mace-polar-1-small`, `mace-polar-1-medium`, and `mace-polar-1-large` models, as well as a PolarMACE checkpoint loaded -with the model name `mace` and a `modelPath`. For a custom checkpoint the method is listed because its type cannot be -known until it is loaded; a non-PolarMACE checkpoint is rejected at that point rather than silently falling back to -mechanical embedding, since the mixed system has already had its ML-MM electrostatics removed. - -The damped real-space and reciprocal-space ML-MM cross energy, including slab and molecule-in-box corrections, and the -external-source feature and energy wrappers are implemented in `graph_longrange`. OpenMM-ML only supplies the external -positions and charges to those blocks, then returns the corresponding MM forces. The normalization is therefore -inherited from the checkpoint's own GTO density and feature bases rather than duplicated in the OpenMM interface. -Interpolation is not supported, because at `lambda_interpolate=0` the ML-MM electrostatics would be missing from the MM -endpoint. - -```python -system = potential.createMixedSystem(topology, mm_system, ml_atoms, embedding='electrostatic') -``` - -### Molecules Spanning the ML-MM Region +#### 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, From 409dcd392d2dcf2da26a8c34007aaab22a20290d Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 17:44:56 +0000 Subject: [PATCH 08/20] Simplify MACE model loading --- openmmml/models/macepotential.py | 28 ++++++++++------------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 5998447..e22cc1a 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -343,14 +343,7 @@ def __init__(self, name: str, modelPath) -> None: self._preloadedModel = None def _loadModel(self, args): - """Load the MACE model, returning it along with the device it is on. - - If createMixedSystem() has already loaded a model in order to inspect - it, that one is handed over here rather than the checkpoint being read a - second time. The handover is consumed on use, so each call after that - loads afresh; holding the model indefinitely would mean a later call - with a different precision converting an already converted model. - """ + """Load a MACE model and place it on the requested device.""" import torch try: from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol, mace_polar @@ -363,24 +356,23 @@ def _loadModel(self, args): return preloaded[0], device if self.name in MACEPotentialImpl.KNOWN_MODELS: - functions = { + loaders = { 'mace_off': mace_off, 'mace_mp': mace_mp, 'mace_omol': mace_omol, 'mace_polar': mace_polar, } - fnName, name, restrictiveLicense, _, _ = MACEPotentialImpl.KNOWN_MODELS[self.name] - model = functions[fnName](model=name, device=device, return_raw_model=True).to(device) - if restrictiveLicense is not None: + loader_name, model_name, restrictive_license, _, _ = self.KNOWN_MODELS[self.name] + model = loaders[loader_name](model=model_name, device=device, return_raw_model=True).to(device) + if restrictive_license is not None: import logging - logging.warning(f'The model {self.name} is distributed under the restrictive {restrictiveLicense} license. Commercial use is not permitted.') + logging.warning(f'The model {self.name} is distributed under the restrictive {restrictive_license} license. Commercial use is not permitted.') elif self.name == "mace": - if self.modelPath is not None: - model = torch.load(self.modelPath, map_location=device) - if hasattr(model, "to"): - model = model.to(device) - else: + if self.modelPath is None: raise ValueError("No modelPath provided for local MACE model.") + model = torch.load(self.modelPath, map_location=device) + if hasattr(model, "to"): + model = model.to(device) else: raise ValueError(f"Unsupported MACE model: {self.name}") if model.__class__.__name__ == "PolarMACE": From 765661b1f775f8b88b78cc802af7065ab256d05b Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 18:54:37 +0000 Subject: [PATCH 09/20] Use interaction energy by default --- openmmml/models/macepotential.py | 45 +++++++++----------------------- 1 file changed, 12 insertions(+), 33 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index e22cc1a..de38251 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -269,17 +269,14 @@ class MACEPotentialImpl(MLPotentialImpl): According to the MACE documentation, 'single' precision is recommended for MD (faster but less accurate), while 'double' precision is recommended for geometry optimization. - By default the reported energy is the full ``energy`` returned by the MACE - model — the same scalar whose gradient w.r.t. positions is reported as the - force, so the resulting potential is exactly conservative. To get only the - message-passing readout component, set ``returnEnergyType='interaction_energy'``: + By default the reported energy is ``interaction_energy``. PolarMACE is an + exception: it automatically uses the full ``energy`` output because that is + the scalar whose gradient w.r.t. positions is reported as the force. >>> system = potential.createSystem(topology, returnEnergyType='interaction_energy') - Note: ``returnEnergyType='interaction_energy'`` is **not** energy/force - consistent for the PolarMACE family, which adds Coulomb / dipole / local- - electron terms to ``total_energy`` whose gradients are in ``forces`` but - which are not in ``interaction_energy``. + PolarMACE therefore uses ``energy`` automatically even when + ``returnEnergyType`` is left at its default. Precision caveat for ``returnEnergyType='energy'``: this key returns the full ``total_energy = e0 + inter_e + extras`` where ``e0`` are the @@ -292,7 +289,7 @@ class MACEPotentialImpl(MLPotentialImpl): magnitude — meV resolution is lost. Use ``precision='double'`` if you need accurate absolute energies, or note that energy differences (e.g. NVE drift) still resolve cleanly because the e0 contribution cancels in - the difference. A runtime warning fires when this regime is detected. + the difference. Attributes ---------- @@ -386,7 +383,7 @@ def addForces( atoms: Optional[Iterable[int]], forceGroup: int, precision: Optional[str] = None, - returnEnergyType: str = "energy", + returnEnergyType: str = "interaction_energy", embedding: str = "mechanical", customNonbondedChargeParameter: Optional[str] = None, **args, @@ -409,12 +406,9 @@ def addForces( If ``None``, the default precision of the model is used. returnEnergyType : str, optional Which scalar from the MACE model output is reported to OpenMM as - the potential energy. Default ``'energy'`` is the same quantity - the force vector is differentiated against, so OpenMM sees a - self-consistent (conservative) potential. ``'interaction_energy'`` - returns only the message-passing readout; for PolarMACE this is - **not** the gradient partner of ``forces`` and will produce - apparent NVE drift / a non-zero finite-difference plateau. + the potential energy. The default is ``'interaction_energy'`` for + ordinary MACE and ``'energy'`` for PolarMACE, whose force gradient + includes additional electrostatic terms. embedding : {"mechanical", "electrostatic"} Which embedding method the caller is implementing. Set by ``createMixedSystem``; there is normally no reason to pass it here @@ -441,6 +435,8 @@ def addForces( # Load the model. model, device = self._loadModel(args) + if model.__class__.__name__ == "PolarMACE": + returnEnergyType = "energy" use_mm_embedding = _should_use_mm_embedding(model, atoms, embedding) @@ -449,23 +445,6 @@ def addForces( includedAtoms = [includedAtoms[i] for i in atoms] atomicNumbers = [atom.element.atomic_number for atom in includedAtoms] - if returnEnergyType == "energy": - try: - e0_max = float(model.atomic_energies_fn.atomic_energies.detach().abs().max()) - except AttributeError: - e0_max = 0.0 - if e0_max > 100.0: # 1 eV per atom is conservative; foundation models far exceed this - import warnings as _w - _w.warn( - f"returnEnergyType='energy' includes per-atom reference " - f"energies (max |e0| = {e0_max:.2f} eV/atom over {int(model.atomic_energies_fn.atomic_energies.numel())} " - f"element entries). Use precision='double' if you need accurate absolute " - f"energies, or pass returnEnergyType='interaction_energy' for the " - f"e0-subtracted readout (note: only 'energy' is gradient-consistent " - f"with PolarMACE — see the docstring).", - stacklevel=2, - ) - modelDefaultDtype = next(model.parameters()).dtype if precision is None: dtype = modelDefaultDtype From 60808ac5d70e4ca0ca7f7f22ad25073c3657b4d0 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 19:00:05 +0000 Subject: [PATCH 10/20] Keep MACE-specific embedding logic electrostatic --- openmmml/models/macepotential.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index de38251..5b4a4ab 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -707,23 +707,13 @@ def _supports_mm_embedding(model) -> bool: return bool(getattr(model, "supports_external_electrostatics", False)) -_SUPPORTED_EMBEDDINGS = ("mechanical", "electrostatic") -_MM_EMBEDDING_MODES = ("electrostatic",) - - def _should_use_mm_embedding(model, atoms: Optional[Iterable[int]], embedding: str) -> bool: - if embedding not in _SUPPORTED_EMBEDDINGS: - raise ValueError( - f"Unsupported embedding mode '{embedding}'. Supported values are " - + ", ".join(repr(m) for m in _SUPPORTED_EMBEDDINGS) - + "." - ) - if embedding not in _MM_EMBEDDING_MODES: + """Validate the only MACE-specific embedding: electrostatic.""" + if embedding == "mechanical": return False + if embedding != "electrostatic": + raise ValueError(f"Unsupported embedding mode '{embedding}'.") if not _supports_mm_embedding(model): - # The mixed system has had its ML-MM Coulomb removed on the assumption - # that the model will supply it, so falling back to mechanical - # embedding here would silently discard those interactions. raise ValueError( f"embedding='{embedding}' requires a model that accepts MM charges " f"and positions (PolarMACE); got {model.__class__.__name__}." From a8b89e1fadd3f55ca7f6f6b9bf6c9c9047f9584d Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 19:03:55 +0000 Subject: [PATCH 11/20] Simplify MACE embedding documentation --- openmmml/models/macepotential.py | 127 +++---------------------------- 1 file changed, 11 insertions(+), 116 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 5b4a4ab..b782e25 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -266,30 +266,14 @@ class MACEPotentialImpl(MLPotentialImpl): >>> system = potential.createSystem(topology, precision='single') By default, the implementation uses the precision of the loaded MACE model. - According to the MACE documentation, 'single' precision is recommended for MD (faster but - less accurate), while 'double' precision is recommended for geometry optimization. + Single precision is faster; double precision is more accurate. By default the reported energy is ``interaction_energy``. PolarMACE is an exception: it automatically uses the full ``energy`` output because that is the scalar whose gradient w.r.t. positions is reported as the force. - >>> system = potential.createSystem(topology, returnEnergyType='interaction_energy') - - PolarMACE therefore uses ``energy`` automatically even when - ``returnEnergyType`` is left at its default. - - Precision caveat for ``returnEnergyType='energy'``: this key returns the - full ``total_energy = e0 + inter_e + extras`` where ``e0`` are the - model's per-atom reference energies. For foundation models (mace-mp, - mace-off, mace-omat, ...) ``e0`` is typically tens of eV per atom, so the - reported scalar for a large ML region can be 10⁴–10⁶ eV in magnitude. - The **forces** stay exact at any scale (they are gradients of this same - scalar), but the **energy** column written into single-precision OpenMM - state files / log lines may carry only ~6–7 significant digits at that - magnitude — meV resolution is lost. Use ``precision='double'`` if you - need accurate absolute energies, or note that energy differences (e.g. - NVE drift) still resolve cleanly because the e0 contribution cancels in - the difference. + PolarMACE automatically uses ``energy`` so its electrostatic forces and + reported energy remain consistent. Attributes ---------- @@ -299,12 +283,7 @@ class MACEPotentialImpl(MLPotentialImpl): The path to the locally trained MACE model if ``name`` is 'mace'. """ - # (Function name, model name, restrictive license name or None, long-range, - # accepts MM charges) - # - # The last flag records whether the model can be given the charges and - # positions of the atoms outside the ML subset, which is what electrostatic - # embedding requires. Only the PolarMACE family can. + # (loader, model name, restrictive license, long-range, accepts MM charges) KNOWN_MODELS = { 'mace-off23-small': ('mace_off', 'small', 'ASL', False, False), 'mace-off23-medium': ('mace_off', 'medium', 'ASL', False, False), @@ -534,65 +513,18 @@ def createMixedSystem(self, **args) -> openmm.System: """Create a mixed system using electrostatic embedding. - The model, rather than the conventional force field, is responsible for - the electrostatic interactions between the atoms within the ML subset - and those outside of it: it is passed the positions and conventional - force field charges of the atoms outside the ML subset, and returns - forces on them alongside the forces on the ML subset. The ML subset can - therefore polarize in response to its surroundings, which mechanical - embedding does not allow. - - This is implemented as the "global charge zero" variant: the - conventional force field charge of every atom in the ML subset is set to - zero. Every Coulomb term involving an ML atom is then zero by - construction, including the reciprocal space part of PME, without any - per-pair exceptions being added. Adding an exception for each ML-MM - pair would instead be incorrect under periodic boundary conditions, - since NonbondedForce evaluates exceptions using plain Cartesian - distances rather than the minimum image convention, so the ML-MM - Lennard-Jones interaction would silently vanish for any pair that is - only within the cutoff across a periodic boundary. Lennard-Jones is - left to the conventional force field and continues to use the ordinary, - periodicity-aware pair list. - - Interactions within the ML subset are excluded entirely, as the model - computes them. Bonded terms that cross the ML/MM boundary are retained. - - Only models that accept MM charges and positions, which at present means - the PolarMACE family, can be used with this embedding method. An error - is raised for any other model rather than falling back to mechanical - embedding, since by that point the ML-MM electrostatics have already - been removed from the conventional force field and a fallback would - simply lose them. - - Because this method has to account for every Coulomb term in the force - field, it requires the System to contain exactly one NonbondedForce: the - MM charges given to the model are read from one, so several would be - ambiguous. - - It also needs to be told about any CustomNonbondedForce, whose energy - expression is arbitrary and cannot be inspected here. Pass - customNonbondedHasCharges=False to declare that it holds no - electrostatics, or True together with customNonbondedChargeParameter - naming the per-particle parameter that holds the charge, which is then - zeroed on the ML atoms exactly as for the NonbondedForce. An error is - raised if the answer is needed and has not been given. - - Note that zeroing that parameter removes the ML terms only if the - expression is multiplicatively separable in the charge, as the usual - q1*q2/r is. That cannot be verified here, so it is the caller's - responsibility. + PolarMACE receives MM positions and charges and computes all ML/MM + electrostatics. The conventional ML charges and ML-region bonded terms + are removed; Lennard-Jones and MM/MM terms remain in the force field. + Exactly one NonbondedForce is required. CustomNonbondedForce charge + handling must be declared with ``customNonbondedHasCharges`` and, when + needed, ``customNonbondedChargeParameter``. """ if embedding != "electrostatic": raise ValueError(f"Unsupported embedding type: {embedding}") - # Check that the model can actually accept MM charges and positions - # before touching the System, so that an unsuitable model is rejected - # with the System left alone rather than stripped of its ML-MM - # electrostatics. This is also the first point at which the check is - # possible, since it needs the loaded checkpoint; the model is handed to - # addForces() below so the checkpoint is only read once. + # Validate before modifying the input system. model, device = self._loadModel(args) if not _supports_mm_embedding(model): @@ -602,22 +534,10 @@ def createMixedSystem(self, ) if interpolate: - # At lambda_interpolate=0 the conventional endpoint would be missing - # the ML-MM Coulomb energy, which is removed from the conventional - # force field outside of the interpolating CustomCVForce and cannot - # be restored from within it. raise ValueError("Electrostatic embedding does not support interpolation.") periodic = system.usesPeriodicBoundaryConditions() - # Electrostatic embedding has to account for every Coulomb term in the - # force field: the ones involving the ML subset are removed here on the - # understanding that the model supplies them. Anything it cannot see is - # either left in place and counted twice, or removed and never replaced, - # and in both cases the result is a wrong energy rather than an error. - # So refuse the cases where the electrostatics cannot be located rather - # than guessing. - nonbondedForces = [f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)] if len(nonbondedForces) > 1: # The MM charges handed to the model are read from a single @@ -632,36 +552,21 @@ def createMixedSystem(self, if customNonbondedHasCharges and customNonbondedChargeParameter is None: raise ValueError("A CustomNonbondedForce includes electrostatic interactions, so customNonbondedChargeParameter must name the per-particle parameter holding the charge.") - # Create the new system with the ML-ML interactions that the model - # computes removed. - newSystem = utilities.removeBonds(system, topology, atoms, True) atomSet = set(atoms) for force in newSystem.getForces(): if isinstance(force, openmm.NonbondedForce): - # Zero the charge of every ML atom, which removes all Coulomb - # interactions involving the ML subset while leaving its - # Lennard-Jones parameters, and the MM-MM interactions, - # untouched. - for atom in atoms: charge, sigma, epsilon = force.getParticleParameters(atom) force.setParticleParameters(atom, 0.0, sigma, epsilon) - # setParticleParameters() does not update the charge products - # that were precomputed for existing exceptions, so the 1-4 - # Coulomb terms crossing the ML/MM boundary have to be zeroed - # separately. - for index in range(force.getNumExceptions()): p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(index) if p1 in atomSet or p2 in atomSet: force.setExceptionParameters(index, p1, p2, 0.0, sigma, epsilon) - # Exclude the ML-ML interactions entirely. - for i in range(len(atoms)): for j in range(i): force.addException(atoms[i], atoms[j], 0, 1, 0, True) @@ -670,12 +575,6 @@ def createMixedSystem(self, elif isinstance(force, openmm.CustomNonbondedForce): - # Zero the named charge parameter on the ML atoms, the same - # trick used for the NonbondedForce above. Unlike there it is - # not guaranteed to work: it removes the ML terms only if the - # energy expression is multiplicatively separable in the charge, - # as q1*q2/r is, and that cannot be checked here. - if customNonbondedChargeParameter is not None: names = [force.getPerParticleParameterName(i) for i in range(force.getNumPerParticleParameters())] @@ -689,10 +588,6 @@ def createMixedSystem(self, utilities.addCustomNonbondedExclusions(force, atoms) - # Add the ML potential, telling it that it is responsible for the - # electrostatic interactions with the atoms outside the ML subset, and - # handing over the model already loaded above. - self._preloadedModel = (model, device) try: self.addForces(topology, newSystem, atoms, forceGroup, embedding=embedding, From e46dd2af03560e314f6aaec6241ac589120e0042 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 19:09:44 +0000 Subject: [PATCH 12/20] Simplify external source preparation --- openmmml/models/macepotential.py | 47 ++++++++++---------------------- 1 file changed, 15 insertions(+), 32 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index b782e25..c9ca9b9 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -42,28 +42,19 @@ def _prepare_external_sources(model, data, compute_force: bool): positions = data.get("mm_positions") charges = data.get("mm_charges") - if ( - positions is None - or charges is None - or positions.numel() == 0 - or charges.numel() == 0 - ): + if positions is None or charges is None: return None ml_positions = data["positions"] - positions = positions.to(device=ml_positions.device, dtype=ml_positions.dtype) - positions = positions.clone().requires_grad_(compute_force) + positions = positions.to(ml_positions).clone().requires_grad_(compute_force) + charges = charges.to(ml_positions).reshape(-1) + if positions.shape[0] != charges.shape[0]: + raise ValueError("MM positions and charges must have the same length.") + width = (int(model.atomic_multipoles_max_l) + 1) ** 2 - features = torch.zeros( - (charges.numel(), width), - dtype=ml_positions.dtype, - device=ml_positions.device, - ) - features[:, 0] = charges.to(features).reshape(-1) - if positions.shape[0] != features.shape[0]: - raise ValueError( - "MM positions and electrostatic sources must have the same length." - ) + features = torch.zeros((len(charges), width), device=ml_positions.device, + dtype=ml_positions.dtype) + features[:, 0] = charges transform = getattr(model, "_charges_to_mul_ir", None) if transform is not None: @@ -71,19 +62,15 @@ def _prepare_external_sources(model, data, compute_force: bool): batch = data.get("mm_source_batch") if batch is None: - if int(data["pbc"].reshape(-1, 3).shape[0]) != 1: + if data["pbc"].reshape(-1, 3).shape[0] != 1: raise ValueError( "mm_source_batch is required for batched PolarMACE inputs." ) - batch = torch.zeros( - positions.shape[0], dtype=torch.long, device=positions.device - ) + batch = torch.zeros(len(positions), dtype=torch.long, device=positions.device) else: batch = batch.to(device=positions.device, dtype=torch.long).reshape(-1) - if batch.shape[0] != positions.shape[0]: - raise ValueError( - "mm_source_batch and mm_positions must have the same length." - ) + if batch.shape[0] != positions.shape[0]: + raise ValueError("mm_source_batch and mm_positions must have the same length.") return {"positions": positions, "features": features, "batch": batch} @@ -734,17 +721,13 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch inputDict["mm_source_batch"] = torch.zeros( len(mmInfo["mm_atoms"]), dtype=torch.long, device=ptr.device ) + # eval and get results results = model(inputDict, compute_force=True) energy = float(results[returnEnergyType].detach())*energyScale forces = (results["forces"]*energyScale*lengthScale).detach().cpu().numpy() mm_forces = results.get("mm_forces") + if mmInfo is not None and mm_forces is None: - # The mixed system has had its ML-MM electrostatics removed on the - # understanding that this model supplies them. A model that returns no - # forces on the MM atoms did not compute them, so continuing would leave - # those interactions missing entirely rather than merely approximated, - # and nothing downstream would report it. The usual cause is a - # PolarMACE checkpoint whose forward does not accept mm_charges. raise ValueError("The model returned no 'mm_forces' although MM charges were supplied; it does not implement electrostatic embedding.") if mm_forces is not None: mm_forces = (mm_forces * energyScale * lengthScale).detach().cpu().numpy() From f2e33259bb00a128a04a31d59500688e930b3af7 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 19:11:06 +0000 Subject: [PATCH 13/20] Simplify MACE adapter setup --- openmmml/models/macepotential.py | 38 +++++--------------------------- 1 file changed, 6 insertions(+), 32 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index c9ca9b9..77ad896 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -388,28 +388,21 @@ def addForces( import torch try: from mace.tools import utils, to_one_hot, atomic_numbers_to_indices - from mace.calculators.foundations_models import mace_off, mace_mp, mace_omol, mace_polar except ImportError as e: raise ImportError(f"Failed to import mace with error: {e}. Install mace with 'pip install mace-torch'.") - try: - from e3nn.util import jit - except ImportError as e: - raise ImportError(f"Failed to import e3nn with error: {e}. Install e3nn with 'pip install e3nn'.") assert returnEnergyType in ["interaction_energy", "energy"], f"Unsupported returnEnergyType: '{returnEnergyType}'. Supported options are 'interaction_energy' or 'energy'." - # Load the model. - model, device = self._loadModel(args) if model.__class__.__name__ == "PolarMACE": returnEnergyType = "energy" use_mm_embedding = _should_use_mm_embedding(model, atoms, embedding) - includedAtoms = list(topology.atoms()) + included_atoms = list(topology.atoms()) if atoms is not None: - includedAtoms = [includedAtoms[i] for i in atoms] - atomicNumbers = [atom.element.atomic_number for atom in includedAtoms] + included_atoms = [included_atoms[i] for i in atoms] + atomic_numbers = [atom.element.atomic_number for atom in included_atoms] modelDefaultDtype = next(model.parameters()).dtype if precision is None: @@ -422,10 +415,6 @@ def addForces( raise ValueError(f"Unsupported precision {precision} for the model. Supported values are 'single' and 'double'.") if dtype != modelDefaultDtype: print(f"Model dtype is {modelDefaultDtype} and requested dtype is {dtype}. The model will be converted to the requested dtype.") - # Actually do the conversion. The previous code only printed the - # warning and left the model untouched, which caused dtype - # mismatches inside e3nn's compiled TensorProduct submodules - # when inputs were passed at the requested dtype. model = model.to(dtype) model_device = device @@ -436,18 +425,12 @@ def addForces( zTable = utils.AtomicNumberTable([int(z) for z in model.atomic_numbers]) nodeAttrs = to_one_hot( - torch.tensor(atomic_numbers_to_indices(atomicNumbers, z_table=zTable), dtype=torch.long, device=model_device).unsqueeze(-1), + torch.tensor(atomic_numbers_to_indices(atomic_numbers, z_table=zTable), dtype=torch.long, device=model_device).unsqueeze(-1), num_classes=len(zTable)) - if atoms is None: - indices = None - else: - indices = np.array(atoms) + indices = None if atoms is None else np.array(atoms) mmInfo = None if use_mm_embedding: - # ML-MM Coulomb is removed by MLPotential.createMixedSystem when - # embedding='electrostatic'. We only need MM positions/charges for - # the PolarMACE input here. mmInfo = _prepareMMEmbedding(system, atoms, customNonbondedChargeParameter) periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions() @@ -475,17 +458,8 @@ def getMLLongRange(self) -> bool | None: return None def getSupportedEmbeddings(self) -> list[str]: - - # Electrostatic embedding requires a model that accepts the charges and - # positions of the atoms outside the ML subset, which of the pretrained - # models only the PolarMACE family does. A custom checkpoint may be a - # PolarMACE model too, but that cannot be known without loading it, so - # the method is offered and createMixedSystem() rejects the checkpoint - # once loaded if it turns out not to be one. - if self.name in MACEPotentialImpl.KNOWN_MODELS: - _, _, _, _, acceptsMMCharges = MACEPotentialImpl.KNOWN_MODELS[self.name] - return ["electrostatic"] if acceptsMMCharges else [] + return ["electrostatic"] if self.KNOWN_MODELS[self.name][4] else [] return ["electrostatic"] def createMixedSystem(self, From e7a14e7692513c02974777e6d792f029df17f7f3 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 10 Sep 2026 21:13:40 +0000 Subject: [PATCH 14/20] Fix PolarMACE energy selection and reject charge offsets --- openmmml/models/macepotential.py | 10 +++++++++- test/TestMACEPotential.py | 17 +++++++++++++++++ test/TestPolarMACEElectrostaticEmbedding.py | 18 ++++++++++++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 77ad896..93b6aa8 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -394,7 +394,7 @@ def addForces( assert returnEnergyType in ["interaction_energy", "energy"], f"Unsupported returnEnergyType: '{returnEnergyType}'. Supported options are 'interaction_energy' or 'energy'." model, device = self._loadModel(args) - if model.__class__.__name__ == "PolarMACE": + if model.__class__.__name__ in ("PolarMACE", "PolarMACEExternalSources"): returnEnergyType = "energy" use_mm_embedding = _should_use_mm_embedding(model, atoms, embedding) @@ -505,6 +505,14 @@ def createMixedSystem(self, # NonbondedForce, so several of them are ambiguous. raise ValueError("Multiple NonbondedForce objects encountered; electrostatic embedding requires exactly one.") + for force in nonbondedForces: + for index in range(force.getNumParticleParameterOffsets()): + if force.getParticleParameterOffset(index)[2] != 0: + raise ValueError("Electrostatic embedding does not support charge parameter offsets.") + for index in range(force.getNumExceptionParameterOffsets()): + if force.getExceptionParameterOffset(index)[2] != 0: + raise ValueError("Electrostatic embedding does not support charge parameter offsets.") + if any(isinstance(f, openmm.CustomNonbondedForce) for f in system.getForces()): # A CustomNonbondedForce's energy expression is arbitrary, so # whether it contains electrostatics cannot be determined here. diff --git a/test/TestMACEPotential.py b/test/TestMACEPotential.py index af46d17..701f62c 100644 --- a/test/TestMACEPotential.py +++ b/test/TestMACEPotential.py @@ -50,6 +50,23 @@ def testPrepareMMEmbedding(): np.testing.assert_allclose(info["mm_charges"], [0.25], atol=1e-12) +@pytest.mark.parametrize("kind", ["ml_particle", "mm_particle", "exception"]) +def testElectrostaticRejectsChargeOffsets(monkeypatch, kind): + system, force = _simple_nonbonded_system() + force.addGlobalParameter("lambda", 1.0) + if kind == "exception": + index = force.addException(0, 2, 0.25, 0.3, 0.0) + force.addExceptionParameterOffset("lambda", index, 0.5, 0.0, 0.0) + else: + force.addParticleParameterOffset("lambda", 0 if kind == "ml_particle" else 2, 0.5, 0.0, 0.0) + before = mm.XmlSerializer.serialize(system) + impl = MACEPotentialImpl("mace", None) + monkeypatch.setattr(impl, "_loadModel", lambda args: (PolarMACE(), "cpu")) + with pytest.raises(ValueError, match="charge parameter offsets"): + impl.createMixedSystem(app.Topology(), system, [0], 0, False, "electrostatic") + assert mm.XmlSerializer.serialize(system) == before + + class _FakeState: def __init__(self, positions_angstrom): self._positions = np.asarray(positions_angstrom, dtype=np.float64) * unit.angstrom diff --git a/test/TestPolarMACEElectrostaticEmbedding.py b/test/TestPolarMACEElectrostaticEmbedding.py index 67ddfc6..ff00400 100644 --- a/test/TestPolarMACEElectrostaticEmbedding.py +++ b/test/TestPolarMACEElectrostaticEmbedding.py @@ -272,6 +272,24 @@ def test_polar_mace_pbc_runs_and_is_finite(polar_mace_model_path): assert np.all(np.isfinite(f)) +@pytest.mark.parametrize("periodic", [False, True]) +def test_polar_mace_default_returns_total_energy(polar_mace_model_path, periodic): + topology, system = _build_topology_and_system(num_mm_waters=2, periodic=periodic) + results = [] + for options in ({}, {"returnEnergyType": "energy"}): + potential = MLPotential("mace", modelPath=polar_mace_model_path) + mixed = potential.createMixedSystem( + topology, system, [0, 1, 2], embedding="electrostatic", **options + ) + context = openmm.Context(mixed, openmm.VerletIntegrator(0.001), + openmm.Platform.getPlatformByName("Reference")) + context.setPositions(_initial_positions(num_mm_waters=2)) + results.append(_energy_and_forces(context)) + del context + np.testing.assert_allclose(results[0][0], results[1][0], rtol=0, atol=1e-9) + np.testing.assert_allclose(results[0][1], results[1][1], rtol=0, atol=1e-8) + + def test_polar_mace_pbc_translation_invariance(polar_mace_model_path): """Energy and forces must be invariant under whole-system translation by one full box vector — the fundamental PBC sanity check.""" From dad64e5a9e25a04e37629b6e9637a7e282c3876b Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Wed, 16 Sep 2026 23:37:11 +0000 Subject: [PATCH 15/20] Pin graph_longrange to the external_field branch The MACE backend requirements still pinned graph_electrostatics v0.4.0, which predates the external-source work: installing from it yields a graph_longrange without external_source_energy / external_source_features, so PolarMACE electrostatic embedding fails at import. Pin the external_field branch instead, and name it (with the install command) in the ImportError raised when the blocks are missing. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN --- devtools/requirements/mace-aimnet-torchmd.txt | 2 +- openmmml/models/macepotential.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/devtools/requirements/mace-aimnet-torchmd.txt b/devtools/requirements/mace-aimnet-torchmd.txt index 0bcef3b..1150f37 100644 --- a/devtools/requirements/mace-aimnet-torchmd.txt +++ b/devtools/requirements/mace-aimnet-torchmd.txt @@ -6,4 +6,4 @@ ase huggingface-hub torchani git+https://github.com/ChengUCB/les.git@v0.2.0 -git+https://github.com/WillBaldwin0/graph_electrostatics.git@v0.4.0 +git+https://github.com/WillBaldwin0/graph_electrostatics.git@external_field diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 93b6aa8..4e81e6c 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -95,8 +95,11 @@ def _enable_polarmace_external_sources(model): ) except ImportError as exc: raise ImportError( - "PolarMACE electrostatic embedding requires a graph_longrange " - "release that provides the external-source energy and feature blocks." + "PolarMACE electrostatic embedding requires the external-source energy " + "and feature blocks from graph_longrange. Install the external_field " + "branch:\n" + " pip install 'git+https://github.com/WillBaldwin0/" + "graph_electrostatics.git@external_field'" ) from exc model.electric_potential_descriptor = ( From fe1fb1abdb7c92fb9bb2a2f4b74f633f52511dd6 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 17 Sep 2026 02:24:51 +0000 Subject: [PATCH 16/20] Tidy the electrostatic embedding validation gates No behaviour change. Verified by fingerprinting the mixed System built before and after (every NonbondedForce parameter and exception, CustomNonbondedForce exclusions and parameters, plus energy and forces from a Reference Context at fixed positions, PBC and non-PBC, sharing one saved PolarMACE model): the two fingerprints are byte-identical, energies -84.29790229105993 and -56.26526271299811 kJ/mol in both. - The `embedding != "electrostatic"` guard is unreachable: MLPotential routes an embedding to this method only when it appears in getSupportedEmbeddings(), and MACE never advertises "mechanical", so that case already falls back to the generic embedding plugin. Kept as an explicit internal invariant and commented as such rather than left looking like user-facing validation. addForces() still handles embedding="mechanical", because the generic plugin calls it. - Model support was validated twice, in createMixedSystem and again in addForces() via _should_use_mm_embedding(). The helper is now the single owner of that message and is called early, so an unsupported model is rejected before any force-field surgery. - Dropped dead work in the NonbondedForce block: the loop zeroing chargeProd for every exception touching an ML atom was followed by an unconditional addException(..., replace=True) over all ML-ML pairs, which overwrote half of it. It now handles ML-MM exceptions only. - New _customNonbondedChargeIndex() resolves the named charge parameter and raises if absent, used both by the early validation and by the force loop, so the message has one source and no force indices are carried across systems. - Copied upstream's comment onto setExceptionsUsePeriodicBoundaryConditions() and ordered the gates cheap-first: interpolation, model, force-field structure, then surgery. Precision parsing moved ahead of model loading so an invalid value fails without loading the model. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN --- openmmml/models/macepotential.py | 63 ++++++++++++++++++++------------ 1 file changed, 39 insertions(+), 24 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index d7f8a1c..a1a7abe 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -396,6 +396,15 @@ def addForces( assert returnEnergyType in ["interaction_energy", "energy"], f"Unsupported returnEnergyType: '{returnEnergyType}'. Supported options are 'interaction_energy' or 'energy'." + if precision is None: + dtype = None + elif precision == "single": + dtype = torch.float32 + elif precision == "double": + dtype = torch.float64 + else: + raise ValueError(f"Unsupported precision {precision} for the model. Supported values are 'single' and 'double'.") + model, device = self._loadModel(args) if model.__class__.__name__ in ("PolarMACE", "PolarMACEExternalSources"): returnEnergyType = "energy" @@ -410,12 +419,6 @@ def addForces( modelDefaultDtype = next(model.parameters()).dtype if precision is None: dtype = modelDefaultDtype - elif precision == "single": - dtype = torch.float32 - elif precision == "double": - dtype = torch.float64 - else: - raise ValueError(f"Unsupported precision {precision} for the model. Supported values are 'single' and 'double'.") if dtype != modelDefaultDtype: print(f"Model dtype is {modelDefaultDtype} and requested dtype is {dtype}. The model will be converted to the requested dtype.") model = model.to(dtype) @@ -491,29 +494,26 @@ def createMixedSystem(self, needed, ``customNonbondedChargeParameter``. """ + # Internal invariant: the framework only calls this method for names + # returned by getSupportedEmbeddings(); this is not user-facing validation. if embedding != "electrostatic": raise ValueError(f"Unsupported embedding type: {embedding}") - # Validate before modifying the input system. - - model, device = self._loadModel(args) - if not _supports_mm_embedding(model): - raise ValueError( - f"embedding='{embedding}' requires a model that accepts MM charges " - f"and positions (PolarMACE); got {model.__class__.__name__}." - ) - + # Interpolation of the model's MM electrostatics is not implemented. if interpolate: raise ValueError("Electrostatic embedding does not support interpolation.") - periodic = system.usesPeriodicBoundaryConditions() + # Validate model support before modifying the input system. + model, device = self._loadModel(args) + _should_use_mm_embedding(model, atoms, embedding) + # Validate the force-field setup before creating the modified system. + # Multiple charge sources would make the MM charges ambiguous. nonbondedForces = [f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)] if len(nonbondedForces) > 1: - # The MM charges handed to the model are read from a single - # NonbondedForce, so several of them are ambiguous. raise ValueError("Multiple NonbondedForce objects encountered; electrostatic embedding requires exactly one.") + # The callback uses fixed charges and cannot follow charge offsets. for force in nonbondedForces: for index in range(force.getNumParticleParameterOffsets()): if force.getParticleParameterOffset(index)[2] != 0: @@ -530,6 +530,14 @@ def createMixedSystem(self, if customNonbondedHasCharges and customNonbondedChargeParameter is None: raise ValueError("A CustomNonbondedForce includes electrostatic interactions, so customNonbondedChargeParameter must name the per-particle parameter holding the charge.") + # A named charge parameter must exist in every custom force we modify. Checked here, before + # the system is copied, so an unusable name fails before any surgery. + if customNonbondedChargeParameter is not None: + for force in system.getForces(): + if isinstance(force, openmm.CustomNonbondedForce): + _customNonbondedChargeIndex(force, customNonbondedChargeParameter) + + periodic = system.usesPeriodicBoundaryConditions() newSystem = utilities.removeBonds(system, topology, atoms, True) atomSet = set(atoms) @@ -542,23 +550,22 @@ def createMixedSystem(self, for index in range(force.getNumExceptions()): p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(index) - if p1 in atomSet or p2 in atomSet: + # ML-ML exceptions are replaced by the all-pairs loop below. + if (p1 in atomSet) != (p2 in atomSet): force.setExceptionParameters(index, p1, p2, 0.0, sigma, epsilon) for i in range(len(atoms)): for j in range(i): force.addException(atoms[i], atoms[j], 0, 1, 0, True) + # This may cause exceptions in the MM region to use PBCs, but + # this should not ordinarily have any significant effects. force.setExceptionsUsePeriodicBoundaryConditions(periodic) elif isinstance(force, openmm.CustomNonbondedForce): if customNonbondedChargeParameter is not None: - names = [force.getPerParticleParameterName(i) - for i in range(force.getNumPerParticleParameters())] - if customNonbondedChargeParameter not in names: - raise ValueError(f"A CustomNonbondedForce has no per-particle parameter '{customNonbondedChargeParameter}'; it defines {names}.") - chargeIndex = names.index(customNonbondedChargeParameter) + chargeIndex = _customNonbondedChargeIndex(force, customNonbondedChargeParameter) for atom in atoms: parameters = list(force.getParticleParameters(atom)) parameters[chargeIndex] = 0.0 @@ -576,6 +583,14 @@ def createMixedSystem(self, return newSystem +def _customNonbondedChargeIndex(force: openmm.CustomNonbondedForce, name: str) -> int: + """Index of the per-particle parameter holding the charge, or raise if the force has no such parameter.""" + names = [force.getPerParticleParameterName(i) for i in range(force.getNumPerParticleParameters())] + if name not in names: + raise ValueError(f"A CustomNonbondedForce has no per-particle parameter '{name}'; it defines {names}.") + return names.index(name) + + def _supports_mm_embedding(model) -> bool: return bool(getattr(model, "supports_external_electrostatics", False)) From 00392e58916514c6d34df75a8a83d6ab08b5bb31 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 17 Sep 2026 02:38:40 +0000 Subject: [PATCH 17/20] Leave the precision logic alone and restore upstream's precision note - Revert the precision-parsing reorder from the previous commit: the block is byte-identical to how it was before the cleanup, parsed after the model is loaded as upstream does. - Restore the sentence this branch had dropped, verbatim from upstream: "According to the MACE documentation, 'single' precision is recommended for MD (faster but less accurate), while 'double' precision is recommended for geometry optimization." It was replaced by a shorter paraphrase in a8b89e1; the diff should not silently reword upstream documentation. Fingerprint of the mixed System (all NonbondedForce parameters and exceptions, CustomNonbondedForce exclusions and parameters, plus energy and forces from a Reference Context at fixed positions, PBC and non-PBC) is still byte-identical to the pre-cleanup baseline. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN --- openmmml/models/macepotential.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index a1a7abe..4db92ef 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -256,7 +256,8 @@ class MACEPotentialImpl(MLPotentialImpl): >>> system = potential.createSystem(topology, precision='single') By default, the implementation uses the precision of the loaded MACE model. - Single precision is faster; double precision is more accurate. + According to the MACE documentation, 'single' precision is recommended for MD (faster but + less accurate), while 'double' precision is recommended for geometry optimization. By default the reported energy is ``interaction_energy``. PolarMACE is an exception: it automatically uses the full ``energy`` output because that is @@ -396,15 +397,6 @@ def addForces( assert returnEnergyType in ["interaction_energy", "energy"], f"Unsupported returnEnergyType: '{returnEnergyType}'. Supported options are 'interaction_energy' or 'energy'." - if precision is None: - dtype = None - elif precision == "single": - dtype = torch.float32 - elif precision == "double": - dtype = torch.float64 - else: - raise ValueError(f"Unsupported precision {precision} for the model. Supported values are 'single' and 'double'.") - model, device = self._loadModel(args) if model.__class__.__name__ in ("PolarMACE", "PolarMACEExternalSources"): returnEnergyType = "energy" @@ -419,6 +411,12 @@ def addForces( modelDefaultDtype = next(model.parameters()).dtype if precision is None: dtype = modelDefaultDtype + elif precision == "single": + dtype = torch.float32 + elif precision == "double": + dtype = torch.float64 + else: + raise ValueError(f"Unsupported precision {precision} for the model. Supported values are 'single' and 'double'.") if dtype != modelDefaultDtype: print(f"Model dtype is {modelDefaultDtype} and requested dtype is {dtype}. The model will be converted to the requested dtype.") model = model.to(dtype) From 14a3abd5ebb25508d0f10f9a29e0571618bd64f3 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 17 Sep 2026 09:36:40 +0000 Subject: [PATCH 18/20] Follow the EMLE embedding's style: camelCase, explicit callback arguments Style alignment with openmmml/embeddings/emleembedding.py (upstream PR #159), which is the reference for how an ML/MM embedding is written here. No behaviour change: the mixed System fingerprint (every NonbondedForce parameter and exception, CustomNonbondedForce exclusions and parameters, plus energy and forces from a Reference Context at fixed positions, PBC and non-PBC, against one saved PolarMACE model) is byte-identical to before, energies -84.29790229105993 and -56.26526271299811 kJ/mol. Tests: 104 passed, 6 skipped. - camelCase for identifiers we own: modelDevice, positionsFull, includedAtoms, useMMEmbedding, mlAtoms/mmAtoms/mmCharges/mmPositions/mmForces, and the private helpers _supportsMMEmbedding, _prepareExternalSources, _enablePolarMACEExternalSources. External API names are untouched and must stay snake_case: the MACE input/result dict keys ("mm_positions", "mm_charges", "mm_source_batch", "external_field", "node_attrs", "fermi_level", "mm_forces"), the compute_* and requires_grad kwargs, and the mace_off/mace_mp/mace_polar model-family strings. - The mmInfo dict no longer crosses into the callback. _prepareMMEmbedding is unpacked at the call site and the values are bound by name, as emleembedding.py does with mlIndices/mmIndices/ mmCharges, so _computeMACE's signature documents its own inputs: def _computeMACE(state, model, ptr, nodeAttrs, batch, pbc, returnEnergyType, charge, multiplicity, periodic, mlIndices=None, mmIndices=None, mmCharges=None) The electrostatic path is selected by `mmIndices is not None` exactly where it previously used `mmInfo is not None`, and still does not call setParticles(). - _shouldUseMMEmbedding both raised and returned a flag, and one call site discarded the flag, so it read as a no-op. It is now _validateMMEmbedding, which only raises; the flag is derived where it is used (`useMMEmbedding = embedding == "electrostatic"`). Its "unsupported mode" message now names what MACE implements and what it delegates. - Wrapped the four error messages this branch added that ran to 122-187 characters, and split the _computeMACE signature over two lines. Upstream's own long lines and the KNOWN_MODELS URL table are left alone. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN --- openmmml/models/macepotential.py | 215 +++++++++++++++++-------------- 1 file changed, 119 insertions(+), 96 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index 4db92ef..cdafc8f 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -37,7 +37,7 @@ import numpy as np -def _prepare_external_sources(model, data, compute_force: bool): +def _prepareExternalSources(model, data, computeForce: bool): import torch positions = data.get("mm_positions") @@ -45,15 +45,15 @@ def _prepare_external_sources(model, data, compute_force: bool): if positions is None or charges is None: return None - ml_positions = data["positions"] - positions = positions.to(ml_positions).clone().requires_grad_(compute_force) - charges = charges.to(ml_positions).reshape(-1) + mlPositions = data["positions"] + positions = positions.to(mlPositions).clone().requires_grad_(computeForce) + charges = charges.to(mlPositions).reshape(-1) if positions.shape[0] != charges.shape[0]: raise ValueError("MM positions and charges must have the same length.") width = (int(model.atomic_multipoles_max_l) + 1) ** 2 - features = torch.zeros((len(charges), width), device=ml_positions.device, - dtype=ml_positions.dtype) + features = torch.zeros((len(charges), width), device=mlPositions.device, + dtype=mlPositions.dtype) features[:, 0] = charges transform = getattr(model, "_charges_to_mul_ir", None) @@ -74,7 +74,7 @@ def _prepare_external_sources(model, data, compute_force: bool): return {"positions": positions, "features": features, "batch": batch} -def _enable_polarmace_external_sources(model): +def _enablePolarMACEExternalSources(model): """Wrap PolarMACE with dynamic MM electrostatic sources in eager mode.""" import torch @@ -140,7 +140,7 @@ def forward( compute_atomic_stresses: bool = False, **kwargs, ): - external = _prepare_external_sources(self.model, data, compute_force) + external = _prepareExternalSources(self.model, data, compute_force) if external is None: return self.model( data, @@ -169,15 +169,15 @@ def forward( "energies and Cartesian forces only." ) - external_kwargs = { + externalKwargs = { "external_feats": external["features"], "external_positions": external["positions"], "external_batch": external["batch"], } self.model.electric_potential_descriptor.set_external_sources( - **external_kwargs + **externalKwargs ) - self.model.coulomb_energy.set_external_sources(**external_kwargs) + self.model.coulomb_energy.set_external_sources(**externalKwargs) try: result = self.model( data, @@ -192,7 +192,7 @@ def forward( **kwargs, ) if compute_force: - ml_gradient, mm_gradient = torch.autograd.grad( + mlGradient, mmGradient = torch.autograd.grad( outputs=[result["energy"]], inputs=[data["positions"], external["positions"]], grad_outputs=[torch.ones_like(result["energy"])], @@ -202,13 +202,13 @@ def forward( ) result["forces"] = ( torch.zeros_like(data["positions"]) - if ml_gradient is None - else -ml_gradient + if mlGradient is None + else -mlGradient ) result["mm_forces"] = ( torch.zeros_like(external["positions"]) - if mm_gradient is None - else -mm_gradient + if mmGradient is None + else -mmGradient ) else: result["mm_forces"] = None @@ -329,11 +329,11 @@ def _loadModel(self, args): 'mace_omol': mace_omol, 'mace_polar': mace_polar, } - loader_name, model_name, restrictive_license, _, _ = self.KNOWN_MODELS[self.name] - model = loaders[loader_name](model=model_name, device=device, return_raw_model=True).to(device) - if restrictive_license is not None: + loaderName, modelName, restrictiveLicense, _, _ = self.KNOWN_MODELS[self.name] + model = loaders[loaderName](model=modelName, device=device, return_raw_model=True).to(device) + if restrictiveLicense is not None: import logging - logging.warning(f'The model {self.name} is distributed under the restrictive {restrictive_license} license. Commercial use is not permitted.') + logging.warning(f'The model {self.name} is distributed under the restrictive {restrictiveLicense} license. Commercial use is not permitted.') elif self.name == "mace": if self.modelPath is None: raise ValueError("No modelPath provided for local MACE model.") @@ -343,7 +343,7 @@ def _loadModel(self, args): else: raise ValueError(f"Unsupported MACE model: {self.name}") if model.__class__.__name__ == "PolarMACE": - model = _enable_polarmace_external_sources(model) + model = _enablePolarMACEExternalSources(model) return model, device def addForces( @@ -401,12 +401,13 @@ def addForces( if model.__class__.__name__ in ("PolarMACE", "PolarMACEExternalSources"): returnEnergyType = "energy" - use_mm_embedding = _should_use_mm_embedding(model, atoms, embedding) + _validateMMEmbedding(model, atoms, embedding) + useMMEmbedding = embedding == "electrostatic" - included_atoms = list(topology.atoms()) + includedAtoms = list(topology.atoms()) if atoms is not None: - included_atoms = [included_atoms[i] for i in atoms] - atomic_numbers = [atom.element.atomic_number for atom in included_atoms] + includedAtoms = [includedAtoms[i] for i in atoms] + atomicNumbers = [atom.element.atomic_number for atom in includedAtoms] modelDefaultDtype = next(model.parameters()).dtype if precision is None: @@ -421,43 +422,46 @@ def addForces( print(f"Model dtype is {modelDefaultDtype} and requested dtype is {dtype}. The model will be converted to the requested dtype.") model = model.to(dtype) - model_device = device + modelDevice = device try: - model_device = next(model.parameters()).device + modelDevice = next(model.parameters()).device except (AttributeError, StopIteration): pass zTable = utils.AtomicNumberTable([int(z) for z in model.atomic_numbers]) nodeAttrs = to_one_hot( - torch.tensor(atomic_numbers_to_indices(atomic_numbers, z_table=zTable), dtype=torch.long, device=model_device).unsqueeze(-1), + torch.tensor(atomic_numbers_to_indices(atomicNumbers, z_table=zTable), dtype=torch.long, device=modelDevice).unsqueeze(-1), num_classes=len(zTable)) - mmInfo = None - if use_mm_embedding: - mmInfo = _prepareMMEmbedding(system, atoms, customNonbondedChargeParameter) + mmIndices = mmCharges = None + if useMMEmbedding: + embeddingData = _prepareMMEmbedding(system, atoms, customNonbondedChargeParameter) + mmIndices = embeddingData["mm_atoms"] + mmCharges = embeddingData["mm_charges"] # The electrostatic path needs FULL-system positions inside the callback: it reads MM # coordinates and returns MM back-reaction forces. PythonForce.setParticles() would hand # the callback only the ML atoms, so that path keeps the explicit index slice/scatter and # does not call setParticles(). The plain ML path uses upstream's restriction instead. - indices = np.array(atoms) if (atoms is not None and mmInfo is not None) else None + mlIndices = np.array(atoms) if (atoms is not None and mmIndices is not None) else None periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions() compute = partial(_computeMACE, model=model, - ptr=torch.tensor([0, nodeAttrs.shape[0]], dtype=torch.long, device=model_device, requires_grad=False), - node_attrs=nodeAttrs.to(dtype), - batch=torch.zeros(nodeAttrs.shape[0], dtype=torch.long, device=model_device, requires_grad=False), - pbc=torch.tensor([periodic, periodic, periodic], dtype=torch.bool, device=model_device, requires_grad=False), + ptr=torch.tensor([0, nodeAttrs.shape[0]], dtype=torch.long, device=modelDevice, requires_grad=False), + nodeAttrs=nodeAttrs.to(dtype), + batch=torch.zeros(nodeAttrs.shape[0], dtype=torch.long, device=modelDevice, requires_grad=False), + pbc=torch.tensor([periodic, periodic, periodic], dtype=torch.bool, device=modelDevice, requires_grad=False), returnEnergyType=returnEnergyType, - charge=torch.tensor([float(args.get('charge', 0))], dtype=dtype, device=model_device, requires_grad=False), - multiplicity=torch.tensor([float(args.get('multiplicity', 1))], dtype=dtype, device=model_device, requires_grad=False), + charge=torch.tensor([float(args.get('charge', 0))], dtype=dtype, device=modelDevice, requires_grad=False), + multiplicity=torch.tensor([float(args.get('multiplicity', 1))], dtype=dtype, device=modelDevice, requires_grad=False), periodic=periodic, - indices=indices, - mmInfo=mmInfo) + mlIndices=mlIndices, + mmIndices=mmIndices, + mmCharges=mmCharges) force = openmm.PythonForce(compute) force.setForceGroup(forceGroup) force.setUsesPeriodicBoundaryConditions(periodic) - if atoms is not None and mmInfo is None: + if atoms is not None and mmIndices is None: force.setParticles(atoms) system.addForce(force) @@ -503,13 +507,16 @@ def createMixedSystem(self, # Validate model support before modifying the input system. model, device = self._loadModel(args) - _should_use_mm_embedding(model, atoms, embedding) + _validateMMEmbedding(model, atoms, embedding) # Validate the force-field setup before creating the modified system. # Multiple charge sources would make the MM charges ambiguous. nonbondedForces = [f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)] if len(nonbondedForces) > 1: - raise ValueError("Multiple NonbondedForce objects encountered; electrostatic embedding requires exactly one.") + raise ValueError( + "Multiple NonbondedForce objects encountered; electrostatic embedding " + "requires exactly one." + ) # The callback uses fixed charges and cannot follow charge offsets. for force in nonbondedForces: @@ -524,9 +531,15 @@ def createMixedSystem(self, # A CustomNonbondedForce's energy expression is arbitrary, so # whether it contains electrostatics cannot be determined here. if customNonbondedHasCharges is None: - raise ValueError("The System contains a CustomNonbondedForce and it is unknown whether it includes electrostatic interactions; pass customNonbondedHasCharges to specify.") + raise ValueError( + "The System contains a CustomNonbondedForce and it is unknown whether it " + "includes electrostatic interactions; pass customNonbondedHasCharges to specify." + ) if customNonbondedHasCharges and customNonbondedChargeParameter is None: - raise ValueError("A CustomNonbondedForce includes electrostatic interactions, so customNonbondedChargeParameter must name the per-particle parameter holding the charge.") + raise ValueError( + "A CustomNonbondedForce includes electrostatic interactions, so " + "customNonbondedChargeParameter must name the per-particle parameter holding the charge." + ) # A named charge parameter must exist in every custom force we modify. Checked here, before # the system is copied, so an unusable name fails before any surgery. @@ -589,17 +602,24 @@ def _customNonbondedChargeIndex(force: openmm.CustomNonbondedForce, name: str) - return names.index(name) -def _supports_mm_embedding(model) -> bool: +def _supportsMMEmbedding(model) -> bool: return bool(getattr(model, "supports_external_electrostatics", False)) -def _should_use_mm_embedding(model, atoms: Optional[Iterable[int]], embedding: str) -> bool: - """Validate the only MACE-specific embedding: electrostatic.""" +def _validateMMEmbedding(model, atoms: Optional[Iterable[int]], embedding: str) -> None: + """Raise if this model and atom selection cannot provide the requested embedding. + + Only "electrostatic" is MACE-specific; "mechanical" is handled by the generic embedding + plugin and needs nothing from us. + """ if embedding == "mechanical": - return False + return if embedding != "electrostatic": - raise ValueError(f"Unsupported embedding mode '{embedding}'.") - if not _supports_mm_embedding(model): + raise ValueError( + f"Unsupported embedding mode '{embedding}'; MACE implements 'electrostatic' " + "and delegates 'mechanical' to the generic embedding." + ) + if not _supportsMMEmbedding(model): raise ValueError( f"embedding='{embedding}' requires a model that accepts MM charges " f"and positions (PolarMACE); got {model.__class__.__name__}." @@ -609,7 +629,6 @@ def _should_use_mm_embedding(model, atoms: Optional[Iterable[int]], embedding: s f"embedding='{embedding}' requires an ML subset; it cannot be used " "with createSystem()." ) - return True def _prepareMMEmbedding(system: openmm.System, atoms: Optional[Iterable[int]], @@ -618,18 +637,18 @@ def _prepareMMEmbedding(system: openmm.System, atoms: Optional[Iterable[int]], if atoms is None: return None - num_particles = int(system.getNumParticles()) - ml_atoms = np.asarray(list(atoms), dtype=np.int64) - ml_set = set(int(i) for i in ml_atoms.tolist()) - mm_atoms = np.asarray( - [i for i in range(num_particles) if i not in ml_set], dtype=np.int64 + numParticles = int(system.getNumParticles()) + mlAtoms = np.asarray(list(atoms), dtype=np.int64) + mlSet = set(int(i) for i in mlAtoms.tolist()) + mmAtoms = np.asarray( + [i for i in range(numParticles) if i not in mlSet], dtype=np.int64 ) # The charges given to the model have to come from wherever the force field # actually keeps them, which is the same force createMixedSystem zeroed the # ML charges in. When that is a CustomNonbondedForce the caller has named # the parameter holding them. - mm_charges = np.empty(len(mm_atoms), dtype=np.float64) + mmCharges = np.empty(len(mmAtoms), dtype=np.float64) if customNonbondedChargeParameter is not None: custom = None @@ -643,12 +662,12 @@ def _prepareMMEmbedding(system: openmm.System, atoms: Optional[Iterable[int]], break if custom is None: raise ValueError(f"No CustomNonbondedForce defines a per-particle parameter '{customNonbondedChargeParameter}'.") - for row, atom_index in enumerate(mm_atoms): - mm_charges[row] = custom.getParticleParameters(int(atom_index))[chargeIndex] + for row, atomIndex in enumerate(mmAtoms): + mmCharges[row] = custom.getParticleParameters(int(atomIndex))[chargeIndex] return { - "ml_atoms": ml_atoms, - "mm_atoms": mm_atoms, - "mm_charges": mm_charges, + "ml_atoms": mlAtoms, + "mm_atoms": mmAtoms, + "mm_charges": mmCharges, } nonbonded = None @@ -661,49 +680,50 @@ def _prepareMMEmbedding(system: openmm.System, atoms: Optional[Iterable[int]], "PolarMACE MM embedding requires a NonbondedForce to source MM charges." ) - for row, atom_index in enumerate(mm_atoms): - charge, _, _ = nonbonded.getParticleParameters(int(atom_index)) - mm_charges[row] = charge.value_in_unit(unit.elementary_charge) + for row, atomIndex in enumerate(mmAtoms): + charge, _, _ = nonbonded.getParticleParameters(int(atomIndex)) + mmCharges[row] = charge.value_in_unit(unit.elementary_charge) return { - "ml_atoms": ml_atoms, - "mm_atoms": mm_atoms, - "mm_charges": mm_charges, + "ml_atoms": mlAtoms, + "mm_atoms": mmAtoms, + "mm_charges": mmCharges, } -def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, periodic, indices=None, mmInfo=None): +def _computeMACE(state, model, ptr, nodeAttrs, batch, pbc, returnEnergyType, charge, + multiplicity, periodic, mlIndices=None, mmIndices=None, mmCharges=None): import torch from mace.data.neighborhood import get_neighborhood energyScale = 96.4853 lengthScale = 10.0 - # With setParticles() (indices is None) the state already holds only the ML atoms; on the + # With setParticles() (mlIndices is None) the state already holds only the ML atoms; on the # electrostatic path it holds the whole system and the ML subset is sliced out here. - positions_full = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom) - numAtoms = positions_full.shape[0] - positions = positions_full if indices is None else positions_full[indices] + positionsFull = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom) + numAtoms = positionsFull.shape[0] + positions = positionsFull if mlIndices is None else positionsFull[mlIndices] if periodic: cell = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) else: cell = np.identity(3, dtype=np.float64) - dtype = node_attrs.dtype + dtype = nodeAttrs.dtype cutoff = float(model.r_max.detach()) edgeIndex, shifts, _, _ = get_neighborhood(positions, cutoff, [periodic, periodic, periodic], cell) - cell_tensor = torch.tensor(cell, dtype=dtype, device=ptr.device) - volume = torch.linalg.det(cell_tensor) + cellTensor = torch.tensor(cell, dtype=dtype, device=ptr.device) + volume = torch.linalg.det(cellTensor) if torch.abs(volume) > 0: - rcell = 2 * torch.pi * torch.linalg.inv(cell_tensor.mT) + rcell = 2 * torch.pi * torch.linalg.inv(cellTensor.mT) else: rcell = torch.zeros((3, 3), dtype=dtype, device=ptr.device) inputDict = { "ptr": ptr, - "node_attrs": node_attrs, + "node_attrs": nodeAttrs, "batch": batch, "pbc": pbc, "positions": torch.tensor(positions, dtype=dtype, device=ptr.device), "edge_index": torch.tensor(edgeIndex, dtype=torch.int64, device=ptr.device), "shifts": torch.tensor(shifts, dtype=dtype, device=ptr.device), - "cell": cell_tensor, + "cell": cellTensor, "rcell": rcell, "volume": volume.reshape(-1), "total_charge": charge, @@ -711,34 +731,37 @@ def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, ch "external_field": torch.zeros((charge.shape[0], 3), dtype=dtype, device=ptr.device), "fermi_level": torch.zeros((1,), dtype=dtype, device=ptr.device) } - # load mm position and charges - if mmInfo is not None: - mm_positions = positions_full[mmInfo["mm_atoms"]] + # Add the MM positions and charges for electrostatic embedding. + if mmIndices is not None: + mmPositions = positionsFull[mmIndices] inputDict["mm_positions"] = torch.tensor( - mm_positions, dtype=dtype, device=ptr.device + mmPositions, dtype=dtype, device=ptr.device ) inputDict["mm_charges"] = torch.tensor( - mmInfo["mm_charges"], dtype=dtype, device=ptr.device + mmCharges, dtype=dtype, device=ptr.device ) inputDict["mm_source_batch"] = torch.zeros( - len(mmInfo["mm_atoms"]), dtype=torch.long, device=ptr.device + len(mmIndices), dtype=torch.long, device=ptr.device ) - # eval and get results + # Evaluate the model and convert the energy and forces to OpenMM units. results = model(inputDict, compute_force=True) energy = float(results[returnEnergyType].detach())*energyScale forces = (results["forces"]*energyScale*lengthScale).detach().cpu().numpy() - mm_forces = results.get("mm_forces") + mmForces = results.get("mm_forces") - if mmInfo is not None and mm_forces is None: - raise ValueError("The model returned no 'mm_forces' although MM charges were supplied; it does not implement electrostatic embedding.") - if mm_forces is not None: - mm_forces = (mm_forces * energyScale * lengthScale).detach().cpu().numpy() + if mmIndices is not None and mmForces is None: + raise ValueError( + "The model returned no 'mm_forces' although MM charges were supplied; " + "it does not implement electrostatic embedding." + ) + if mmForces is not None: + mmForces = (mmForces * energyScale * lengthScale).detach().cpu().numpy() # Scatter ML and MM forces back to the full system. - if indices is not None: + if mlIndices is not None: f = np.zeros((numAtoms, 3), dtype=(np.float64 if dtype == torch.float64 else np.float32)) - f[indices] = forces - if mmInfo is not None and mm_forces is not None: - f[mmInfo["mm_atoms"]] += mm_forces.astype(f.dtype, copy=False) + f[mlIndices] = forces + if mmIndices is not None and mmForces is not None: + f[mmIndices] += mmForces.astype(f.dtype, copy=False) forces = f return energy, forces From fc68adba5d36c85a52ee5165aedbec8e32ff5a7c Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 17 Sep 2026 12:03:53 +0000 Subject: [PATCH 19/20] Drop the tests from this PR Removes the two test files this branch added (TestPolarMACEElectrostaticEmbedding.py, TestElectrostaticEmbeddingPME.py) and restores test/TestMACEPotential.py and test/TestMechanicalEmbedding.py to upstream. The test/ tree is now byte-identical to upstream and the PR is source-only: macepotential.py plus the requirements pin. Tests will be added back once their scope is agreed; the previous set was ~1200 lines across four files, most of it single-assertion checks on the force-field surgery at a granularity the rest of this repository does not use. This also fixes a real breakage: TestMACEPotential.py imported _should_use_mm_embedding, which the camelCase pass renamed, so that file failed at collection. It was not in the set I had been running, so the "104 passed" figure reported earlier never covered it. Restoring the upstream file removes the stale import; 129 tests now collect from the two upstream files. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN --- test/TestElectrostaticEmbeddingPME.py | 508 -------------------- test/TestMACEPotential.py | 192 -------- test/TestMechanicalEmbedding.py | 37 -- test/TestPolarMACEElectrostaticEmbedding.py | 484 ------------------- 4 files changed, 1221 deletions(-) delete mode 100644 test/TestElectrostaticEmbeddingPME.py delete mode 100644 test/TestPolarMACEElectrostaticEmbedding.py diff --git a/test/TestElectrostaticEmbeddingPME.py b/test/TestElectrostaticEmbeddingPME.py deleted file mode 100644 index a8ad78e..0000000 --- a/test/TestElectrostaticEmbeddingPME.py +++ /dev/null @@ -1,508 +0,0 @@ -"""Unit tests for the periodic NonbondedForce / PME path of -``MLPotential.createMixedSystem(embedding='electrostatic')``. - -These tests exercise the MM-side surgery only. A no-op MLPotentialImpl -is registered so the tests do not depend on MACE, torch, or GPU. - -Slice 4 of docs/codex-plans/pbc-electrostatic-embedding-small-plan.md. -""" -from __future__ import annotations - -import numpy as np -import openmm -import openmm.app as app -import openmm.unit as unit -import pytest - -from openmmml import MLPotential -from openmmml.mlpotential import MLPotentialImpl, MLPotentialImplFactory -from openmmml.models.macepotential import MACEPotentialImpl - - -# --------------------------------------------------------------------------- -# Register a no-op MLPotentialImpl for tests -# --------------------------------------------------------------------------- - -class _NoopImpl(MACEPotentialImpl): - """The real electrostatic embedding with the model evaluation stubbed out. - - Electrostatic embedding is a MACE-specific embedding method, so these tests - inherit MACEPotentialImpl.createMixedSystem() to exercise the actual - nonbonded surgery under test. Only addForces() is stubbed, which is what - would otherwise require a MACE checkpoint and the PolarMACE stack; the - surgery runs before it and is unaffected. - """ - - def __init__(self): - super().__init__("mace", None) - - def _loadModel(self, args): - # createMixedSystem() loads the model to check that it accepts MM - # charges. Stand in for a real checkpoint with an object of the class - # it looks for, so the surgery under test runs without the PolarMACE - # stack being installed. - class PolarMACE: - supports_external_electrostatics = True - - pass - - return PolarMACE(), "cpu" - - def addForces(self, topology, system, atoms, forceGroup, **args): - return - - -class _NoopFactory(MLPotentialImplFactory): - def createImpl(self, name, **args): - return _NoopImpl() - - -MLPotential.registerImplFactory("noop_test_impl", _NoopFactory()) - - -# --------------------------------------------------------------------------- -# Test system construction -# --------------------------------------------------------------------------- - -# 4 atoms: ML = {0, 1}, MM = {2, 3}. Periodic 2 nm cubic box. -_PARAMS = [ - # (mass, charge_e, sigma_nm, epsilon_kj) - (12.0, 0.6, 0.30, 0.20), # ML 0 - (12.0, -0.4, 0.32, 0.25), # ML 1 - (16.0, -0.8, 0.31, 0.65), # MM 2 - (1.0, 0.4, 0.10, 0.05), # MM 3 -] -_BOX_NM = 2.0 -_ML_ATOMS = [0, 1] -_MM_ATOMS = [2, 3] - - -def _build_periodic_system(use_pme=True, with_boundary_bonded=True, with_pre_existing_exception=False): - system = openmm.System() - box_vec = _BOX_NM * unit.nanometer - system.setDefaultPeriodicBoxVectors( - openmm.Vec3(_BOX_NM, 0, 0) * unit.nanometer, - openmm.Vec3(0, _BOX_NM, 0) * unit.nanometer, - openmm.Vec3(0, 0, _BOX_NM) * unit.nanometer, - ) - - nonbonded = openmm.NonbondedForce() - if use_pme: - nonbonded.setNonbondedMethod(openmm.NonbondedForce.PME) - nonbonded.setCutoffDistance(0.6 * unit.nanometer) - else: - nonbonded.setNonbondedMethod(openmm.NonbondedForce.NoCutoff) - for mass, charge, sigma, epsilon in _PARAMS: - system.addParticle(mass) - nonbonded.addParticle( - charge * unit.elementary_charge, - sigma * unit.nanometer, - epsilon * unit.kilojoule_per_mole, - ) - - # A pre-existing 1-3 ML-MM exception lets us check that sigma/epsilon - # are read from the exception rather than synthesized via Lorentz-Berthelot. - if with_pre_existing_exception: - nonbonded.addException( - 1, 2, - 0.123 * unit.elementary_charge * unit.elementary_charge, - 0.345 * unit.nanometer, - 0.678 * unit.kilojoule_per_mole, - ) - - system.addForce(nonbonded) - - if with_boundary_bonded: - # ML-ML bond + ML-MM bond (boundary) + MM-MM bond - bonds = openmm.HarmonicBondForce() - bonds.addBond(0, 1, 0.15 * unit.nanometer, 1000.0 * unit.kilojoule_per_mole / unit.nanometer ** 2) # ML-ML - bonds.addBond(1, 2, 0.15 * unit.nanometer, 800.0 * unit.kilojoule_per_mole / unit.nanometer ** 2) # ML-MM (boundary) - bonds.addBond(2, 3, 0.10 * unit.nanometer, 500.0 * unit.kilojoule_per_mole / unit.nanometer ** 2) # MM-MM - system.addForce(bonds) - - angles = openmm.HarmonicAngleForce() - angles.addAngle(0, 1, 2, 1.9, 100.0 * unit.kilojoule_per_mole / unit.radian ** 2) # ML-ML-MM (boundary) - angles.addAngle(1, 2, 3, 1.9, 100.0 * unit.kilojoule_per_mole / unit.radian ** 2) # ML-MM-MM (boundary) - system.addForce(angles) - - return system - - -def _build_topology(): - topology = app.Topology() - chain = topology.addChain() - res = topology.addResidue("MOL", chain) - elements = [app.element.carbon, app.element.carbon, app.element.oxygen, app.element.hydrogen] - for i, el in enumerate(elements): - topology.addAtom(f"A{i}", el, res) - topology.setPeriodicBoxVectors( - unit.Quantity(np.diag([_BOX_NM, _BOX_NM, _BOX_NM]), unit.nanometer) - ) - return topology - - -def _get_nonbonded(system): - for force in system.getForces(): - if isinstance(force, openmm.NonbondedForce): - return force - raise AssertionError("No NonbondedForce in system") - - -def _read_particle(force, i): - charge, sigma, epsilon = force.getParticleParameters(i) - return ( - charge.value_in_unit(unit.elementary_charge), - sigma.value_in_unit(unit.nanometer), - epsilon.value_in_unit(unit.kilojoule_per_mole), - ) - - -def _read_exception(force, i): - p1, p2, chargeProd, sigma, epsilon = force.getExceptionParameters(i) - return ( - int(p1), int(p2), - chargeProd.value_in_unit(unit.elementary_charge * unit.elementary_charge), - sigma.value_in_unit(unit.nanometer), - epsilon.value_in_unit(unit.kilojoule_per_mole), - ) - - -def _all_exceptions(force): - return { - tuple(sorted((p1, p2))): (cp, s, e) - for (p1, p2, cp, s, e) in (_read_exception(force, i) for i in range(force.getNumExceptions())) - } - - -def _make_mixed_system(**system_kwargs): - system = _build_periodic_system(**system_kwargs) - topology = _build_topology() - potential = MLPotential("noop_test_impl") - return potential.createMixedSystem( - topology, system, _ML_ATOMS, embedding="electrostatic" - ), system - - -# --------------------------------------------------------------------------- -# Slice 4 unit tests -# --------------------------------------------------------------------------- - -def test_pme_method_preserved(): - """The PME setting on NonbondedForce must survive the surgery.""" - mixed, _ = _make_mixed_system() - nb = _get_nonbonded(mixed) - assert nb.getNonbondedMethod() == openmm.NonbondedForce.PME - - -def test_ml_charges_zeroed(): - """ML particle charges must be zero so reciprocal-space ML-* contribs vanish.""" - mixed, _ = _make_mixed_system() - nb = _get_nonbonded(mixed) - for i in _ML_ATOMS: - charge, _, _ = _read_particle(nb, i) - assert charge == pytest.approx(0.0, abs=1e-12) - - -def test_mm_charges_unchanged(): - """MM particle charges must be untouched by the surgery.""" - mixed, original = _make_mixed_system() - nb_mixed = _get_nonbonded(mixed) - nb_orig = _get_nonbonded(original) - for i in _MM_ATOMS: - new_charge, _, _ = _read_particle(nb_mixed, i) - old_charge, _, _ = _read_particle(nb_orig, i) - assert new_charge == pytest.approx(old_charge, abs=1e-12) - - -def test_lj_parameters_preserved_for_all_particles(): - """sigma/epsilon must be untouched (LJ stays in MM force field).""" - mixed, original = _make_mixed_system() - nb_mixed = _get_nonbonded(mixed) - nb_orig = _get_nonbonded(original) - for i in range(nb_orig.getNumParticles()): - _, s_new, e_new = _read_particle(nb_mixed, i) - _, s_old, e_old = _read_particle(nb_orig, i) - assert s_new == pytest.approx(s_old, abs=1e-12) - assert e_new == pytest.approx(e_old, abs=1e-12) - - -def test_ml_mm_pairs_get_no_exception(): - """ML-MM pairs must be left on the ordinary pair list. - - ML-MM Coulomb is removed by zeroing the ML particle charges, not by adding - an exception per ML-MM pair. An exception would also be wrong: OpenMM - evaluates exceptions at the plain Cartesian distance rather than the - minimum image one, so under PBC the ML-MM Lennard-Jones interaction would - silently disappear for any pair that is only within the cutoff across a - periodic boundary. - """ - mixed, _ = _make_mixed_system(with_pre_existing_exception=False) - nb = _get_nonbonded(mixed) - excs = _all_exceptions(nb) - for ml in _ML_ATOMS: - for mm in _MM_ATOMS: - key = tuple(sorted((ml, mm))) - assert key not in excs, f"Unexpected ML-MM exception for {key}" - - -def test_ml_charges_zeroed_and_lj_untouched(): - """The Coulomb removal is done by zeroing the ML particle charges, which - leaves their Lennard-Jones parameters, and so ML-MM LJ, intact.""" - mixed, _ = _make_mixed_system(with_pre_existing_exception=False) - nb = _get_nonbonded(mixed) - for ml in _ML_ATOMS: - charge, sigma, epsilon = _read_particle(nb, ml) - assert charge == pytest.approx(0.0, abs=1e-15) - assert sigma == pytest.approx(_PARAMS[ml][2], rel=1e-12) - assert epsilon == pytest.approx(_PARAMS[ml][3], rel=1e-12) - for mm in _MM_ATOMS: - charge, sigma, epsilon = _read_particle(nb, mm) - assert charge == pytest.approx(_PARAMS[mm][1], rel=1e-12) - assert sigma == pytest.approx(_PARAMS[mm][2], rel=1e-12) - assert epsilon == pytest.approx(_PARAMS[mm][3], rel=1e-12) - - -def test_ml_mm_pre_existing_exception_keeps_lj_zeroes_charge(): - """If a pre-existing ML-MM exception had nonzero chargeProd, surgery must - zero it but keep the explicit sigma/epsilon.""" - mixed, _ = _make_mixed_system(with_pre_existing_exception=True) - nb = _get_nonbonded(mixed) - excs = _all_exceptions(nb) - cp, sigma, epsilon = excs[(1, 2)] - assert cp == pytest.approx(0.0, abs=1e-15) - assert sigma == pytest.approx(0.345, rel=1e-12) - assert epsilon == pytest.approx(0.678, rel=1e-12) - - -def test_ml_ml_exceptions_zero_charge_and_zero_lj(): - """ML-ML pairs are fully internal to MACE — both Coulomb and LJ must be - zeroed in the MM force field.""" - mixed, _ = _make_mixed_system() - nb = _get_nonbonded(mixed) - excs = _all_exceptions(nb) - key = tuple(sorted(_ML_ATOMS)) - assert key in excs - cp, _, epsilon = excs[key] - assert cp == pytest.approx(0.0, abs=1e-15) - assert epsilon == pytest.approx(0.0, abs=1e-15) - - -def test_mm_mm_exceptions_untouched(): - """MM-MM pairs must keep whatever the original system specified - (here: nothing — no exceptions added by surgery).""" - mixed, original = _make_mixed_system() - nb_mixed = _get_nonbonded(mixed) - nb_orig = _get_nonbonded(original) - excs_orig = _all_exceptions(nb_orig) - excs_mixed = _all_exceptions(nb_mixed) - mm_key = tuple(sorted(_MM_ATOMS)) - assert excs_mixed.get(mm_key) == excs_orig.get(mm_key) - - -def test_boundary_bonded_terms_preserved(): - """Boundary bonded terms (HarmonicBond/Angle that connect ML-MM) must - survive — they remain classical in electrostatic embedding.""" - mixed, _ = _make_mixed_system(with_boundary_bonded=True) - bond_force = next(f for f in mixed.getForces() if isinstance(f, openmm.HarmonicBondForce)) - bond_pairs = set() - for i in range(bond_force.getNumBonds()): - p1, p2, _, _ = bond_force.getBondParameters(i) - bond_pairs.add(tuple(sorted((int(p1), int(p2))))) - # ML-MM boundary bond and MM-MM bond must remain. - assert (1, 2) in bond_pairs - assert (2, 3) in bond_pairs - # ML-internal bond must be removed. - assert (0, 1) not in bond_pairs - - angle_force = next(f for f in mixed.getForces() if isinstance(f, openmm.HarmonicAngleForce)) - angle_triples = set() - for i in range(angle_force.getNumAngles()): - p1, p2, p3, _, _ = angle_force.getAngleParameters(i) - angle_triples.add((int(p1), int(p2), int(p3))) - # Boundary angles (touching at least one MM atom) preserved. - assert (0, 1, 2) in angle_triples - assert (1, 2, 3) in angle_triples - - -def test_reciprocal_space_pme_matches_mm_only_reference(): - """PME reciprocal-space contribution from the mixed-system NonbondedForce - must equal that of an MM-only reference (where ML charges are explicitly - zeroed but the system is otherwise identical).""" - mixed, original = _make_mixed_system(with_boundary_bonded=False) - - # Strip the bonded forces so only NonbondedForce contributes (we are - # comparing the reciprocal-space + direct-space Coulomb/LJ summed term). - def _strip_to_nonbonded(system): - # Remove non-NonbondedForce forces by index, descending. - for idx in reversed(range(system.getNumForces())): - if not isinstance(system.getForce(idx), openmm.NonbondedForce): - system.removeForce(idx) - return system - - _strip_to_nonbonded(mixed) - - # Build MM-only reference by zeroing ML charges in a fresh original system. - mm_only = _build_periodic_system(with_boundary_bonded=False) - nb_ref = _get_nonbonded(mm_only) - for i in _ML_ATOMS: - _, sigma, epsilon = nb_ref.getParticleParameters(i) - nb_ref.setParticleParameters(i, 0.0 * unit.elementary_charge, sigma, epsilon) - # Add the same exception structure (LB ML-MM zero-charge + ML-ML LJ-zero) - # so that direct-space exclusion subtractions match between the two systems. - # That is what the surgery itself produced; copy it onto the reference. - nb_mixed = _get_nonbonded(mixed) - existing_pairs = set( - tuple(sorted((int(p1), int(p2)))) - for i in range(nb_ref.getNumExceptions()) - for (p1, p2, *_rest) in [nb_ref.getExceptionParameters(i)] - ) - for i in range(nb_mixed.getNumExceptions()): - p1, p2, cp, sigma, epsilon = nb_mixed.getExceptionParameters(i) - if tuple(sorted((int(p1), int(p2)))) in existing_pairs: - continue - nb_ref.addException(int(p1), int(p2), cp, sigma, epsilon) - _strip_to_nonbonded(mm_only) - - positions = [ - openmm.Vec3(0.20, 0.30, 0.40), - openmm.Vec3(0.45, 0.30, 0.40), - openmm.Vec3(0.70, 0.30, 0.40), - openmm.Vec3(0.80, 0.30, 0.40), - ] * unit.nanometer - - platform = openmm.Platform.getPlatformByName("Reference") - ctx_mixed = openmm.Context(mixed, openmm.VerletIntegrator(0.001), platform) - ctx_ref = openmm.Context(mm_only, openmm.VerletIntegrator(0.001), platform) - ctx_mixed.setPositions(positions) - ctx_ref.setPositions(positions) - - e_mixed = ctx_mixed.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) - e_ref = ctx_ref.getState(getEnergy=True).getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) - - assert e_mixed == pytest.approx(e_ref, rel=1e-6, abs=1e-6) - - -# --------------------------------------------------------------------------- -# Electrostatics the surgery cannot account for. -# -# The ML-MM Coulomb is removed from the force field on the understanding that -# the model supplies it. Any Coulomb term this method cannot find is either left -# in place and counted twice, or removed and never replaced. Both give a wrong -# energy with no error, so these cases are refused rather than guessed at. -# --------------------------------------------------------------------------- - - -def _minimal_topology(numParticles): - from openmm.app import element - - topology = app.Topology() - chain = topology.addChain() - residue = topology.addResidue("X", chain) - for i in range(numParticles): - topology.addAtom(f"H{i}", element.hydrogen, residue) - return topology - - -def _createMixedSystem(system, **args): - impl = _NoopImpl() - return impl.createMixedSystem( - _minimal_topology(system.getNumParticles()), system, _ML_ATOMS, 0, False, - "electrostatic", **args, - ) - - -def test_multiple_nonbonded_forces_rejected(): - """The MM charges given to the model come from one NonbondedForce, so more - than one is ambiguous: the surgery would zero ML charges in all of them - while the model saw the charges of only the first.""" - system = _build_periodic_system() - system.addForce(openmm.NonbondedForce()) - with pytest.raises(ValueError, match="Multiple NonbondedForce"): - _createMixedSystem(system) - - -def _addCoulombCustomNonbondedForce(system): - force = openmm.CustomNonbondedForce("138.935456*q1*q2/r") - force.addPerParticleParameter("q") - for index in range(system.getNumParticles()): - force.addParticle([_PARAMS[index][1]]) - system.addForce(force) - return system - - -def test_custom_nonbonded_force_requires_an_answer(): - """A CustomNonbondedForce's energy expression is arbitrary, so whether it - carries electrostatics cannot be determined here and must be declared.""" - system = _addCoulombCustomNonbondedForce(_build_periodic_system()) - with pytest.raises(ValueError, match="unknown whether it includes electrostatic"): - _createMixedSystem(system) - - -def test_custom_nonbonded_force_with_charges_needs_the_parameter_name(): - """Declaring that it does carry electrostatics is not enough on its own: - the charge cannot be zeroed without knowing which parameter holds it.""" - system = _addCoulombCustomNonbondedForce(_build_periodic_system()) - with pytest.raises(ValueError, match="must name the per-particle parameter"): - _createMixedSystem(system, customNonbondedHasCharges=True) - - -def test_custom_nonbonded_force_unknown_parameter_name_rejected(): - """Naming a parameter the force does not define is an error, not a no-op - that would leave the electrostatics in place.""" - system = _addCoulombCustomNonbondedForce(_build_periodic_system()) - with pytest.raises(ValueError, match="no per-particle parameter"): - _createMixedSystem(system, customNonbondedHasCharges=True, - customNonbondedChargeParameter="charge") - - -def test_custom_nonbonded_charge_parameter_zeroed_on_ml_atoms(): - """Naming the charge parameter zeroes it on the ML atoms, which removes the - ML-MM Coulomb the custom force would otherwise still contribute.""" - system = _addCoulombCustomNonbondedForce(_build_periodic_system()) - mixed = _createMixedSystem(system, customNonbondedHasCharges=True, - customNonbondedChargeParameter="q") - custom = next(f for f in mixed.getForces() if isinstance(f, openmm.CustomNonbondedForce)) - for atom in _ML_ATOMS: - assert custom.getParticleParameters(atom)[0] == pytest.approx(0.0, abs=1e-12) - for atom in _MM_ATOMS: - assert custom.getParticleParameters(atom)[0] == pytest.approx(_PARAMS[atom][1], rel=1e-12) - - -def test_custom_nonbonded_force_without_charges_accepted(): - """Declaring it carries none proceeds, with ML-ML excluded as usual.""" - system = _addCoulombCustomNonbondedForce(_build_periodic_system()) - mixed = _createMixedSystem(system, customNonbondedHasCharges=False) - custom = next(f for f in mixed.getForces() if isinstance(f, openmm.CustomNonbondedForce)) - exclusions = { - tuple(sorted(custom.getExclusionParticles(i))) - for i in range(custom.getNumExclusions()) - } - assert tuple(sorted(_ML_ATOMS)) in exclusions - - -def test_custom_nonbonded_charges_reach_the_model(): - """When the electrostatics live in the CustomNonbondedForce, the charges the - model is given must come from there too. - - Reading them from the NonbondedForce in that case hands the model zeros - while the surgery has already removed the real ML-MM Coulomb, so the - interaction disappears rather than being computed by the model. - """ - from openmmml.models.macepotential import _prepareMMEmbedding - - system = openmm.System() - for _ in range(4): - system.addParticle(1.0) - lj = openmm.NonbondedForce() # Lennard-Jones only, no charges - for _, _, sigma, epsilon in _PARAMS: - lj.addParticle(0.0, sigma, epsilon) - system.addForce(lj) - coulomb = openmm.CustomNonbondedForce("138.935456*q1*q2/r") - coulomb.addPerParticleParameter("q") - for _, charge, _, _ in _PARAMS: - coulomb.addParticle([charge]) - system.addForce(coulomb) - - expected = [_PARAMS[i][1] for i in _MM_ATOMS] - charges = _prepareMMEmbedding(system, _ML_ATOMS, "q")["mm_charges"] - np.testing.assert_allclose(charges, expected, atol=1e-12) diff --git a/test/TestMACEPotential.py b/test/TestMACEPotential.py index 701f62c..4470278 100644 --- a/test/TestMACEPotential.py +++ b/test/TestMACEPotential.py @@ -5,175 +5,14 @@ import openmm.app as app import openmm.unit as unit import pytest -import torch from openmmml import MLPotential -from openmmml.models.macepotential import ( - MACEPotentialImpl, - _computeMACE, - _prepareMMEmbedding, - _should_use_mm_embedding, -) mace = pytest.importorskip("mace", reason="mace is not installed") platform_ints = range(mm.Platform.getNumPlatforms()) # Get the path to the test data test_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") - -def _simple_nonbonded_system(): - system = mm.System() - nonbonded = mm.NonbondedForce() - params = [ - (1.0, 0.30, 0.20), - (-0.5, 0.40, 0.50), - (0.25, 0.50, 0.80), - ] - for charge_e, sigma_nm, epsilon_kj in params: - system.addParticle(12.0) - nonbonded.addParticle( - charge_e * unit.elementary_charge, - sigma_nm * unit.nanometer, - epsilon_kj * unit.kilojoule_per_mole, - ) - system.addForce(nonbonded) - return system, nonbonded - - -def testPrepareMMEmbedding(): - """The MM complement and its charges are read off the NonbondedForce.""" - system, nonbonded = _simple_nonbonded_system() - info = _prepareMMEmbedding(system, [0, 1]) - assert info is not None - np.testing.assert_array_equal(info["ml_atoms"], [0, 1]) - np.testing.assert_array_equal(info["mm_atoms"], [2]) - np.testing.assert_allclose(info["mm_charges"], [0.25], atol=1e-12) - - -@pytest.mark.parametrize("kind", ["ml_particle", "mm_particle", "exception"]) -def testElectrostaticRejectsChargeOffsets(monkeypatch, kind): - system, force = _simple_nonbonded_system() - force.addGlobalParameter("lambda", 1.0) - if kind == "exception": - index = force.addException(0, 2, 0.25, 0.3, 0.0) - force.addExceptionParameterOffset("lambda", index, 0.5, 0.0, 0.0) - else: - force.addParticleParameterOffset("lambda", 0 if kind == "ml_particle" else 2, 0.5, 0.0, 0.0) - before = mm.XmlSerializer.serialize(system) - impl = MACEPotentialImpl("mace", None) - monkeypatch.setattr(impl, "_loadModel", lambda args: (PolarMACE(), "cpu")) - with pytest.raises(ValueError, match="charge parameter offsets"): - impl.createMixedSystem(app.Topology(), system, [0], 0, False, "electrostatic") - assert mm.XmlSerializer.serialize(system) == before - - -class _FakeState: - def __init__(self, positions_angstrom): - self._positions = np.asarray(positions_angstrom, dtype=np.float64) * unit.angstrom - - def getPositions(self, asNumpy=False): - return self._positions - - -class _FakeModel: - def __init__(self, dtype=torch.float32): - self.r_max = torch.tensor(3.0, dtype=dtype) - self.dtype = dtype - - def __call__(self, input_dict, compute_force=True): - del compute_force - n_ml = input_dict["positions"].shape[0] - forces = torch.tensor( - [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]][:n_ml], - dtype=self.dtype, - device=input_dict["positions"].device, - ) - out = { - "interaction_energy": torch.tensor([2.5], dtype=self.dtype, device=forces.device), - "forces": forces, - } - if "mm_positions" in input_dict: - out["mm_forces"] = torch.tensor( - [[0.5, 0.25, -0.75]], - dtype=self.dtype, - device=forces.device, - ) - return out - - -class PolarMACE: - supports_external_electrostatics = True - - -class MACE: - pass - - -def testComputeMACEScattersMMForces(): - ptr = torch.tensor([0, 2], dtype=torch.long) - node_attrs = torch.ones((2, 1), dtype=torch.float32) - batch = torch.zeros(2, dtype=torch.long) - pbc = torch.tensor([False, False, False], dtype=torch.bool) - mm_info = {"mm_atoms": np.array([2], dtype=np.int64), "mm_charges": np.array([0.25], dtype=np.float64)} - - energy, forces = _computeMACE( - state=_FakeState([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]]), - model=_FakeModel(dtype=torch.float32), - ptr=ptr, - node_attrs=node_attrs, - batch=batch, - pbc=pbc, - returnEnergyType="interaction_energy", - charge=torch.tensor([0.0], dtype=torch.float32), - multiplicity=torch.tensor([1.0], dtype=torch.float32), - indices=np.array([0, 1], dtype=np.int64), - periodic=False, - mmInfo=mm_info, - ) - - assert energy == pytest.approx(2.5 * 96.4853) - expected = np.array( - [ - [1.0, 2.0, 3.0], - [-1.0, -2.0, -3.0], - [0.5, 0.25, -0.75], - ], - dtype=np.float32, - ) * (96.4853 * 10.0) - np.testing.assert_allclose(forces, expected, rtol=1e-6, atol=1e-6) - - -def testShouldUseMMEmbedding(): - assert _should_use_mm_embedding(PolarMACE(), [0, 1], "electrostatic") - assert not _should_use_mm_embedding(PolarMACE(), [0, 1], "mechanical") - assert not _should_use_mm_embedding(MACE(), [0, 1], "mechanical") - with pytest.raises(ValueError, match="Unsupported embedding mode"): - _should_use_mm_embedding(PolarMACE(), [0, 1], "bad-mode") - - -def testElectrostaticEmbeddingRejectsUnsupportedCases(): - """Electrostatic embedding must fail loudly rather than falling back to - mechanical embedding: by the time the potential is asked to add its forces, - the ML-MM electrostatics have already been removed from the MM force field, - so a silent fallback would simply lose them.""" - with pytest.raises(ValueError, match="requires a model that accepts MM charges"): - _should_use_mm_embedding(MACE(), [0, 1], "electrostatic") - with pytest.raises(ValueError, match="requires an ML subset"): - _should_use_mm_embedding(PolarMACE(), None, "electrostatic") - - -def testOnlyPolarModelsOfferElectrostaticEmbedding(): - """Electrostatic embedding needs a model that accepts MM charges, which of - the pretrained models only the PolarMACE family does. A custom checkpoint - may be one, so it is offered there too and checked once the model loads.""" - for name, (_, _, _, _, acceptsMMCharges) in MACEPotentialImpl.KNOWN_MODELS.items(): - embeddings = MLPotential(name).getSupportedEmbeddings() - assert ('electrostatic' in embeddings) == acceptsMMCharges, name - assert acceptsMMCharges == name.startswith('mace-polar'), name - custom = MLPotential('mace', modelPath='unused-until-forces-are-added.model') - assert 'electrostatic' in custom.getSupportedEmbeddings() - - @pytest.mark.parametrize("platform_int", list(platform_ints)) class TestMACE: @@ -266,34 +105,3 @@ def testPrecisionApplied(self, platform_int, precision): "Energy is not finite under precision {}".format(precision) assert np.isclose(energyML, self.refEnergy['mace-off23-small'], rtol=1e-6),\ "Energy is not close to reference under precision {}".format(precision) - - -def testMissingMMForcesIsAnError(): - """A model given MM charges must return forces on the MM atoms.""" - class _NoMMForces(_FakeModel): - def __call__(self, inputDict, compute_force=True): - out = super().__call__(inputDict, compute_force) - out.pop("mm_forces", None) - return out - - mmInfo = { - "mm_atoms": np.array([2], dtype=np.int64), - "mm_charges": np.array([0.25], dtype=np.float64), - } - with pytest.raises(ValueError, match="no 'mm_forces'"): - _computeMACE( - state=_FakeState( - [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [3.0, 0.0, 0.0]] - ), - model=_NoMMForces(dtype=torch.float32), - ptr=torch.tensor([0, 2], dtype=torch.long), - node_attrs=torch.ones((2, 1), dtype=torch.float32), - batch=torch.zeros(2, dtype=torch.long), - pbc=torch.tensor([False, False, False], dtype=torch.bool), - returnEnergyType="interaction_energy", - charge=torch.tensor([0.0], dtype=torch.float32), - multiplicity=torch.tensor([1.0], dtype=torch.float32), - indices=np.array([0, 1], dtype=np.int64), - periodic=False, - mmInfo=mmInfo, - ) diff --git a/test/TestMechanicalEmbedding.py b/test/TestMechanicalEmbedding.py index 0257437..ac01f50 100644 --- a/test/TestMechanicalEmbedding.py +++ b/test/TestMechanicalEmbedding.py @@ -489,40 +489,3 @@ def testLinkAtomInfo(self, platform_int): assert mixed_info["oldToNew"] == list(range(original_count)) for i in range(mixed_system.getNumParticles()): assert mixed_system.isVirtualSite(i) == (i >= original_count) - - -def testCustomNonbondedForce(): - """Mechanical embedding excludes ML-ML CustomNonbondedForce pairs.""" - numParticles = 4 - mlAtoms = [0, 1] - system = openmm.System() - nonbonded = openmm.NonbondedForce() - custom = openmm.CustomNonbondedForce("4*epsilon*((sigma/r)^12-(sigma/r)^6)") - custom.addPerParticleParameter("sigma") - custom.addPerParticleParameter("epsilon") - for _ in range(numParticles): - system.addParticle(1.0) - nonbonded.addParticle(0.0, 0.3, 0.2) - custom.addParticle([0.3, 0.2]) - system.addForce(nonbonded) - system.addForce(custom) - - topology = openmm.app.Topology() - chain = topology.addChain() - residue = topology.addResidue("X", chain) - for i in range(numParticles): - topology.addAtom(f"H{i}", openmm.app.element.hydrogen, residue) - - mixed = MLPotential("mace-off23-small").createMixedSystem( - topology, system, mlAtoms - ) - mixedCustom = next( - force - for force in mixed.getForces() - if isinstance(force, openmm.CustomNonbondedForce) - ) - exclusions = { - tuple(sorted(mixedCustom.getExclusionParticles(i))) - for i in range(mixedCustom.getNumExclusions()) - } - assert tuple(sorted(mlAtoms)) in exclusions diff --git a/test/TestPolarMACEElectrostaticEmbedding.py b/test/TestPolarMACEElectrostaticEmbedding.py deleted file mode 100644 index ff00400..0000000 --- a/test/TestPolarMACEElectrostaticEmbedding.py +++ /dev/null @@ -1,484 +0,0 @@ -"""End-to-end tests for PolarMACE electrostatic embedding via openmmml. - -Covers both: - -* Non-periodic plumbing: MM positions/charges flow into PolarMACE, - electrostatic embedding shifts the ML energy and creates a back-reaction - force on the MM atoms (Slice 4 cross-check, fills the gap that the only - existing electrostatic-embedding tests in TestMACEPotential.py exercise - the helpers in isolation rather than the full plumbing). - -* PBC: ``MLPotential('mace', modelPath=...).createMixedSystem( - embedding='electrostatic')`` is translation-invariant under PBC -- - shifting all positions by one full box vector must leave the potential - energy and per-atom forces unchanged. (Slice 5 of - docs/codex-plans/pbc-electrostatic-embedding-small-plan.md.) - -The tests train nothing -- a tiny random PolarMACE is built and serialized -in a module-scoped fixture so the suite runs in seconds on CPU. -""" -from __future__ import annotations - -import os -import tempfile - -import numpy as np -import openmm -import openmm.app as app -import openmm.unit as unit -import pytest -import torch - -torch.serialization.add_safe_globals([slice]) - -mace = pytest.importorskip("mace", reason="mace is not installed") - -# OpenMM-ML installs its external-source adapter when this checkpoint is loaded. -# The checkpoint itself stays a stock PolarMACE and does not need the mlmm -# monkey-patch that these tests previously imported. - -from e3nn import o3 # noqa: E402 - -from openmmml import MLPotential # noqa: E402 -from openmmml.models import macepotential # noqa: E402 -from mace.modules import interaction_classes # noqa: E402 -from mace.modules.extensions import PolarMACE # noqa: E402 - - -_BOX_NM = 1.4 # 14 Angstrom box (large enough for r_max=4 A and PME) -_DTYPE = torch.float64 - - -# --------------------------------------------------------------------------- -# PolarMACE model + save to tempfile -# --------------------------------------------------------------------------- - -def _build_polar_mace(device: torch.device, dtype: torch.dtype) -> PolarMACE: - fixedpoint_update_config = { - "type": "AgnosticEmbeddedOneBodyVariableUpdate", - "potential_embedding_cls": "AgnosticChargeBiasedLinearPotentialEmbedding", - "nonlinearity_cls": "MLPNonLinearity", - } - field_readout_config = {"type": "OneBodyMLPFieldReadout"} - return PolarMACE( - r_max=4.0, - num_bessel=4, - num_polynomial_cutoff=3, - max_ell=1, - interaction_cls=interaction_classes[ - "RealAgnosticResidualNonLinearInteractionBlock" - ], - interaction_cls_first=interaction_classes[ - "RealAgnosticResidualNonLinearInteractionBlock" - ], - num_interactions=2, - num_elements=2, - hidden_irreps=o3.Irreps("4x0e + 4x1o"), - MLP_irreps=o3.Irreps("8x0e"), - atomic_energies=torch.zeros(2, dtype=dtype, device=device), - avg_num_neighbors=3.0, - atomic_numbers=[1, 8], - correlation=1, - gate=torch.nn.functional.silu, - radial_MLP=[16, 16], - radial_type="bessel", - kspace_cutoff_factor=1.0, - atomic_multipoles_max_l=1, - atomic_multipoles_smearing_width=1.0, - field_feature_max_l=1, - field_feature_widths=[1.0], - field_feature_norms=[1.0, 1.0], - num_recursion_steps=1, - field_si=False, - include_electrostatic_self_interaction=False, - add_local_electron_energy=True, - field_dependence_type="AgnosticEmbeddedOneBodyVariableUpdate", - final_field_readout_type="OneBodyMLPFieldReadout", - return_electrostatic_potentials=False, - heads=["Default"], - field_norm_factor=1.0, - fixedpoint_update_config=fixedpoint_update_config, - field_readout_config=field_readout_config, - ).to(device=device, dtype=dtype) - - -@pytest.fixture(scope="module") -def polar_mace_model_path(): - torch.manual_seed(7) - model = _build_polar_mace(torch.device("cpu"), _DTYPE) - model.eval() - tmpdir = tempfile.mkdtemp(prefix="polar_mace_pbc_test_") - path = os.path.join(tmpdir, "polar_mace_test.pt") - torch.save(model, path) - yield path - try: - os.remove(path) - os.rmdir(tmpdir) - except OSError: - pass - - -# --------------------------------------------------------------------------- -# Periodic OpenMM topology + system: 1 ML water + N MM waters -# --------------------------------------------------------------------------- - -# TIP3P-ish charges and LJ for the test (units consistent with OpenMM defaults). -_O_CHARGE = -0.834 -_H_CHARGE = 0.417 -_O_SIGMA_NM = 0.31507 -_O_EPS_KJ = 0.6364 -_H_SIGMA_NM = 1.0e-3 -_H_EPS_KJ = 0.0 - - -def _add_water(system, nonbonded, bonds, angles, masses_charges_lj): - """Append a 3-atom water to system+forces; returns particle indices.""" - indices = [] - for mass, charge, sigma, epsilon in masses_charges_lj: - idx = system.addParticle(mass) - nonbonded.addParticle( - charge * unit.elementary_charge, - sigma * unit.nanometer, - epsilon * unit.kilojoule_per_mole, - ) - indices.append(idx) - o, h1, h2 = indices - bonds.addBond(o, h1, 0.09572 * unit.nanometer, 4.5e5 * unit.kilojoule_per_mole / unit.nanometer ** 2) - bonds.addBond(o, h2, 0.09572 * unit.nanometer, 4.5e5 * unit.kilojoule_per_mole / unit.nanometer ** 2) - angles.addAngle(h1, o, h2, 1.824, 460.0 * unit.kilojoule_per_mole / unit.radian ** 2) - nonbonded.addException( - o, h1, - 0.0 * unit.elementary_charge * unit.elementary_charge, - 0.5 * (_O_SIGMA_NM + _H_SIGMA_NM) * unit.nanometer, - 0.0 * unit.kilojoule_per_mole, - ) - nonbonded.addException( - o, h2, - 0.0 * unit.elementary_charge * unit.elementary_charge, - 0.5 * (_O_SIGMA_NM + _H_SIGMA_NM) * unit.nanometer, - 0.0 * unit.kilojoule_per_mole, - ) - nonbonded.addException( - h1, h2, - 0.0 * unit.elementary_charge * unit.elementary_charge, - _H_SIGMA_NM * unit.nanometer, - 0.0 * unit.kilojoule_per_mole, - ) - return indices - - -def _build_topology_and_system(num_mm_waters=3, periodic=True, mm_oxygen_charge=_O_CHARGE): - topology = app.Topology() - chain = topology.addChain() - - # ML water (residue 0). - res = topology.addResidue("HOH", chain) - a_o = topology.addAtom("O", app.element.oxygen, res) - a_h1 = topology.addAtom("H1", app.element.hydrogen, res) - a_h2 = topology.addAtom("H2", app.element.hydrogen, res) - topology.addBond(a_o, a_h1) - topology.addBond(a_o, a_h2) - - for _ in range(num_mm_waters): - res = topology.addResidue("HOH", chain) - a_o = topology.addAtom("O", app.element.oxygen, res) - a_h1 = topology.addAtom("H1", app.element.hydrogen, res) - a_h2 = topology.addAtom("H2", app.element.hydrogen, res) - topology.addBond(a_o, a_h1) - topology.addBond(a_o, a_h2) - - if periodic: - topology.setPeriodicBoxVectors( - unit.Quantity(np.diag([_BOX_NM, _BOX_NM, _BOX_NM]), unit.nanometer) - ) - - system = openmm.System() - if periodic: - system.setDefaultPeriodicBoxVectors( - openmm.Vec3(_BOX_NM, 0, 0) * unit.nanometer, - openmm.Vec3(0, _BOX_NM, 0) * unit.nanometer, - openmm.Vec3(0, 0, _BOX_NM) * unit.nanometer, - ) - nonbonded = openmm.NonbondedForce() - if periodic: - nonbonded.setNonbondedMethod(openmm.NonbondedForce.PME) - nonbonded.setCutoffDistance(0.5 * unit.nanometer) - else: - nonbonded.setNonbondedMethod(openmm.NonbondedForce.NoCutoff) - bonds = openmm.HarmonicBondForce() - angles = openmm.HarmonicAngleForce() - - ml_o = (15.999, _O_CHARGE, _O_SIGMA_NM, _O_EPS_KJ) - ml_h = (1.008, _H_CHARGE, _H_SIGMA_NM, _H_EPS_KJ) - mm_o = (15.999, mm_oxygen_charge, _O_SIGMA_NM, _O_EPS_KJ) - mm_h_charge = -mm_oxygen_charge / 2.0 - mm_h = (1.008, mm_h_charge, _H_SIGMA_NM, _H_EPS_KJ) - - _add_water(system, nonbonded, bonds, angles, [ml_o, ml_h, ml_h]) - for _ in range(num_mm_waters): - _add_water(system, nonbonded, bonds, angles, [mm_o, mm_h, mm_h]) - - system.addForce(nonbonded) - system.addForce(bonds) - system.addForce(angles) - - return topology, system - - -def _initial_positions(num_mm_waters): - """Return positions in nanometers. ML water at ~box centre, MM waters scattered.""" - base_o = np.array([ - [0.55, 0.55, 0.55], # ML O - [0.20, 0.20, 0.20], - [0.95, 0.20, 0.30], - [0.30, 0.95, 0.95], - [0.95, 0.95, 0.95], - [0.10, 0.50, 0.10], - ])[: 1 + num_mm_waters] - H1_OFFSET = np.array([0.09572, 0.0, 0.0]) - H2_OFFSET = np.array([-0.0240, 0.0927, 0.0]) # 104.5 deg, 0.09572 nm bond - coords = [] - for o in base_o: - coords.append(o) - coords.append(o + H1_OFFSET) - coords.append(o + H2_OFFSET) - return np.array(coords) * unit.nanometer - - -# --------------------------------------------------------------------------- -# Integration tests -# --------------------------------------------------------------------------- - -def _energy_and_forces(context): - state = context.getState(getEnergy=True, getForces=True) - e = state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole) - f = state.getForces(asNumpy=True).value_in_unit(unit.kilojoule_per_mole / unit.nanometer) - return float(e), np.asarray(f) - - -def test_polar_mace_pbc_runs_and_is_finite(polar_mace_model_path): - """Smoke test: createMixedSystem(embedding='electrostatic') with PolarMACE - under PBC returns a finite energy.""" - topology, mm_system = _build_topology_and_system(num_mm_waters=2) - potential = MLPotential("mace", modelPath=polar_mace_model_path) - mixed_system = potential.createMixedSystem( - topology, mm_system, [0, 1, 2], embedding="electrostatic" - ) - platform = openmm.Platform.getPlatformByName("Reference") - context = openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) - context.setPositions(_initial_positions(num_mm_waters=2)) - e, f = _energy_and_forces(context) - assert np.isfinite(e) - assert np.all(np.isfinite(f)) - - -@pytest.mark.parametrize("periodic", [False, True]) -def test_polar_mace_default_returns_total_energy(polar_mace_model_path, periodic): - topology, system = _build_topology_and_system(num_mm_waters=2, periodic=periodic) - results = [] - for options in ({}, {"returnEnergyType": "energy"}): - potential = MLPotential("mace", modelPath=polar_mace_model_path) - mixed = potential.createMixedSystem( - topology, system, [0, 1, 2], embedding="electrostatic", **options - ) - context = openmm.Context(mixed, openmm.VerletIntegrator(0.001), - openmm.Platform.getPlatformByName("Reference")) - context.setPositions(_initial_positions(num_mm_waters=2)) - results.append(_energy_and_forces(context)) - del context - np.testing.assert_allclose(results[0][0], results[1][0], rtol=0, atol=1e-9) - np.testing.assert_allclose(results[0][1], results[1][1], rtol=0, atol=1e-8) - - -def test_polar_mace_pbc_translation_invariance(polar_mace_model_path): - """Energy and forces must be invariant under whole-system translation by - one full box vector — the fundamental PBC sanity check.""" - topology, mm_system = _build_topology_and_system(num_mm_waters=2) - potential = MLPotential("mace", modelPath=polar_mace_model_path) - mixed_system = potential.createMixedSystem( - topology, mm_system, [0, 1, 2], embedding="electrostatic" - ) - platform = openmm.Platform.getPlatformByName("Reference") - context = openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) - - pos = _initial_positions(num_mm_waters=2) - context.setPositions(pos) - e0, f0 = _energy_and_forces(context) - - # Translate all atoms by +x box length. Forces are reported per-atom in the - # same order, so direct comparison is valid. - pos_shifted = ( - pos.value_in_unit(unit.nanometer) + np.array([_BOX_NM, 0.0, 0.0]) - ) * unit.nanometer - context.setPositions(pos_shifted) - e1, f1 = _energy_and_forces(context) - - assert e1 == pytest.approx(e0, rel=1e-6, abs=1e-4) - np.testing.assert_allclose(f1, f0, rtol=1e-5, atol=1e-3) - - -def test_polar_mace_pbc_ml_atom_near_boundary(polar_mace_model_path): - """If an ML atom sits near a periodic boundary, its energy must match the - energy from the unwrapped image (translation by one full box must give - the same result).""" - topology, mm_system = _build_topology_and_system(num_mm_waters=2) - potential = MLPotential("mace", modelPath=polar_mace_model_path) - mixed_system = potential.createMixedSystem( - topology, mm_system, [0, 1, 2], embedding="electrostatic" - ) - platform = openmm.Platform.getPlatformByName("Reference") - context = openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) - - # Put the ML water O near x=0 (boundary) with H atoms straddling the box. - pos_np = _initial_positions(num_mm_waters=2).value_in_unit(unit.nanometer) - pos_np[0] = np.array([0.05, 0.55, 0.55]) # O near +x boundary edge - pos_np[1] = pos_np[0] + np.array([0.09572, 0.0, 0.0]) # H still inside box - pos_np[2] = pos_np[0] + np.array([-0.024, 0.0927, 0.0]) # H near boundary - context.setPositions(pos_np * unit.nanometer) - e_near, f_near = _energy_and_forces(context) - assert np.isfinite(e_near) - assert np.all(np.isfinite(f_near)) - - # Translate everything by one full box: identical energy and forces. - context.setPositions((pos_np + np.array([_BOX_NM, 0.0, 0.0])) * unit.nanometer) - e_shift, f_shift = _energy_and_forces(context) - assert e_shift == pytest.approx(e_near, rel=1e-6, abs=1e-4) - np.testing.assert_allclose(f_shift, f_near, rtol=1e-5, atol=1e-3) - - -# --------------------------------------------------------------------------- -# Non-PBC tests: MACE electrostatic plumbing through openmmml -# --------------------------------------------------------------------------- - -def _build_nonpbc_mixed(potential, num_mm_waters, embedding, mm_oxygen_charge=_O_CHARGE): - topology, mm_system = _build_topology_and_system( - num_mm_waters=num_mm_waters, periodic=False, mm_oxygen_charge=mm_oxygen_charge - ) - mixed_system = potential.createMixedSystem( - topology, mm_system, [0, 1, 2], embedding=embedding - ) - platform = openmm.Platform.getPlatformByName("Reference") - return mixed_system, openmm.Context(mixed_system, openmm.VerletIntegrator(0.001), platform) - - -def _nonpbc_initial_positions(num_mm_waters): - base_o = np.array([ - [0.0, 0.0, 0.0], # ML O at origin - [0.50, 0.0, 0.0], # MM water on +x - [0.0, 0.50, 0.0], # MM water on +y - [0.0, 0.0, 0.50], # MM water on +z - ])[: 1 + num_mm_waters] - H1_OFFSET = np.array([0.09572, 0.0, 0.0]) - H2_OFFSET = np.array([-0.0240, 0.0927, 0.0]) - coords = [] - for o in base_o: - coords.append(o) - coords.append(o + H1_OFFSET) - coords.append(o + H2_OFFSET) - return np.array(coords) * unit.nanometer - - -def test_polar_mace_electrostatic_nonpbc_smoke(polar_mace_model_path): - """Non-PBC: createMixedSystem(embedding='electrostatic') with PolarMACE - runs end-to-end and produces finite energy/forces.""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - _, ctx = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="electrostatic") - ctx.setPositions(_nonpbc_initial_positions(num_mm_waters=2)) - e, f = _energy_and_forces(ctx) - assert np.isfinite(e) - assert np.all(np.isfinite(f)) - - -def test_polar_mace_electrostatic_changes_energy_vs_mechanical(polar_mace_model_path): - """Mechanical embedding does NOT pass MM charges into MACE; electrostatic - does. With nonzero MM charges, the two modes must give different ML energies - -- the only path that exercises the openmmml -> mm_charges plumbing.""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - positions = _nonpbc_initial_positions(num_mm_waters=2) - - _, ctx_mech = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="mechanical") - ctx_mech.setPositions(positions) - e_mech, _ = _energy_and_forces(ctx_mech) - - _, ctx_elec = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="electrostatic") - ctx_elec.setPositions(positions) - e_elec, _ = _energy_and_forces(ctx_elec) - - # Energies must not coincide -- with TIP3P-ish MM charges there is a real - # ML-MM Coulomb piece that mechanical embedding doesn't see. - assert abs(e_elec - e_mech) > 1e-3 - - -def test_polar_mace_electrostatic_zero_mm_charges_matches_mechanical(polar_mace_model_path): - """If MM charges are zero, the ML-MM Coulomb term vanishes and electrostatic - embedding should reproduce mechanical embedding (within float64 round-off).""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - positions = _nonpbc_initial_positions(num_mm_waters=2) - - _, ctx_mech = _build_nonpbc_mixed( - potential, num_mm_waters=2, embedding="mechanical", mm_oxygen_charge=0.0 - ) - ctx_mech.setPositions(positions) - e_mech, f_mech = _energy_and_forces(ctx_mech) - - _, ctx_elec = _build_nonpbc_mixed( - potential, num_mm_waters=2, embedding="electrostatic", mm_oxygen_charge=0.0 - ) - ctx_elec.setPositions(positions) - e_elec, f_elec = _energy_and_forces(ctx_elec) - - assert e_elec == pytest.approx(e_mech, rel=1e-5, abs=1e-3) - np.testing.assert_allclose(f_elec, f_mech, rtol=1e-4, atol=1e-2) - - -def test_polar_mace_electrostatic_mm_atoms_receive_back_reaction(polar_mace_model_path): - """The MACE mm_forces path must scatter forces onto MM atoms. With nonzero - MM charges, MM atoms must feel a non-trivial force; with zero MM charges - they must not (modulo MM-internal forces in the unmodified pieces of the - system, which we factor out by subtraction).""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - positions = _nonpbc_initial_positions(num_mm_waters=2) - - _, ctx_charged = _build_nonpbc_mixed( - potential, num_mm_waters=2, embedding="electrostatic", mm_oxygen_charge=_O_CHARGE - ) - ctx_charged.setPositions(positions) - _, f_charged = _energy_and_forces(ctx_charged) - - _, ctx_neutral = _build_nonpbc_mixed( - potential, num_mm_waters=2, embedding="electrostatic", mm_oxygen_charge=0.0 - ) - ctx_neutral.setPositions(positions) - _, f_neutral = _energy_and_forces(ctx_neutral) - - # MM atoms are indices 3..8 (2 waters); their force differs between charged - # and neutral runs, and that difference is exactly the ML-MM back reaction - # routed through openmmml. - delta_mm = f_charged[3:] - f_neutral[3:] - assert np.linalg.norm(delta_mm) > 1e-2 - - -def test_polar_mace_electrostatic_mm_charge_displacement_changes_ml_force( - polar_mace_model_path, -): - """Moving an MM charge changes the field on the ML region and therefore - the force on the ML atoms. This is the cleanest end-to-end check that the - field embedding is wired through openmmml into MACE.""" - potential = MLPotential("mace", modelPath=polar_mace_model_path) - - pos_a = _nonpbc_initial_positions(num_mm_waters=2) - # Move the second MM water (atoms 6, 7, 8) by +0.05 nm in x. - pos_b_np = pos_a.value_in_unit(unit.nanometer).copy() - pos_b_np[6:9, 0] += 0.05 - pos_b = pos_b_np * unit.nanometer - - _, ctx = _build_nonpbc_mixed(potential, num_mm_waters=2, embedding="electrostatic") - ctx.setPositions(pos_a) - _, f_a = _energy_and_forces(ctx) - ctx.setPositions(pos_b) - _, f_b = _energy_and_forces(ctx) - - # ML forces (atoms 0..2) must have changed: the ML region polarization - # responds to the shifted MM field. - delta_ml = f_b[:3] - f_a[:3] - assert np.linalg.norm(delta_ml) > 1e-3 From 66903284b6ca075ddc3a2851407054e29ffe9a95 Mon Sep 17 00:00:00 2001 From: CheukHinHoJerry Date: Thu, 17 Sep 2026 12:31:22 +0000 Subject: [PATCH 20/20] clean up --- openmmml/models/macepotential.py | 118 ++++++++++++++----------------- 1 file changed, 52 insertions(+), 66 deletions(-) diff --git a/openmmml/models/macepotential.py b/openmmml/models/macepotential.py index cdafc8f..71524ee 100644 --- a/openmmml/models/macepotential.py +++ b/openmmml/models/macepotential.py @@ -384,10 +384,7 @@ def addForces( ``createMixedSystem``; there is normally no reason to pass it here directly. ``mechanical`` (the default) does not pass MM positions or charges into MACE. ``electrostatic`` passes them into PolarMACE - and scatters the returned ``mm_forces`` back onto the MM atoms; the - removal of the classical ML-MM Coulomb that this assumes is done by - ``createMixedSystem``, not here. It is an error to request it for a - model that cannot accept MM charges. + and scatters the returned ``mm_forces`` back onto the MM atoms. """ import torch try: @@ -438,10 +435,7 @@ def addForces( embeddingData = _prepareMMEmbedding(system, atoms, customNonbondedChargeParameter) mmIndices = embeddingData["mm_atoms"] mmCharges = embeddingData["mm_charges"] - # The electrostatic path needs FULL-system positions inside the callback: it reads MM - # coordinates and returns MM back-reaction forces. PythonForce.setParticles() would hand - # the callback only the ML atoms, so that path keeps the explicit index slice/scatter and - # does not call setParticles(). The plain ML path uses upstream's restriction instead. + mlIndices = np.array(atoms) if (atoms is not None and mmIndices is not None) else None periodic = (topology.getPeriodicBoxVectors() is not None) or system.usesPeriodicBoundaryConditions() @@ -496,12 +490,9 @@ def createMixedSystem(self, needed, ``customNonbondedChargeParameter``. """ - # Internal invariant: the framework only calls this method for names - # returned by getSupportedEmbeddings(); this is not user-facing validation. if embedding != "electrostatic": raise ValueError(f"Unsupported embedding type: {embedding}") - # Interpolation of the model's MM electrostatics is not implemented. if interpolate: raise ValueError("Electrostatic embedding does not support interpolation.") @@ -509,8 +500,7 @@ def createMixedSystem(self, model, device = self._loadModel(args) _validateMMEmbedding(model, atoms, embedding) - # Validate the force-field setup before creating the modified system. - # Multiple charge sources would make the MM charges ambiguous. + # Validate the force-field setup nonbondedForces = [f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)] if len(nonbondedForces) > 1: raise ValueError( @@ -518,7 +508,7 @@ def createMixedSystem(self, "requires exactly one." ) - # The callback uses fixed charges and cannot follow charge offsets. + # Electrostatic embedding does not support following charge offsets. for force in nonbondedForces: for index in range(force.getNumParticleParameterOffsets()): if force.getParticleParameterOffset(index)[2] != 0: @@ -607,11 +597,6 @@ def _supportsMMEmbedding(model) -> bool: def _validateMMEmbedding(model, atoms: Optional[Iterable[int]], embedding: str) -> None: - """Raise if this model and atom selection cannot provide the requested embedding. - - Only "electrostatic" is MACE-specific; "mechanical" is handled by the generic embedding - plugin and needs nothing from us. - """ if embedding == "mechanical": return if embedding != "electrostatic": @@ -631,74 +616,75 @@ def _validateMMEmbedding(model, atoms: Optional[Iterable[int]], embedding: str) ) -def _prepareMMEmbedding(system: openmm.System, atoms: Optional[Iterable[int]], - customNonbondedChargeParameter: Optional[str] = None): - """Extract the MM complement and its charges from the system's NonbondedForce.""" +def _prepareMMEmbedding( + system: openmm.System, + atoms: Optional[Iterable[int]], + customNonbondedChargeParameter: Optional[str] = None, +): + """Extract MM atoms and charges from a standard or custom nonbonded force.""" if atoms is None: return None - numParticles = int(system.getNumParticles()) mlAtoms = np.asarray(list(atoms), dtype=np.int64) - mlSet = set(int(i) for i in mlAtoms.tolist()) + mlSet = set(mlAtoms) mmAtoms = np.asarray( - [i for i in range(numParticles) if i not in mlSet], dtype=np.int64 + [i for i in range(system.getNumParticles()) if i not in mlSet], + dtype=np.int64, ) - # The charges given to the model have to come from wherever the force field - # actually keeps them, which is the same force createMixedSystem zeroed the - # ML charges in. When that is a CustomNonbondedForce the caller has named - # the parameter holding them. - mmCharges = np.empty(len(mmAtoms), dtype=np.float64) - if customNonbondedChargeParameter is not None: - custom = None for force in system.getForces(): - if isinstance(force, openmm.CustomNonbondedForce): - names = [force.getPerParticleParameterName(i) - for i in range(force.getNumPerParticleParameters())] - if customNonbondedChargeParameter in names: - custom = force - chargeIndex = names.index(customNonbondedChargeParameter) - break - if custom is None: - raise ValueError(f"No CustomNonbondedForce defines a per-particle parameter '{customNonbondedChargeParameter}'.") - for row, atomIndex in enumerate(mmAtoms): - mmCharges[row] = custom.getParticleParameters(int(atomIndex))[chargeIndex] - return { - "ml_atoms": mlAtoms, - "mm_atoms": mmAtoms, - "mm_charges": mmCharges, - } - - nonbonded = None - for force in system.getForces(): - if isinstance(force, openmm.NonbondedForce): - nonbonded = force - break - if nonbonded is None: - raise ValueError( - "PolarMACE MM embedding requires a NonbondedForce to source MM charges." + if not isinstance(force, openmm.CustomNonbondedForce): + continue + names = [ + force.getPerParticleParameterName(i) + for i in range(force.getNumPerParticleParameters()) + ] + if customNonbondedChargeParameter in names: + chargeIndex = names.index(customNonbondedChargeParameter) + break + else: + raise ValueError( + "No CustomNonbondedForce defines a per-particle parameter " + f"{customNonbondedChargeParameter!r}." + ) + + def getCharge(i): + return force.getParticleParameters(i)[chargeIndex] + + else: + force = next( + (f for f in system.getForces() + if isinstance(f, openmm.NonbondedForce)), + None, ) + if force is None: + raise ValueError( + "PolarMACE MM embedding requires a NonbondedForce " + "to source MM charges." + ) - for row, atomIndex in enumerate(mmAtoms): - charge, _, _ = nonbonded.getParticleParameters(int(atomIndex)) - mmCharges[row] = charge.value_in_unit(unit.elementary_charge) + def getCharge(i): + return force.getParticleParameters(i)[0].value_in_unit( + unit.elementary_charge + ) return { "ml_atoms": mlAtoms, "mm_atoms": mmAtoms, - "mm_charges": mmCharges, + "mm_charges": np.asarray( + [getCharge(int(i)) for i in mmAtoms], + dtype=np.float64, + ), } -def _computeMACE(state, model, ptr, nodeAttrs, batch, pbc, returnEnergyType, charge, +def _computeMACE(state, model, ptr, node_attrs, batch, pbc, returnEnergyType, charge, multiplicity, periodic, mlIndices=None, mmIndices=None, mmCharges=None): import torch from mace.data.neighborhood import get_neighborhood energyScale = 96.4853 lengthScale = 10.0 - # With setParticles() (mlIndices is None) the state already holds only the ML atoms; on the - # electrostatic path it holds the whole system and the ML subset is sliced out here. positionsFull = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom) numAtoms = positionsFull.shape[0] positions = positionsFull if mlIndices is None else positionsFull[mlIndices] @@ -706,7 +692,7 @@ def _computeMACE(state, model, ptr, nodeAttrs, batch, pbc, returnEnergyType, cha cell = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom) else: cell = np.identity(3, dtype=np.float64) - dtype = nodeAttrs.dtype + dtype = node_attrs.dtype cutoff = float(model.r_max.detach()) edgeIndex, shifts, _, _ = get_neighborhood(positions, cutoff, [periodic, periodic, periodic], cell) cellTensor = torch.tensor(cell, dtype=dtype, device=ptr.device) @@ -717,7 +703,7 @@ def _computeMACE(state, model, ptr, nodeAttrs, batch, pbc, returnEnergyType, cha rcell = torch.zeros((3, 3), dtype=dtype, device=ptr.device) inputDict = { "ptr": ptr, - "node_attrs": nodeAttrs, + "node_attrs": node_attrs, "batch": batch, "pbc": pbc, "positions": torch.tensor(positions, dtype=dtype, device=ptr.device),