feat(mta): implement link atom injection and force redistribution

Per-step link atom handling in the force provider:
- After gathering ML atom positions, compute link atom positions via
  LinkFrontierAtom::setPositions/getLinkPosition and append to the model input
- Build neighbor list pairs between link atoms and ML atoms within cutoff
- After model evaluation, redistribute link atom forces to real atoms via
  LinkFrontierAtom::spreadForce (pattern from NNPot torchmodel.cpp)
- Strip link atoms from arrays after force scatter

Also stores link frontier and MM charges in MetatomicData from options.

Electrostatic embedding data collection is in place (MM charges stored)
but the per-step add_data() call to pass them to the model is deferred
until a model that accepts point_charges input is available.

Author: HaoZeke

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

@@ -164,6 +164,13 @@ 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;
+    //! MM point charges for electrostatic embedding (global atom index -> charge).
+    std::vector<real> mmCharges;
+    //! Number of link atoms (appended after real ML atoms in positions_).
+    int32_t numLinkAtoms = 0;
+
     //! Non-conservative mode: forces/stress predicted directly, no backward pass.
     bool nonConservative = false;
     //! Output keys for non-conservative forces and stress.
@@ -486,6 +493,26 @@ MetatomicForceProvider::MetatomicForceProvider(const MetatomicOptions& options,
                         data_->nc_stress_key.c_str());
     }
 
+    // Store link frontier and MM charges from preprocessing
+    data_->linkFrontier = options_.params_.linkFrontier_;
+    data_->mmCharges    = options_.params_.mmCharges_;
+    data_->numLinkAtoms = static_cast<int32_t>(data_->linkFrontier.size());
+
+    if (!data_->linkFrontier.empty())
+    {
+        GMX_LOG(logger_.info)
+                .asParagraph()
+                .appendTextFormatted("Metatomic: %d link atoms at ML/MM boundary",
+                                     data_->numLinkAtoms);
+    }
+    if (!data_->mmCharges.empty())
+    {
+        GMX_LOG(logger_.info)
+                .asParagraph()
+                .appendTextFormatted("Metatomic: electrostatic embedding with %zu charges",
+                                     data_->mmCharges.size());
+    }
+
     GMX_LOG(logger_.info)
             .asParagraph()
             .appendText("MetatomicForceProvider initialization complete.");
@@ -1153,6 +1180,49 @@ void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, F
     }
     copy_mat(inputs.box_, box_);
 
+    // Inject link atoms at ML/MM boundary bonds.  Link atoms are appended
+    // after the real ML atoms in positions_/atomNumbers_.  Their forces will
+    // be redistributed to the real atoms after model evaluation.
+    if (data_->numLinkAtoms > 0)
+    {
+        for (auto& link : data_->linkFrontier)
+        {
+            // Get embedded (ML) and MM atom positions from GROMACS coordinates.
+            // link.getEmbeddedIndex() and link.getMMIndex() are global indices.
+            const int embGmxLocal = options_.params_.mtaAtoms_->localIndex()[
+                    std::distance(options_.params_.mtaIndices_.begin(),
+                                  std::find(options_.params_.mtaIndices_.begin(),
+                                            options_.params_.mtaIndices_.end(),
+                                            link.getEmbeddedIndex()))];
+            const RVec posEmb = inputs.x_[embGmxLocal];
+            // MM atom: use global index to find in full coordinate array
+            // For single rank, global index == local index
+            const RVec posMM = inputs.x_[link.getMMIndex()];
+
+            link.setPositions(posEmb, posMM);
+            RVec linkPos = link.getLinkPosition();
+
+            // Set input indices: embedded atom -> its MTA model index,
+            // link atom -> the index we're about to append
+            int32_t embMtaIdx = -1;
+            for (int32_t m = 0; m < numLocalMta_; ++m)
+            {
+                if (mtaToGlobalMta_[m] == link.getEmbeddedIndex())
+                {
+                    embMtaIdx = m;
+                    break;
+                }
+            }
+            int32_t linkMtaIdx = numLocalMta_ + static_cast<int32_t>(
+                    &link - &data_->linkFrontier[0]);
+            link.setInputIndices(embMtaIdx, linkMtaIdx);
+
+            positions_.push_back(linkPos);
+            atomNumbers_.push_back(link.linkAtomNumber());
+        }
+        numLocalMta_ += data_->numLinkAtoms;
+    }
+
     // 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
@@ -1337,6 +1407,51 @@ void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, F
                 nPairs = outIdx;
                 nlSamplesBuffer_.resize(nPairs * 5);
                 nlVectorsBuffer_.resize(nPairs * 3);
+
+                // Add pairs between link atoms and ML atoms within cutoff
+                if (data_->numLinkAtoms > 0)
+                {
+                    const int32_t numRealMta = numLocalMta_ - data_->numLinkAtoms;
+                    for (int32_t li = 0; li < data_->numLinkAtoms; ++li)
+                    {
+                        int32_t linkIdx = numRealMta + li;
+                        RVec    linkPos = positions_[linkIdx];
+                        for (int32_t mi = 0; mi < numRealMta; ++mi)
+                        {
+                            RVec    mlPos = positions_[mi];
+                            double  dx    = linkPos[0] - mlPos[0];
+                            double  dy    = linkPos[1] - mlPos[1];
+                            double  dz    = linkPos[2] - mlPos[2];
+                            double  r2    = dx * dx + dy * dy + dz * dz;
+                            if (r2 < cutoff2)
+                            {
+                                // link -> ml pair
+                                nlSamplesBuffer_.push_back(linkIdx);
+                                nlSamplesBuffer_.push_back(mi);
+                                nlSamplesBuffer_.push_back(0);
+                                nlSamplesBuffer_.push_back(0);
+                                nlSamplesBuffer_.push_back(0);
+                                nlVectorsBuffer_.push_back(-dx);
+                                nlVectorsBuffer_.push_back(-dy);
+                                nlVectorsBuffer_.push_back(-dz);
+                                nPairs++;
+                                if (fullList)
+                                {
+                                    // ml -> link pair (reverse)
+                                    nlSamplesBuffer_.push_back(mi);
+                                    nlSamplesBuffer_.push_back(linkIdx);
+                                    nlSamplesBuffer_.push_back(0);
+                                    nlSamplesBuffer_.push_back(0);
+                                    nlSamplesBuffer_.push_back(0);
+                                    nlVectorsBuffer_.push_back(dx);
+                                    nlVectorsBuffer_.push_back(dy);
+                                    nlVectorsBuffer_.push_back(dz);
+                                    nPairs++;
+                                }
+                            }
+                        }
+                    }
+                }
             }
 
             // Debug: dump global pair indices + distances for comparison
@@ -1608,8 +1723,44 @@ void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, F
     }
     else
     {
-        // Serial: apply forces directly
-        for (int32_t i = 0; i < numLocalMta_; i++)
+        // Redistribute link atom forces to real atoms before scattering.
+        // Link atoms occupy indices [numRealMta, numLocalMta_) in the force array.
+        if (data_->numLinkAtoms > 0)
+        {
+            t_pbc pbc;
+            set_pbc(&pbc, *options_.params_.pbcType_, inputs.box_);
+
+            const int32_t numRealMta = numLocalMta_ - data_->numLinkAtoms;
+            for (int32_t li = 0; li < data_->numLinkAtoms; ++li)
+            {
+                auto& link = data_->linkFrontier[li];
+                int32_t linkIdx = numRealMta + li;
+                RVec fL = { static_cast<real>(forceData[3 * linkIdx]),
+                            static_cast<real>(forceData[3 * linkIdx + 1]),
+                            static_cast<real>(forceData[3 * linkIdx + 2]) };
+                auto [fEmb, fMM] = link.spreadForce(fL, pbc);
+
+                // Add redistributed force to embedded (ML) atom
+                int32_t embIdx = link.getInputIndexEmb();
+                if (embIdx >= 0 && embIdx < numRealMta)
+                {
+                    int32_t gmxIdx = mtaToGmxLocal_[embIdx];
+                    outputs->forceWithVirial_.force_[gmxIdx][0] += fEmb[0];
+                    outputs->forceWithVirial_.force_[gmxIdx][1] += fEmb[1];
+                    outputs->forceWithVirial_.force_[gmxIdx][2] += fEmb[2];
+                }
+
+                // Add redistributed force to MM atom
+                int32_t mmGlobalIdx = link.getMMIndex();
+                outputs->forceWithVirial_.force_[mmGlobalIdx][0] += fMM[0];
+                outputs->forceWithVirial_.force_[mmGlobalIdx][1] += fMM[1];
+                outputs->forceWithVirial_.force_[mmGlobalIdx][2] += fMM[2];
+            }
+        }
+
+        // Serial: apply forces for real ML atoms (excluding link atoms)
+        const int32_t numRealMta = numLocalMta_ - data_->numLinkAtoms;
+        for (int32_t i = 0; i < numRealMta; i++)
         {
             int32_t gmxIdx = mtaToGmxLocal_[i];
             outputs->forceWithVirial_.force_[gmxIdx][0] += static_cast<real>(forceData[3 * i]);
@@ -1620,6 +1771,14 @@ void MetatomicForceProvider::calculateForces(const ForceProviderInput& inputs, F
 
     forceScatterTimer.stop();
 
+    // Strip link atoms from arrays (they were appended in this step)
+    if (data_->numLinkAtoms > 0)
+    {
+        numLocalMta_ -= data_->numLinkAtoms;
+        positions_.resize(numLocalMta_);
+        atomNumbers_.resize(numLocalMta_);
+    }
+
     // Restore original state.  Backward ghost atoms have no GROMACS local
     // buffer index, and backward pairs should not persist across steps.
     if (useNewtonNL)