fix(mta): correct threading comment, model-driven embedding, link docs

- 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

Author: HaoZeke

src/gromacs/applied_forces/metatomic/metatomic_forceprovider.cpp

@@ -241,11 +241,15 @@ 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.
-    // 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.
+    // Synchronize PyTorch's thread count with GROMACS's ntomp.
+    //
+    // at::set_num_threads updates PyTorch's cached thread count AND
+    // MKL's thread pool (via mkl_set_num_threads), while GROMACS's
+    // omp_set_num_threads only updates the OpenMP ICV.  Without this
+    // call, PyTorch/MKL may retain the init-time default (all cores).
+    //
+    // For real MPI, ntomp defaults to all cores for single-rank and to
+    // the per-rank count for multi-rank, so the behavior is unchanged.
     {
         int ntomp = gmx_omp_nthreads_get(ModuleMultiThread::Default);
         at::set_num_threads(std::max(1, ntomp));
@@ -1166,10 +1170,13 @@ 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.
+    // Link atom setup: boundary MM atoms are replaced with hydrogen link
+    // atoms in the model input.  Positions are overwritten here (in C++),
+    // then recomputed inside the autograd graph (via torch ops) so that
+    // forces are automatically correct via the chain rule.
+    //
+    // The user sees this as: "N embedded atoms, M of which are link atoms
+    // (hydrogen caps at ML/MM boundary bonds)."
     if (!data_->linkFrontier.empty())
     {
         bool needTypesRebuild = false;

src/gromacs/applied_forces/metatomic/metatomic_options.cpp

@@ -86,7 +86,6 @@ 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
 {
@@ -176,8 +175,6 @@ void MetatomicOptions::initMdpTransform(IKeyValueTreeTransformRules* rules)
             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)
@@ -197,7 +194,6 @@ void MetatomicOptions::initMdpOptions(IOptionsContainerWithSections* options)
     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
@@ -235,8 +231,6 @@ void MetatomicOptions::buildMdpOutput(KeyValueTreeObjectBuilder* builder) const
                 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);
     }
 }
 
@@ -274,24 +268,19 @@ void MetatomicOptions::modifyTopology(gmx_mtop_t* top)
         return;
     }
 
-    // Collect MM charges before preprocessing (which may modify charges)
-    if (params_.electrostaticEmbedding)
+    // Always collect MM charges (cheap). The force provider will pass them
+    // to the model only if the model declares point_charges in its
+    // requested_inputs (model-driven, not user-controlled).
+    for (const auto& molblock : top->molblock)
     {
-        for (const auto& molblock : top->molblock)
+        const auto& moltype = top->moltype[molblock.type];
+        for (int m = 0; m < molblock.nmol; m++)
         {
-            const auto& moltype = top->moltype[molblock.type];
-            for (int m = 0; m < molblock.nmol; m++)
+            for (int a = 0; a < moltype.atoms.nr; a++)
             {
-                for (int a = 0; a < moltype.atoms.nr; a++)
-                {
-                    params_.mmCharges_.push_back(moltype.atoms.atom[a].q);
-                }
+                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)

src/gromacs/applied_forces/metatomic/metatomic_options.h

@@ -105,8 +105,6 @@ struct MetatomicParameters
 
     //! Enable ONIOM link atoms at cut bonds between ML and MM regions
     bool linkAtoms = false;
-    //! Enable electrostatic embedding (pass MM point charges to ML model)
-    bool electrostaticEmbedding = false;
 
     std::vector<Index>            mtaIndices_;
     std::vector<Index>            mmIndices_;