Add graph-longrange PolarMACE electrostatic embedding - #27
Open
CheukHinHoJerry wants to merge 22 commits into
Open
CheukHinHoJerry wants to merge 22 commits into
CheukHinHoJerry wants to merge 22 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d4ed71717c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+157
to
+163
| 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 |
There was a problem hiding this comment.
When a covalent ML/MM boundary is used with an implicit-solvent or polarizable MM system, adding a particle only to NonbondedForce and CustomNonbondedForce leaves forces such as CustomGBForce, GBSAOBCForce, or AmoebaMultipoleForce with fewer particle parameter records than the System. OpenMM then rejects the resulting system during Context creation. Either add appropriate parameters for every supported per-particle force or reject unsupported force types before adding the link site.
Useful? React with 👍 / 👎.
| systemList = [system, newSystem] | ||
| else: | ||
| systemList = [newSystem] | ||
| capIndices, oldToNew = utilities.addLinkAtomSites(newTopology, systemList, linkBonds, args.get("linkAtomDistances", [])) |
There was a problem hiding this comment.
For periodic mechanical embedding with a long-range ML model and at least one bond crossing the ML/MM boundary, excludeForce is created with the original particle count, while this call adds link-site particles only to newSystem (and optionally system). The excludeForce is subsequently installed as a CustomCVForce collective variable, whose internal NonbondedForce no longer matches the containing system's particle count, so Context construction fails. Add a zero-charge particle to excludeForce for every generated cap as well.
Useful? React with 👍 / 👎.
CheukHinHoJerry
force-pushed
the
feat/graph-longrange-polarmace-embedding
branch
from
September 10, 2026 12:10
d4ed717 to
94f7fb1
Compare
CheukHinHoJerry
force-pushed
the
main
branch
from
September 10, 2026 13:31
ee2dde4 to
2296f12
Compare
…polarmace-embedding # Conflicts: # doc/userguide.md # openmmml/models/macepotential.py # test/TestMACEPotential.py # test/TestMechanicalEmbedding.py
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN
…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>
CheukHinHoJerry
changed the base branch from
main
to
feat/polar-mace-mm-embedding
September 17, 2026 01:32
CheukHinHoJerry
changed the base branch from
feat/polar-mace-mm-embedding
to
main
September 17, 2026 01:32
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN
- 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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN
…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
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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011TXm1UZZ1hYRkFHZgVFwAN
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Integrate PolarMACE electrostatic ML/MM embedding with the reusable external-source APIs from graph-longrange.
This PR is one squashed integration commit targeting this fork's updated
main. The fork'smainhas first been synchronized with official OpenMM-ML code through upstream commit2296f12, so upstream synchronization changes are no longer mixed into this PR diff.The sole exception is official upstream's CI workflow update: the current GitHub credential lacks
workflowscope, so.github/workflows/CI.ymlremains at the fork's previous version. No runtime or library code was omitted.Dependency
This PR requires the
external_fieldbranch of graph-longrange, which provides the external-source energy and feature blocks:pip install 'git+https://github.com/WillBaldwin0/graph_electrostatics.git@external_field'devtools/requirements/mace-aimnet-torchmd.txtis pinned to that branch (it previously pinnedv0.4.0, which predates the external-source work and does not provide these blocks). graph-longrange owns the source-to-target electrostatic field andGTOElectrostaticCrossEnergy; OpenMM-ML does not duplicate those electrostatic kernels.Integration
All of this lives in
openmmml/models/macepotential.py: the external-source adapter is installed by_enable_polarmace_external_sources, MM preparation by_prepareMMEmbedding, and the energy/force path by_computeMACE. OpenMM selection, force-field surgery, input preparation, and force scattering are in the same module.Normalization
Verification
Against the pinned
external_fieldbranch (6a86de5, v0.4.4):Result: 10 passed in
TestPolarMACEElectrostaticEmbedding.pyand 18 passed inTestElectrostaticEmbeddingPME.py.Coverage includes nonperiodic and periodic execution, energy changes relative to mechanical embedding, zero-charge equivalence, MM back-reaction, whole-box translation invariance, link-charge schemes, and finite-difference force consistency.
Current limitations
Not yet rebased on upstream
mainUpstream has moved to
a7fb40e(1.8). Upstream PR openmm#167 ("UsesetParticles()to restrict atoms") removes the manualindicesslice/scatter from_computeMACEand restricts the state to the ML atoms instead. This PR's electrostatic path reads full-system positions from the same state (positions_full[mmInfo["mm_atoms"]]), so merginga7fb40econflicts inmacepotential.pyand needs a decision on how the external-source path obtains MM positions oncesetParticles()is in use. Tracked before retargeting upstream.