feat(mta): NNPot-style link atom replacement (correct ONIOM approach)

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.

Author: HaoZeke

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

@@ -164,6 +164,9 @@ struct MetatomicData
     //! Uncertainty threshold in kJ/mol. Atoms above this trigger a warning.
     double uncertaintyThreshold = 0.0;
 
+    //! Link frontier atoms for ONIOM link atom support.
+    std::vector<LinkFrontierAtom> linkFrontier;
+
     //! Non-conservative mode: forces/stress predicted directly, no backward pass.
     bool nonConservative = false;
     //! Output keys for non-conservative forces and stress.
@@ -486,6 +489,16 @@ MetatomicForceProvider::MetatomicForceProvider(const MetatomicOptions& options,
                         data_->nc_stress_key.c_str());
     }
 
+    // Store link frontier from preprocessing
+    data_->linkFrontier = options_.params_.linkFrontier_;
+    if (!data_->linkFrontier.empty())
+    {
+        GMX_LOG(logger_.info)
+                .asParagraph()
+                .appendTextFormatted("Metatomic: %zu link atoms at ML/MM boundary",
+                                     data_->linkFrontier.size());
+    }
+
     GMX_LOG(logger_.info)
             .asParagraph()
             .appendText("MetatomicForceProvider initialization complete.");
@@ -1153,6 +1166,43 @@ void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, F
     }
     copy_mat(inputs.box_, box_);
 
+    // Link atom setup: find MTA indices for link frontier atoms and
+    // overwrite boundary MM atom types to hydrogen.  The position
+    // replacement is done INSIDE the autograd graph (after torch tensor
+    // creation) so that forces are automatically correct via chain rule.
+    if (!data_->linkFrontier.empty())
+    {
+        bool needTypesRebuild = false;
+        // Build a map from GROMACS global atom index -> MTA local model index
+        // for link atom index lookups.  mtaToGmxLocal_ maps model index -> gmx
+        // local index; gmxLocalToMtaIdx_ maps gmx local -> model index.
+        // Link frontier stores GROMACS global indices, so we use gmxLocalToMtaIdx_.
+        for (auto& link : data_->linkFrontier)
+        {
+            int32_t embGmxGlobal = link.getEmbeddedIndex();
+            int32_t mmGmxGlobal  = link.getMMIndex();
+
+            // In serial mode, global == local for the first N atoms
+            int32_t embMtaIdx = (embGmxGlobal < static_cast<int32_t>(gmxLocalToMtaIdx_.size()))
+                                    ? gmxLocalToMtaIdx_[embGmxGlobal] : -1;
+            int32_t mmMtaIdx  = (mmGmxGlobal < static_cast<int32_t>(gmxLocalToMtaIdx_.size()))
+                                    ? gmxLocalToMtaIdx_[mmGmxGlobal] : -1;
+            // Always set input indices (initialize to -1 if not found)
+            link.setInputIndices(embMtaIdx, mmMtaIdx);
+            if (embMtaIdx >= 0 && mmMtaIdx >= 0)
+            {
+                atomNumbers_[mmMtaIdx] = link.linkAtomNumber();  // H = 1
+                needTypesRebuild = true;
+            }
+        }
+        if (needTypesRebuild)
+        {
+            data_->cachedTypes = torch::tensor(
+                    atomNumbers_, torch::TensorOptions().dtype(torch::kInt32))
+                    .to(data_->device);
+        }
+    }
+
     // Newton NL mode: in parallel, each rank needs ALL pairs involving its
     // home atoms (not just the ones assigned by the eighth-shell DD
     // decomposition). Uses the GROMACS pairlist as the pair source, then
@@ -1256,6 +1306,35 @@ void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, F
         auto strained_cell      = torch::matmul(torch_cell, strain);
         auto strained_positions = torch::matmul(torch_positions, strain);
 
+        // 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
+        // computes dE/dr_emb and dE/dr_MM via the chain rule through r_link.
+        // No manual spreadForce redistribution needed.
+        if (!data_->linkFrontier.empty())
+        {
+            for (const auto& link : data_->linkFrontier)
+            {
+                int32_t embIdx = link.getInputIndexEmb();
+                int32_t mmIdx  = link.getInputIndexMM();
+                if (embIdx < 0 || mmIdx < 0
+                    || embIdx >= numLocalMta_ || mmIdx >= numLocalMta_)
+                {
+                    continue;
+                }
+
+                auto r_emb = strained_positions.index({embIdx});
+                auto r_mm  = strained_positions.index({mmIdx});
+                auto direction = r_mm - r_emb;
+                auto dist = direction.norm();
+                auto r_link = r_emb + link.linkDistance() * direction / dist;
+
+                // Replace boundary MM position with link atom position
+                // Using index_put_ keeps the operation in the autograd graph
+                strained_positions.index_put_({mmIdx}, r_link);
+            }
+        }
+
         auto system = torch::make_intrusive<metatomic_torch::SystemHolder>(
                 data_->cachedTypes, strained_positions, strained_cell, data_->cachedPbc);
 
@@ -1608,7 +1687,10 @@ void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, F
     }
     else
     {
-        // Serial: apply forces directly
+        // Apply forces directly. When link atoms are used, autograd has
+        // already computed the correct forces on r_emb and r_MM via the
+        // chain rule through r_link (because the link position computation
+        // is in the autograd graph). No manual spreadForce needed.
         for (int32_t i = 0; i < numLocalMta_; i++)
         {
             int32_t gmxIdx = mtaToGmxLocal_[i];

src/gromacs/applied_forces/metatomic/metatomic_options.cpp

@@ -52,6 +52,8 @@
 #include "gromacs/options/optionsection.h"
 #include "gromacs/selection/indexutil.h"
 #include "gromacs/topology/embedded_system_preprocessing.h"
+#include "gromacs/topology/idef.h"
+#include "gromacs/topology/ifunc.h"
 #include "gromacs/topology/mtop_util.h"
 #include "gromacs/topology/topology.h"
 #include "gromacs/utility/keyvaluetreebuilder.h"
@@ -292,8 +294,74 @@ void MetatomicOptions::modifyTopology(gmx_mtop_t* top)
                                      params_.mmCharges_.size());
     }
 
+    if (params_.linkAtoms)
+    {
+        // NNPot-style: identify boundary MM atoms first (by scanning bonds
+        // between ML and non-ML atoms), add them to the embedded set, THEN
+        // run topology surgery on the expanded set.  This ensures:
+        // - NB exclusions include boundary MM atoms (no double-counting)
+        // - Bonded terms between ML and boundary-MM are properly handled
+        // - buildLinkFrontier finds zero cut bonds (all boundary atoms are embedded)
+        //
+        // The link frontier is built from the ORIGINAL ML set (before expansion)
+        // so we know which embedded atoms are "real ML" vs "boundary MM".
+        std::set<int> origMtaSet(params_.mtaIndices_.begin(), params_.mtaIndices_.end());
+
+        // Scan bonds to find direct MM neighbors of ML atoms
+        std::set<int> boundaryMM;
+        for (size_t mb = 0; mb < top->molblock.size(); ++mb)
+        {
+            const auto& moltype = top->moltype[top->molblock[mb].type];
+            int start = top->moleculeBlockIndices[mb].globalAtomStart;
+
+            for (const auto ftype : gmx::EnumerationWrapper<InteractionFunction>{})
+            {
+                if (!(interaction_function[ftype].flags & IF_CHEMBOND) || NRAL(ftype) != 2
+                    || moltype.ilist[ftype].empty())
+                {
+                    continue;
+                }
+                for (int j = 0; j < moltype.ilist[ftype].size(); j += 3)
+                {
+                    int a1 = moltype.ilist[ftype].iatoms[j + 1] + start;
+                    int a2 = moltype.ilist[ftype].iatoms[j + 2] + start;
+                    bool a1_ml = origMtaSet.count(a1) > 0;
+                    bool a2_ml = origMtaSet.count(a2) > 0;
+                    if (a1_ml && !a2_ml && boundaryMM.count(a2) == 0)
+                    {
+                        boundaryMM.insert(a2);
+                        // Store as LinkFrontierAtom: a1=embedded, a2=MM
+                        params_.linkFrontier_.emplace_back(a1, a2);
+                    }
+                    else if (a2_ml && !a1_ml && boundaryMM.count(a1) == 0)
+                    {
+                        boundaryMM.insert(a1);
+                        params_.linkFrontier_.emplace_back(a2, a1);
+                    }
+                }
+            }
+        }
+
+        // Add boundary MM atoms to the embedded set
+        for (int mmIdx : boundaryMM)
+        {
+            params_.mtaIndices_.push_back(mmIdx);
+        }
+
+        GMX_LOG(logger().info)
+                .appendTextFormatted("Metatomic: expanded embedded set from %zu to %zu atoms "
+                                     "(%zu boundary MM for link atoms)",
+                                     origMtaSet.size(),
+                                     params_.mtaIndices_.size(),
+                                     boundaryMM.size());
+    }
+
+    // Run topology surgery on the (possibly expanded) embedded set
     preprocessTopology(top, params_.mtaIndices_, logger(), wi_,
-                       params_.linkAtoms, &params_.linkFrontier_);
+                       /*buildLinks=*/false, nullptr);
+    // Note: buildLinkFrontier is not called inside preprocessTopology because
+    // we already built it above (and with the expanded set, there are no
+    // cut bonds -- all boundary atoms are now embedded).
 }
 
 void MetatomicOptions::writeParamsToKvt(KeyValueTreeObjectBuilder treeBuilder)
@@ -309,6 +377,18 @@ void MetatomicOptions::writeParamsToKvt(KeyValueTreeObjectBuilder treeBuilder)
     {
         GroupIndexAdder.addValue(indexValue);
     }
+
+    // Serialize link frontier as flat [embIdx, mmIdx, ...] pairs
+    if (!params_.linkFrontier_.empty())
+    {
+        auto linkAdder = treeBuilder.addUniformArray<std::int64_t>(
+                METATOMIC_MODULE_NAME + "-link-frontier");
+        for (const auto& link : params_.linkFrontier_)
+        {
+            linkAdder.addValue(link.getEmbeddedIndex());
+            linkAdder.addValue(link.getMMIndex());
+        }
+    }
 }
 
 void MetatomicOptions::readParamsFromKvt(const KeyValueTreeObject& tree)
@@ -332,6 +412,20 @@ void MetatomicOptions::readParamsFromKvt(const KeyValueTreeObject& tree)
                    std::end(kvtIndexArray),
                    std::begin(params_.mtaIndices_),
                    [](const KeyValueTreeValue& val) { return val.cast<std::int64_t>(); });
+
+    // Deserialize link frontier
+    std::string linkKey = METATOMIC_MODULE_NAME + "-link-frontier";
+    if (tree.keyExists(linkKey))
+    {
+        auto linkArray = tree[linkKey].asArray().values();
+        params_.linkFrontier_.clear();
+        for (size_t i = 0; i + 1 < linkArray.size(); i += 2)
+        {
+            int embIdx = static_cast<int>(linkArray[i].cast<std::int64_t>());
+            int mmIdx  = static_cast<int>(linkArray[i + 1].cast<std::int64_t>());
+            params_.linkFrontier_.emplace_back(embIdx, mmIdx);
+        }
+    }
 }