[WIP] EMLE embedding by epretti · Pull Request #159 · openmm/openmm-ml

epretti · 2026-07-16T00:17:45Z

Adds support for EMLE embedding.

(Work in progress, requires #158 and needs updates, review, more testing, and possibly an EMLE-Engine release.)

…rostatic path Upstream openmm#167 restricts PythonForce to the ML atoms via setParticles(), so the callback state no longer carries MM coordinates. PolarMACE electrostatic embedding needs them in the same forward pass (the MM field polarises the ML density) and returns MM back-reaction forces, so that path keeps the explicit index slice/scatter and does not call setParticles(). The plain ML path uses upstream's restriction unchanged; with indices=None and mmInfo=None _computeMACE is behaviourally identical to upstream. Same pattern as the EMLE embedding PR (openmm#159): a full-system PythonForce that indexes ML and MM regions itself. Verified on OpenMM 8.6.1 (isolated env) with graph_longrange@external_field: TestPolarMACEElectrostaticEmbedding + TestElectrostaticEmbeddingPME + TestMechanicalEmbedding: 104 passed, 6 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…ents Style alignment with openmmml/embeddings/emleembedding.py (upstream PR openmm#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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN

JMorado · 2026-09-21T09:30:23Z

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


Hopefully this won't be needed once we release the base emle-engine package on conda-forge.

JMorado · 2026-09-21T09:33:00Z

+| 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. |


Note that alphaMode will change in the next release. 'species' will become 'fixed', and 'reference' will become 'flexible'. See https://github.com/chemle/emle-engine/blob/devel/emle/models/_emle.py#L116-L122

I'll update this documentation, and also change the code to not pass the keyword argument to EMLE if the user doesn't specify it; this way, the default won't be hard-coded in OpenMM-ML.

JMorado · 2026-09-21T09:35:44Z

+| `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. |
+


Do we also want to support the other embedding schemes available in emle-engine? Please see: https://github.com/chemle/emle-engine/blob/devel/emle/models/_emle.py#L101-L114. They have been used in some recent papers.

I wasn't sure if there were benefits to exposing anything other than electrostatic, but if people are using the others, we can support them.

JMorado · 2026-09-21T09:41:31Z

+        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)")


I'd probably suggest simplifying this bit and leaving self.name always as emle. A modelPath can be optionally accepted if the user wants to use a bespoke EMLE model. I think this is closer to the logic used for example in the mace implementation.

In MACE we have mace to load a custom model and mace-... for various pretrained models. I was thinking along these lines where emle-engine would be used to load a custom model with the EMLE code, while emle would give the default pretrained EMLE model, and if there were a future pretrained model like EMLE2 or whatever, it could have another name. But if that doesn't make sense for EMLE, e.g., you're planning on exposing any future models only through new method keyword arguments, etc., I could make it just emle.

JMorado · 2026-09-21T09:49:09Z

+        cell = state.getPeriodicBoxVectors(asNumpy=True).value_in_unit(unit.angstrom)
+        cellTensor = torch.tensor(cell, dtype=dtype, device=device)
+    else:
+        cellTensor = None


See chemle/emle-engine#91 for the support we discussed for nonperiodic systems.

JMorado · 2026-09-21T09:51:09Z

Many thanks @epretti, looks very good! I'm also pinging @lohedges in cases he has any further feedback.

JMorado · 2026-09-21T10:10:43Z

+# TODO: EMLE energy values have a lot of numerical noise.  Is this expected?
+atol = 0.06


I might have mentioned this to you already but I think 0.05 kJ/mol is about the single precision limit for Hartree energies, so it might be worth running the tests in double precision.

JMorado · 2026-09-21T10:18:50Z

+            'emle-engine = openmmml.embeddings.emleembedding:EMLEEmbeddingFactory',
+            'emle = openmmml.embeddings.emleembedding:EMLEEmbeddingFactory',


This could be made just one as I mentioned previously.

JMorado · 2026-09-21T11:17:26Z

+# TODO: This script should ideally compute reference energies for EMLE without
+# using OpenMM-ML or calling the EMLE model directly.


See the attached zip file for an updated version of this file that uses Sire-EMLE. It also contains the output, which I have confirmed matches the current expected_energy value in TestEMLEEmbedding.py

data_emle.zip

Thanks! Is there any chance of doing periodic and non-periodic cases (the latter using the same system but with NoCutoff or CutoffNonPeriodic as the nonbondedMethod)? I had a non-periodic test also at first, then removed it when the support for the EMLE preprocessing was added. Now that the preprocessing supports periodic and non-periodic systems, I might want to add a test case back for it if possible.

epretti · 2026-09-22T22:28:55Z

@JMorado Regarding your comment from our other conversation:

I chose and documented default values for the cutoff and switching distances if none are specified. These come from the cutoff in the Sire-EMLE example and default switching distance fraction (I set up the OpenMM-ML interface so that the switching distance is given explicitly since this is more consistent with OpenMM). If you think it is better to not specify defaults, such that users must specify their own values, let me know.

I think the default should be the same as that used for the NonbondedForce (or CustomNonbondeForce in the system). Would it make sense to use getCutoffDistance and getSwitchingDistance to determine these? This cutoff will be used for the ML-MM electrostatic interaction, so to me it makes sense that it shares the same values as the remainder of the nonbonded interactions in the system.]

This would make sense if we could guarantee that they would have meaningful values that actually are in use, but that is not always the case. The cutoff distance is normally in use unless a user is simulating a non-periodic system and has specified NoCutoff. We could forbid this or just require a user to explicitly specify a cutoff for EMLE in this case. The switching distance will also always be set, but often not to a meaningful value, since switching is turned off by default in NonbondedForce unless a user asks for it. I would be inclined to argue that the defaults, if we are going to provide them, should be whatever was in use when EMLE was trained. But I’ll defer to you as more knowledgeable on the best choice here.

As for the non-periodic preprocessing option, I’ve seen that this was merged into the “devel” branch.

Which branch is the release going to come from? I’ve been trying to test against the “devel” branch now to try this option out. I found using this branch that the EMLE model object can no longer be pickled, which breaks the OpenMM-ML interpolation feature, since it needs to clone its PythonForce which involves serializing and deserializing it. I don’t know if this was intentional or not.

epretti added 5 commits July 15, 2026 16:22

Refactor embedding and interpolation routines

eb198b4

EMLE implementation

09a327e

Debug CI

ec4bb0e

Debug CI

1b47c4c

Merge branch 'main' into emle

294aea0

CheukHinHoJerry mentioned this pull request Sep 17, 2026

Add PolarMACE electrostatic ML/MM embedding via graph-longrange #169

Draft

epretti added 3 commits September 18, 2026 09:39

Merge branch 'main' into emle

66c2492

Update and try to get CI working with pip install only

876bad8

Add EMLE preprocessing (with cutoff and switching function)

fc95e6a

JMorado reviewed Sep 21, 2026

View reviewed changes

		# TODO: EMLE energy values have a lot of numerical noise. Is this expected?
		atol = 0.06

		'emle-engine = openmmml.embeddings.emleembedding:EMLEEmbeddingFactory',
		'emle = openmmml.embeddings.emleembedding:EMLEEmbeddingFactory',

		# TODO: This script should ideally compute reference energies for EMLE without
		# using OpenMM-ML or calling the EMLE model directly.

Conversation

epretti commented Jul 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JMorado commented Sep 21, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

epretti commented Sep 22, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants