fix(ml/mm): do ONIOM

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.
@@ -238,35 +241,14 @@ MetatomicForceProvider::MetatomicForceProvider(const MetatomicOptions& options,
     // For GPU devices, CPU overhead is minimal so we keep 1 thread to avoid
     // oversubscription with GROMACS threads.  For CPU devices, model inference
     // (matmuls, convolutions) benefits from multi-threading.
-    if (data_->device.is_cpu())
+    // Set PyTorch thread count to match GROMACS OpenMP threads.
+    // This is critical for thread-MPI builds: PyTorch's default (all cores)
+    // conflicts with thread-MPI's internal threading, causing incorrect
+    // forces and simulation blow-up.  For real MPI with multiple ranks,
+    // this also prevents oversubscription.
     {
-#if GMX_THREAD_MPI
-        // Thread-MPI: ranks share a process.  PyTorch's global thread pool
-        // would be contended by all ranks calling forward() concurrently,
-        // so keep at 1 to avoid oversubscription.
-        if (mpiComm_.isParallel())
-        {
-            at::set_num_threads(1);
-        }
-#else
-        // Real MPI (or no MPI): each rank is a separate process.
-        // Use the GROMACS-assigned OpenMP thread count so that PyTorch
-        // can parallelize matrix operations within each rank's allocation.
-        if (mpiComm_.isParallel())
-        {
-            int ntomp = gmx_omp_nthreads_get(ModuleMultiThread::Default);
-            at::set_num_threads(std::max(1, ntomp));
-        }
-        // Serial: let PyTorch use its default (all cores)
-#endif
-    }
-    else
-    {
-        // GPU/other device: model runs on accelerator, CPU work is minimal.
-        if (mpiComm_.isParallel())
-        {
-            at::set_num_threads(1);
-        }
+        int ntomp = gmx_omp_nthreads_get(ModuleMultiThread::Default);
+        at::set_num_threads(std::max(1, ntomp));
     }
 
     // JIT fusion: dynamic strategy with depth limit of 10 improves CPU
@@ -507,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.");
@@ -1174,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
@@ -1277,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);
 
@@ -1629,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"
@@ -83,13 +85,20 @@ static const std::string VARIANT_ENERGY_UQ_TAG     = "variant-energy-uq";
 static const std::string NON_CONSERVATIVE_TAG      = "non-conservative";
 static const std::string VARIANT_NC_FORCES_TAG     = "variant-nc-forces";
 static const std::string VARIANT_NC_STRESS_TAG     = "variant-nc-stress";
+static const std::string LINK_ATOMS_TAG            = "link-atoms";
+static const std::string ELECTROSTATIC_EMBEDDING_TAG = "electrostatic-embedding";
 
 namespace
 {
 // TODO(rg): this is duplicated from the nnpotoptions
 
 //! \brief Helper function to preprocess topology for MTA
-void preprocessTopology(gmx_mtop_t* mtop, ArrayRef<const Index> mtaIndices, const MDLogger& logger, WarningHandler* wi)
+void preprocessTopology(gmx_mtop_t*                    mtop,
+                        ArrayRef<const Index>           mtaIndices,
+                        const MDLogger&                 logger,
+                        WarningHandler*                 wi,
+                        bool                            buildLinks,
+                        std::vector<LinkFrontierAtom>*  linkFrontierOut)
 {
     // convert mtaIndices to set for faster lookup
     std::set<int> mtaIndicesSet(mtaIndices.begin(), mtaIndices.end());
@@ -116,15 +125,24 @@ void preprocessTopology(gmx_mtop_t* mtop, ArrayRef<const Index> mtaIndices, cons
     // 4) Make F_CONNBOND between atoms within QM region
     modifyEmbeddedTwoCenterInteractions(mtop, mtaIndicesSet, isMTABlock, logger);
 
-    // 5) Remove angles and settles containing 2 or more QM atoms
+    // 5) Remove angles and settles containing all-ML atoms (ONIOM)
     modifyEmbeddedThreeCenterInteractions(mtop, mtaIndicesSet, isMTABlock, logger);
 
-    // 6) Remove dihedrals containing 3 or more QM atoms
+    // 6) Remove dihedrals containing all-ML atoms (ONIOM)
     modifyEmbeddedFourCenterInteractions(mtop, mtaIndicesSet, isMTABlock, logger);
 
     // 7) Check for constrained bonds in subsystem
     checkConstrainedBonds(mtop, mtaIndicesSet, isMTABlock, wi);
 
+    // 8) Build link frontier atoms at ML/MM boundary bonds
+    if (buildLinks && linkFrontierOut != nullptr)
+    {
+        *linkFrontierOut = buildLinkFrontier(mtop, mtaIndicesSet, isMTABlock, logger);
+        GMX_LOG(logger.info)
+                .appendTextFormatted("Number of link frontier atoms: %zu",
+                                     linkFrontierOut->size());
+    }
+
     // finalize topology
     mtop->finalize();
 }
@@ -156,6 +174,10 @@ void MetatomicOptions::initMdpTransform(IKeyValueTreeTransformRules* rules)
             rules, stringIdentityTransform, METATOMIC_MODULE_NAME, VARIANT_NC_FORCES_TAG);
     addMdpTransformFromString<std::string>(
             rules, stringIdentityTransform, METATOMIC_MODULE_NAME, VARIANT_NC_STRESS_TAG);
+    addMdpTransformFromString<bool>(
+            rules, &fromStdString<bool>, METATOMIC_MODULE_NAME, LINK_ATOMS_TAG);
+    addMdpTransformFromString<bool>(
+            rules, &fromStdString<bool>, METATOMIC_MODULE_NAME, ELECTROSTATIC_EMBEDDING_TAG);
 }
 
 void MetatomicOptions::initMdpOptions(IOptionsContainerWithSections* options)
@@ -174,6 +196,8 @@ void MetatomicOptions::initMdpOptions(IOptionsContainerWithSections* options)
     section.addOption(BooleanOption(NON_CONSERVATIVE_TAG.c_str()).store(&params_.nonConservative));
     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));
 }
 
 void MetatomicOptions::buildMdpOutput(KeyValueTreeObjectBuilder* builder) const
@@ -209,6 +233,10 @@ void MetatomicOptions::buildMdpOutput(KeyValueTreeObjectBuilder* builder) const
                 builder, METATOMIC_MODULE_NAME, VARIANT_NC_FORCES_TAG, params_.variantNcForces);
         addMdpOutputValue<std::string>(
                 builder, METATOMIC_MODULE_NAME, VARIANT_NC_STRESS_TAG, params_.variantNcStress);
+        addMdpOutputValue<bool>(
+                builder, METATOMIC_MODULE_NAME, LINK_ATOMS_TAG, params_.linkAtoms);
+        addMdpOutputValue<bool>(
+                builder, METATOMIC_MODULE_NAME, ELECTROSTATIC_EMBEDDING_TAG, params_.electrostaticEmbedding);
     }
 }
 
@@ -245,7 +273,95 @@ void MetatomicOptions::modifyTopology(gmx_mtop_t* top)
     {
         return;
     }
-    preprocessTopology(top, params_.mtaIndices_, logger(), wi_);
+
+    // Collect MM charges before preprocessing (which may modify charges)
+    if (params_.electrostaticEmbedding)
+    {
+        for (const auto& molblock : top->molblock)
+        {
+            const auto& moltype = top->moltype[molblock.type];
+            for (int m = 0; m < molblock.nmol; m++)
+            {
+                for (int a = 0; a < moltype.atoms.nr; a++)
+                {
+                    params_.mmCharges_.push_back(moltype.atoms.atom[a].q);
+                }
+            }
+        }
+        GMX_LOG(logger().info)
+                .appendTextFormatted("Metatomic: collected %zu point charges for "
+                                     "electrostatic embedding",
+                                     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_,
+                       /*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)
@@ -261,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)
@@ -284,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);
+        }
+    }
 }