fix(ml/mm): do ONIOM by HaoZeke · Pull Request #11 · metatensor/gromacs

HaoZeke · 2026-03-22T13:19:02Z

This GitHub repository is only used as a mirror of https://gitlab.com/gromacs/gromacs/.

Please, do not create pull requests on GitHub.

If you want to contribute, please take a look at our Developer Guide:

https://manual.gromacs.org/current/dev-manual/contribute.html

Much more stable now. Basically when in thread MPI, PyTorch must match the same thread count, without it the conserved energy drift for *nompi* is insane.

Now we have a drfit of around 0.5 kJ/mol over 2000 steps / 1 ps. MPI build was always fine, drift of 0.08 kJ/mol over 500 steps...

The right fix for lab-cosmo/atomistic-cookbook#240

Also for ONIOM, the old code removed angles with >= 2 ML atoms and dihedrals with >= 3 ML atoms, when it should only remove interactions where ALL atoms are in the ML region. Boundary interactions (involving both ML and MM atoms) must stay in the MM force field since the ML model can't see the MM atoms. For the ACE-ALA-NME test system this makes no difference (all internal interactions have all atoms as ML), but it matters for systems with intra-molecular ML/MM boundaries.

Basically we do (E = E_MM(full) + E_ML(solute) - E_MM(solute)), where E_MM(solute) is subtracted via topology surgery (removing bonded/NB terms within the ML region). This gives a well-defined total energy with conservative forces (via autograd backward pass).

~~Still need to handle link atoms and electrostatic embedding in a followup, but for systems with explicit capping groups (ACE/NME) this is sufficient.~~

The previous topology surgery removed too many interactions at the ML/MM boundary, which is inconsistent with the ONIOM subtractive scheme: E_ONIOM = E_MM(real) + E_ML(model) - E_MM(model) where E_MM(model) only includes interactions where ALL atoms are in the ML subsystem. Boundary interactions (involving both ML and MM atoms) should remain in E_MM(real) and are NOT double-counted because the ML model cannot see the MM atoms. Changes to removal thresholds: - Angles: remove only if ALL 3 atoms are ML (was >= 2). Boundary angles (2 ML + 1 MM) are now kept in E_MM(real). - Dihedrals: remove only if ALL 4 atoms are ML (was >= 3). Boundary dihedrals (3 ML + 1 MM, 2 ML + 2 MM) are now kept. - Bonds: unchanged (both atoms must be ML, already correct). - NB exclusions: unchanged (ML-ML pairs excluded, already correct). - SETTLE: unchanged (removed if >= 2 ML atoms, water constraint). Reference: Chung et al., Chem. Rev. 115, 5678 (2015).

Thread-MPI builds with PyTorch metatomic produce incorrect forces when PyTorch's default thread count (all cores) conflicts with thread-MPI's internal threading. This causes simulations to blow up within ~100 steps (conserved energy drift > 2000 kJ/mol) while identical real-MPI builds are perfectly stable (drift 0.08 kJ/mol). Fix: unconditionally set at::set_num_threads to match GROMACS ntomp. This replaces the previous conditional logic that left thread count at the PyTorch default for single-rank runs. The ntomp value from gmx_omp_nthreads_get is always correct: it reflects the user's -ntomp setting or the auto-detected value.

Infrastructure for ONIOM link atoms and electrostatic embedding: - New MDP options: metatomic-link-atoms (bool), metatomic-electrostatic-embedding (bool) - When link-atoms=yes: buildLinkFrontier() runs during preprocessing, link frontier atoms stored in MetatomicParameters - When electrostatic-embedding=yes: MM point charges collected from topology before preprocessing and stored for per-step use - Both default to no for backward compatibility The force provider changes (injecting link atoms into the System, passing charges via add_data, and redistributing link atom forces) follow in the next commit.

PicoCentauri · 2026-03-22T14:15:26Z

    section.addOption(StringOption(VARIANT_NC_FORCES_TAG.c_str()).store(&params_.variantNcForces));
    section.addOption(StringOption(VARIANT_NC_STRESS_TAG.c_str()).store(&params_.variantNcStress));
+    section.addOption(BooleanOption(LINK_ATOMS_TAG.c_str()).store(&params_.linkAtoms));
+    section.addOption(BooleanOption(ELECTROSTATIC_EMBEDDING_TAG.c_str()).store(&params_.electrostaticEmbedding));


This we might check how the nnpot is doing it. I would basically add the charges as a metatomic.System property.

Only if the model requests them as an input!

Replicate the NNPot algorithm for link atoms: boundary MM atoms are included in the metatomic atom set and their positions/types are REPLACED with link atoms each step. This is fundamentally different from the previous broken approach that APPENDED link atoms as extra atoms, causing the ML model to see chemically nonsensical fragments. Algorithm: 1. Preprocessing: buildLinkFrontier() finds boundary bonds, then boundary MM atom indices are added to mtaIndices_ 2. Per-step: gatherAtomPositions() gathers boundary MM positions. Then positions are replaced with link atom positions and atomic numbers overwritten with H (Z=1) 3. Model evaluation: the model sees ML atoms + H caps (at boundary MM slots), not ML atoms + extra random H atoms 4. Force redistribution: link atom forces are split between the embedded ML atom and the real MM atom via spreadForce() 5. Force scatter: all forces (including redistributed) applied to GROMACS atoms normally No explicit E_MM(model) subtraction needed: topology surgery removes internal ML terms, and the link atom replacement handles boundary terms. This matches the NNPot ONIOM implementation.

Luthaf · 2026-03-23T08:15:41Z

Basically when in thread MPI, PyTorch must match the same thread count, without it the conserved energy drift for nompi is insane.

Do you know why this fails otherwise?

Luthaf · 2026-03-23T08:16:44Z

-        {
-            at::set_num_threads(1);
-        }
+        int ntomp = gmx_omp_nthreads_get(ModuleMultiThread::Default);


Why does this no longer use all core on a MPI rank when doing real MPI?

- Threading comment: explain dual OpenMP runtime root cause (libomp from GROMACS + libgomp from PyTorch, independent global state) - Remove metatomic-electrostatic-embedding MDP option (model-driven now, per Luthaf review). Charges always collected; force provider passes them only if the model requests point_charges input. - Improve link atom documentation comments

- Threading: document PyTorch cached-thread-count + MKL synchronization as the reason for unconditional at::set_num_threads (replaces incorrect dual-OMP-runtime explanation) - Remove metatomic-electrostatic-embedding MDP option; charges always collected but only passed if model requests point_charges input - Improve link atom documentation comments

PicoCentauri

Did we know how this affects the thread mpi performance?

PicoCentauri · 2026-05-20T06:21:53Z

+.. mdp:: metatomic-oniom
+
+   (false) Enable the subtractive ONIOM topology preprocessing for the selected
+   Metatomic input group. When enabled, bonded and non-bonded
+   molecular-mechanics terms represented by the model are removed from the
+   classical force-field contribution and the Metatomic model contribution is
+   added through the applied-forces framework. When disabled, the Metatomic
+   contribution is added on top of the unmodified molecular-mechanics topology.
+
+.. mdp:: metatomic-link-atoms
+
+   (false) Enable hydrogen link atoms at bonds crossing the Metatomic/MM
+   boundary. Direct MM neighbors bonded to selected Metatomic atoms are included
+   in the embedded atom set and represented as hydrogen caps during model
+   evaluation. This requires :mdp:`metatomic-oniom` and a model trained for
+   link-atom-capped boundary chemistry. If a boundary MM atom is bonded to
+   multiple selected Metatomic atoms, each cut bond is represented by a distinct
+   hydrogen cap in the model input.


Are the options exclusive?

I wonder why we can't always turn on the ONIOM option. And maybe we should add a link to some deeper infos what ONIOM is.

PicoCentauri · 2026-05-20T06:25:20Z

+
+        // Link atom position replacement INSIDE the autograd graph.
+        // r_link = r_emb + d_link * (r_MM - r_emb) / |r_MM - r_emb|
+        // By computing this with torch operations, autograd automatically


hail autograd

PicoCentauri · 2026-05-20T06:26:49Z

+
+            if (!extraLinkPositions.empty())
+            {
+                strained_positions = torch::cat(


I had some AD problems in the past with cat. But maybe it is fixed...

PicoCentauri · 2026-05-20T06:27:33Z

+Metatomic models can declare additional inputs through their
+``requested_inputs`` metadata. When a model requests ``charges`` or a
+``charges/<variant>`` input, |Gromacs| provides a per-atom Metatomic
+``TensorMap`` with ``quantity = charge`` and ``unit = e``. Charge inputs are
+model-controlled; no :ref:`mdp` option enables or disables them.


We should also link the page where we give all additional inputs that we support.

HaoZeke added 2 commits March 20, 2026 19:34

HaoZeke requested review from Luthaf, PicoCentauri and ceriottm March 22, 2026 13:19

HaoZeke changed the title ~~fix(oniom): setup and fix thread mpi issues~~ fix(qm/mm): setup and fix thread mpi issues Mar 22, 2026

HaoZeke force-pushed the fix/oniom-subtraction branch from 6312b85 to 9f1b24b Compare March 22, 2026 14:03

PicoCentauri reviewed Mar 22, 2026

View reviewed changes

HaoZeke force-pushed the fix/oniom-subtraction branch 11 times, most recently from f7127b5 to b3c3830 Compare March 22, 2026 16:17

HaoZeke force-pushed the fix/oniom-subtraction branch from b3c3830 to 2314cd5 Compare March 22, 2026 16:20

Luthaf reviewed Mar 23, 2026

View reviewed changes

Comment thread src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp Outdated

Luthaf requested changes Mar 23, 2026

View reviewed changes

Comment thread src/gromacs/applied_forces/metatomic/metatomic_options.cpp Outdated

HaoZeke force-pushed the fix/oniom-subtraction branch 2 times, most recently from 5aa837c to 02637c1 Compare March 23, 2026 21:37

HaoZeke force-pushed the fix/oniom-subtraction branch from 02637c1 to 16ba984 Compare March 24, 2026 00:10

HaoZeke added 18 commits May 6, 2026 01:53

fix(mta): provide requested charge inputs

89616ec

merge: update metatomic base

dbb491e

test(mta): align ONIOM topology expectations

4c2bc65

test(threadmpi): wait on Windows thread handles

0674569

fix(mta): map link atoms by global atom index

b21a99b

chore(mta): avoid shadowed model index

2799a41

fix(mta): support multiple link caps per boundary atom

08b4864

fix(mta): use active link distance value

5cafa17

fix(mta): honor PBC shifts for link caps

e4f532e

test(topology): cover shifted link force spread

42270f2

fix(mta): redistribute link cap output forces

eaae9b7

test(topology): use tolerance for link force spread

cb7505c

fix(mta): single-own link cap outputs

8d92223

test(mta): require opt-in ONIOM preprocessing

787a5f9

test(mta): fix ONIOM option test scaffolding

ea93e8a

feat(mta): make ONIOM preprocessing opt in

fa82917

test(mta): update ONIOM option refdata

9ffa3fe

fix(mta): build model pairlist from all plain pairs

d905082

PicoCentauri reviewed May 20, 2026

View reviewed changes

HaoZeke changed the title ~~fix(qm/mm): setup and fix thread mpi issues~~ fix(ml/mm): do ONIOM May 24, 2026

Navigation Menu

Uh oh!

fix(ml/mm): do ONIOM - #11

fix(ml/mm): do ONIOM#11
HaoZeke wants to merge 23 commits into
metatomicfrom
fix/oniom-subtraction

HaoZeke commented Mar 22, 2026 •

edited

Loading

Uh oh!

PicoCentauri Mar 22, 2026

Uh oh!

Luthaf Mar 23, 2026

Uh oh!

Luthaf commented Mar 23, 2026

Uh oh!

Luthaf Mar 23, 2026

Uh oh!

Uh oh!

Uh oh!

PicoCentauri left a comment

Uh oh!

PicoCentauri May 20, 2026

Uh oh!

PicoCentauri May 20, 2026

Uh oh!

PicoCentauri May 20, 2026

Uh oh!

PicoCentauri May 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

HaoZeke commented Mar 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

PicoCentauri Mar 22, 2026

Choose a reason for hiding this comment

Uh oh!

Luthaf Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Luthaf commented Mar 23, 2026

Uh oh!

Luthaf Mar 23, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

PicoCentauri left a comment

Choose a reason for hiding this comment

Uh oh!

PicoCentauri May 20, 2026

Choose a reason for hiding this comment

Uh oh!

PicoCentauri May 20, 2026

Choose a reason for hiding this comment

Uh oh!

PicoCentauri May 20, 2026

Choose a reason for hiding this comment

Uh oh!

PicoCentauri May 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

HaoZeke commented Mar 22, 2026 •

edited

Loading