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

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

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

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.

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.

- 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

modifyTopology scanned bonds for ML/MM boundary atoms using only the molecule block's globalAtomStart. It runs before splitEmbeddedBlocks, so a block can still hold many copies of one molecule type (e.g. 34 lipids), and only the first copy's boundary bonds were found; link atoms for all other molecules were silently missing. Visit every molecule in the block with its own atom offset. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

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

HaoZeke force-pushed the fix/oniom-subtraction branch 12 times, most recently 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 3 times, most recently from 02637c1 to 16ba984 Compare March 24, 2026 00:10

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

PicoCentauri force-pushed the fix/oniom-subtraction branch from 98c23d2 to e6de5ca Compare September 17, 2026 12:44

HaoZeke and others added 21 commits September 18, 2026 13:18

fix(mta): provide requested charge inputs

94a19a9

test(mta): align ONIOM topology expectations

67d2d9f

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

339d0af

chore(mta): avoid shadowed model index

044c927

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

1261e03

fix(mta): use active link distance value

d881737

fix(mta): honor PBC shifts for link caps

b763b32

test(topology): cover shifted link force spread

86bd7d3

fix(mta): redistribute link cap output forces

17f8202

test(topology): use tolerance for link force spread

f84df6f

fix(mta): single-own link cap outputs

28aa9d1

test(mta): require opt-in ONIOM preprocessing

6392e09

test(mta): fix ONIOM option test scaffolding

34a62f8

feat(mta): make ONIOM preprocessing opt in

0141227

test(mta): update ONIOM option refdata

fceab77

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

67ef5a3

PicoCentauri force-pushed the fix/oniom-subtraction branch from e6de5ca to 4154684 Compare September 18, 2026 11:18

Navigation Menu

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

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

HaoZeke commented Mar 22, 2026 •

edited

Loading

Uh oh!

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!

Uh oh!

Uh oh!

PicoCentauri May 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

Uh oh!

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!

Uh oh!

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