[WIP] EMLE embedding

.github/workflows/CI.yml

@@ -16,7 +16,7 @@ jobs:
       fail-fast: false
       matrix:
         configuration: ["cpu", "cuda"]
-        requirements: ["mace-aimnet-torchmd", "nequip", "fennix", "orb"]
+        requirements: ["mace-aimnet-torchmd", "nequip", "fennix", "orb", "emle"]
         include:
           - configuration: "cpu"
             runner: "ubuntu-latest"
@@ -34,6 +34,8 @@ jobs:
             python-version: "3.14"
           - requirements: "orb"
             python-version: "3.12"
+          - requirements: "emle"
+            python-version: "3.14"
 
     steps:
       - uses: actions/checkout@v2

devtools/requirements/emle.txt

@@ -0,0 +1,7 @@
+pytest
+pygit2
+mace-torch
+loguru
+torchani==2.2.4
+setuptools<81
+git+https://github.com/chemle/emle-engine.git@6cb9bbd14008788ae83bd9d39a27b11e0e3407f7

doc/userguide.md

@@ -409,6 +409,27 @@ To avoid double-counting bonded interactions between the MM force field and ML p
 - 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.
 
+### EMLE
+
+OpenMM-ML supports the EMLE (Electrostatic Machine Learning Embedding) method through an interface with the
+[EMLE-Engine](https://chemle.github.io/emle-engine/) package.  The following embedding names are supported.
+
+| Name | Model |
+| --- | --- |
+| `emle` | Default pretrained [EMLE](https://github.com/chemle/emle-models) model. |
+| `emle-engine` | Use a custom EMLE-Engine model loaded from a local file. |
+
+The following extra keyword arguments are recognized by the embedding:
+
+| Name | Model |
+| --- | --- |
+| `embeddingModelPath` | Path to a local model, only used (and required) if `emle-engine` is given as the embedding name. |
+| `alphaMode` | The mode for computing atomic polarizabilities: can be `'species'` (default) or `'reference'`.  See the [EMLE-Engine API documentation](https://chemle.github.io/emle-engine/api/index_models.html#emle.models.EMLE) for more details. |
+| `cutoffDistance` | The cutoff distance for EMLE.  Must be an `openmm.unit.Quantity` with distance units.  Beyond this distance from the ML region, MM atoms will not have an effect.  The default value is 0.75 nm. |
+| `switchingDistance` | The switching distance for EMLE.  Must be an `openmm.unit.Quantity` with distance units.  Between this distance and the cutoff distance from the ML region, the effect of MM atoms will go smoothly to zero.  The default value is 0.6 nm. |
+| `precision` | `'single'` for single precision or `'double'` for double precision. |
+| `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. |
+
 ## Other Packages
 
 OpenMM-ML is based on a plugin architecture, allowing other packages to provide their own interfaces to it.  The

openmmml/embeddings/__init__.py

@@ -1 +1 @@
-from . import mechanicalembedding
+from . import mechanicalembedding, emleembedding

openmmml/embeddings/emleembedding.py

@@ -0,0 +1,225 @@
+"""
+emleembedding.py: Provides the EMLE (Electrostatic Machine Learning Embedding)
+method as an embedding plugin.
+
+This is part of the OpenMM molecular simulation toolkit originating from
+Simbios, the NIH National Center for Physics-Based Simulation of
+Biological Structures at Stanford, funded under the NIH Roadmap for
+Medical Research, grant U54 GM072970. See https://simtk.org.
+
+Portions copyright (c) 2026 Stanford University and the Authors.
+Authors: Evan Pretti
+Contributors:
+
+Permission is hereby granted, free of charge, to any person obtaining a
+copy of this software and associated documentation files (the "Software"),
+to deal in the Software without restriction, including without limitation
+the rights to use, copy, modify, merge, publish, distribute, sublicense,
+and/or sell copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+
+from openmmml.mlpotential import MLPotentialImpl, Embedding, EmbeddingFactory
+from openmmml.embeddings import utilities
+import openmm
+import openmm.app
+import openmm.unit as unit
+import numpy as np
+from functools import partial
+
+class EMLEEmbeddingFactory(EmbeddingFactory):
+    """This is the factory that creates EMLEEmbedding objects."""
+
+    def createEmbedding(self, name: str, **args) -> Embedding:
+        return EMLEEmbedding(name)
+
+
+class EMLEEmbedding(Embedding):
+    """EMLE (Electrostatic Machine Learning Embedding).  This embedding method
+    can be used with any ML potential to perform ML/MM simulations.  The ML-ML
+    interactions will be computed with the ML potential of choice, the MM-MM
+    interactions with a conventional force field, and the remaining interactions
+    with EMLE.
+    """
+
+    def __init__(self, name: str):
+        """
+        Initialize the EMLEEmbedding.
+
+        Parameters
+        ----------
+        name : str
+            The name of the EMLE model.  `emle` (the only builtin model) uses
+            the standard EMLE model.  `emle-engine` allows you to use a custom
+            model supported by the EMLE library.
+        """
+
+        self.name = name
+
+    def createMixedSystem(self,
+                          potential: MLPotentialImpl,
+                          topology: openmm.app.Topology,
+                          system: openmm.System,
+                          atoms: list[int],
+                          forceGroup: int,
+                          interpolate: bool,
+                          **args) -> openmm.System:
+
+        # Make sure that we can import the EMLE library.
+
+        try:
+            from emle.models import EMLE
+        except ImportError:
+            raise ImportError("Failed to import emle-engine: for installation instructions, visit https://github.com/chemle/emle-engine")
+        import torch
+
+        # Get the model path to pass to EMLE.
+
+        if self.name == "emle":
+            modelPath = None
+        elif self.name == "emle-engine":
+            try:
+                modelPath = args["embeddingModelPath"]
+            except KeyError:
+                raise ValueError("For the emle-engine embedding method, an embeddingModelPath must be provided")
+        else:
+            raise ValueError(f"Unrecognized embedding name {self.name!r} for EMLE (recognized names are emle, emle-engine)")
+
+        # Extract additional options to pass to EMLE.
+
+        precision = args.get("precision", None)
+        alphaMode = args.get("alphaMode", "species")
+        cutoffDistance = args.get("cutoffDistance", 0.75 * unit.nanometer)
+        switchingDistance = args.get("switchingDistance", 0.6 * unit.nanometer)
+
+        # Create the new system with ML-ML interactions to be computed by the ML
+        # potential removed.
+
+        periodic = system.usesPeriodicBoundaryConditions()
+        newSystem = utilities.removeBonds(system, topology, atoms, True)
+        numAtoms = newSystem.getNumParticles()
+
+        allCharges = [0.0] * numAtoms
+        for force in newSystem.getForces():
+            if isinstance(force, openmm.NonbondedForce):
+                # Get charges on all particles.
+
+                for atom in range(numAtoms):
+                    charge, _, _ = force.getParticleParameters(atom)
+                    allCharges[atom] += charge.value_in_unit(unit.elementary_charge)
+
+                # All of the LJ interactions in the ML region should be zeroed.
+                # The ML-ML and ML-MM electrostatics should both be zeroed.
+
+                for atom in atoms:
+                    _, sigma, epsilon = force.getParticleParameters(atom)
+                    force.setParticleParameters(atom, 0, sigma, epsilon)
+
+                for iAtom1 in range(len(atoms)):
+                    for iAtom2 in range(iAtom1):
+                        force.addException(atoms[iAtom1], atoms[iAtom2], 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):
+                utilities.makeCustomNonbondedExclusions(force, atoms)
+
+        # Create a PythonForce to compute the EMLE interaction.
+
+        device = potential._getTorchDevice(args)
+
+        if precision is None:
+            # This is the default used by the EMLE library if None is given.
+            dtype = torch.get_default_dtype()
+        elif precision == "single":
+            dtype = torch.float32
+        elif precision == "double":
+            dtype = torch.float64
+        else:
+            raise ValueError(f"Unsupported precision {precision} for the embedding. Supported values are 'single' and 'double'.")
+
+        mlAtomSet = set(atoms)
+        mmAtomList = sorted(set(range(numAtoms)) - mlAtomSet)
+        topologyAtoms = list(topology.atoms())
+
+        atomicNumbers = torch.tensor([topologyAtoms[atom].element.atomic_number for atom in atoms], device=device, dtype=int)
+        mlCharge = sum(allCharges[atom] for atom in atoms)
+        mmCharges = torch.tensor([allCharges[atom] for atom in mmAtomList], device=device)
+        mlIndices = torch.tensor(atoms, device=device, dtype=int)
+        mmIndices = torch.tensor(mmAtomList, device=device, dtype=int)
+
+        # Trying to pass a float charge to EMLE will give an error.
+
+        mlChargeRounded = round(mlCharge)
+        if not np.isclose(mlChargeRounded, mlCharge):
+            raise ValueError(f"Non-integer charge on the ML region {mlCharge} unsupported by EMLE")
+
+        emleCutoff = cutoffDistance.value_in_unit(unit.angstrom)
+        emleSwitchWidth = 1 - switchingDistance / cutoffDistance
+        if not 0 < emleSwitchWidth < 1:
+            raise ValueError("Switching distance must be between 0 and cutoff distance")
+
+        model = EMLE(model=modelPath, method="electrostatic", alpha_mode=alphaMode, cutoff=emleCutoff, switch_width=emleSwitchWidth, device=device, dtype=dtype)
+        energyScale = (1.0 * unit.hartree / unit.item).value_in_unit(unit.kilojoule_per_mole)
+        emleForce = openmm.PythonForce(partial(
+            _computeEMLE,
+            atomicNumbers=atomicNumbers,
+            mlCharge=mlChargeRounded,
+            mmCharges=mmCharges,
+            mlIndices=mlIndices,
+            mmIndices=mmIndices,
+            periodic=periodic,
+            device=device,
+            dtype=dtype,
+            model=model,
+            energyScale=energyScale,
+        ))
+        emleForce.setForceGroup(forceGroup)
+        emleForce.setUsesPeriodicBoundaryConditions(periodic)
+
+        if interpolate:
+            interpolator = utilities.InterpolationHelper()
+            interpolator.addMLPotentialTerms(potential, topology, atoms, forceGroup, **args)
+            interpolator.addMLTerm(emleForce)
+            interpolator.addMMBondedTerms(system, topology, atoms)
+            interpolator.setupNonbonded(newSystem, system)
+            interpolator.setupInterpolation(newSystem)
+
+        else:
+            potential.addForces(topology, newSystem, atoms, forceGroup, **args)
+            newSystem.addForce(emleForce)
+
+        return newSystem
+
+def _computeEMLE(state, atomicNumbers, mlCharge, mmCharges, mlIndices, mmIndices, periodic, device, dtype, model, energyScale):
+    import torch
+
+    positions = state.getPositions(asNumpy=True).value_in_unit(unit.angstrom)
+    positionsTensor = torch.tensor(positions, dtype=dtype, device=device, requires_grad=True)
+
+    if periodic:
+        cell = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom)
+        cellTensor = torch.tensor(cell, dtype=dtype, device=device)
+    else:
+        cellTensor = None
+
+    energy = energyScale * model(atomicNumbers, mmCharges, positionsTensor[mlIndices], positionsTensor[mmIndices], cellTensor, mlCharge, preprocess=True, use_switching_function=True)
+    energy = energy.sum()
+    # For unknown reasons, retain_graph=True appears necessary when calling EMLE
+    # even though we are only calling backward() one time.
+    energy.backward(retain_graph=True)
+    return energy.item(), (-positionsTensor.grad).numpy(force=True)

setup.py

@@ -84,6 +84,8 @@
         ],
         'openmmml.embeddings': [
             'mechanical = openmmml.embeddings.mechanicalembedding:MechanicalEmbeddingFactory',
+            'emle-engine = openmmml.embeddings.emleembedding:EMLEEmbeddingFactory',
+            'emle = openmmml.embeddings.emleembedding:EMLEEmbeddingFactory',
         ]
     }
 )

test/TestEMLEEmbedding.py

@@ -0,0 +1,52 @@
+import numpy as np
+import openmm
+import openmm.app
+import os
+import pytest
+
+from openmmml import MLPotential
+
+emle = pytest.importorskip("emle", reason="emle is not installed")
+mace = pytest.importorskip("mace", reason="mace is not installed")
+platform_ints = range(openmm.Platform.getNumPlatforms())
+# Get the path to the test data
+test_data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+
+# TODO: EMLE energy values have a lot of numerical noise.  Is this expected?
+atol = 0.06
+
+@pytest.mark.parametrize("platform_int", list(platform_ints))
+class TestEMLEEmbedding:
+
+    @pytest.mark.parametrize("interpolate", (False, True))
+    def testEmbedding(self, platform_int, interpolate):
+        pdb = openmm.app.PDBFile(os.path.join(test_data_dir, "alanine-dipeptide", "alanine-dipeptide-explicit.pdb"))
+
+        subset = [atom.index for atom in pdb.topology.atoms() if atom.residue.chain.index == 0]
+
+        mm_force_field = openmm.app.ForceField("amber19-all.xml", "amber19/tip3pfb.xml")
+        ml_potential = MLPotential("mace-off23-small")
+        mm_system = mm_force_field.createSystem(pdb.topology, nonbondedMethod=openmm.app.PME)
+        mixed_system = ml_potential.createMixedSystem(pdb.topology, mm_system, subset, embedding="emle", interpolate=interpolate)
+
+        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)
+
+        mm_context.setPositions(pdb.positions)
+        mixed_context.setPositions(pdb.positions)
+
+        mm_energy = mm_context.getState(energy=True).getPotentialEnergy().value_in_unit(openmm.unit.kilojoule_per_mole)
+
+        # Reference energies are computed with EMLECalculator from EMLE-Engine
+        expected_energy = -33863.30429558904
+
+        if interpolate:
+            for lambda_value in (0.0, 0.25, 0.5, 0.75, 1.0):
+                mixed_context.setParameter("lambda_interpolate", lambda_value)
+                mixed_energy = mixed_context.getState(energy=True).getPotentialEnergy().value_in_unit(openmm.unit.kilojoule_per_mole)
+                assert np.isclose(mixed_energy, expected_energy * lambda_value + mm_energy * (1 - lambda_value), rtol=0, atol=atol)
+
+        else:
+            mixed_energy = mixed_context.getState(energy=True).getPotentialEnergy().value_in_unit(openmm.unit.kilojoule_per_mole)
+            assert np.isclose(mixed_energy, expected_energy, rtol=0, atol=atol)