add atom_modify map yes

Author: OutisLi

deepmd/pt/entrypoints/freeze_pt2.py

@@ -74,6 +74,22 @@ def _get_model_ntypes(model: torch.nn.Module) -> int:
     return int(descriptor.get_ntypes())
 
 
+def _model_has_message_passing(model: torch.nn.Module) -> bool:
+    """Return whether the regular .pt2 graph requires a real atom mapping."""
+    for obj in (
+        model,
+        getattr(model, "atomic_model", None),
+        model.get_descriptor() if hasattr(model, "get_descriptor") else None,
+    ):
+        if obj is None or not hasattr(obj, "has_message_passing"):
+            continue
+        try:
+            return bool(obj.has_message_passing())
+        except (AttributeError, NotImplementedError):
+            continue
+    return False
+
+
 def _strip_shape_assertions(graph_module: torch.nn.Module) -> None:
     """Remove deferred shape assertions from spin export graphs.
 
@@ -243,6 +259,8 @@ def _collect_metadata(
         "dim_aparam": int(model.get_dim_aparam()),
         "dim_chg_spin": int(model.get_dim_chg_spin()),
         "mixed_types": bool(model.mixed_types()),
+        "has_message_passing": _model_has_message_passing(model),
+        "has_comm_artifact": False,
         "has_default_fparam": bool(model.has_default_fparam()),
         "default_fparam": _to_py_list(model.get_default_fparam()),
         "default_chg_spin": _to_py_list(model.get_default_chg_spin()),

deepmd/pt/model/model/sezm_model.py

@@ -2265,7 +2265,7 @@ def post_process_output_dens(
         )
 
     # =========================================================================
-    # Charge/Spin Condition Metadata
+    # Metadata
     # =========================================================================
 
     def has_chg_spin_ebd(self) -> bool:
@@ -2284,6 +2284,10 @@ def get_default_chg_spin(self) -> torch.Tensor | None:
         """Return default charge/spin conditions as a tensor."""
         return self.atomic_model.get_default_chg_spin()
 
+    def has_message_passing(self) -> bool:
+        """Return whether the descriptor performs message passing."""
+        return self.atomic_model.has_message_passing()
+
     # =========================================================================
     # Mode Management
     # =========================================================================

examples/water/dpa4/lmp/README.md

@@ -59,6 +59,9 @@ Step   PotEng       KinEng      TotEng     Temp
   entries `"O"` and `"H"` respectively. When the element names are
   omitted, the mapping falls back to the `type_map` order stored in
   the `.pt2` metadata.
+- `atom_modify map yes` keeps the ghost / periodic-image to local-atom
+  mapping explicit for `.pt2` graph inference. GNN-style `.pt2` models
+  fail fast when this atom map is required but absent.
 - The 500-step `pretrained.pt` is intended as a smoke test, not a
   physically accurate water potential. Retrain with a longer schedule
   for production.

examples/water/dpa4/lmp/in.lammps

@@ -9,6 +9,7 @@
 units           metal
 boundary        p p p
 atom_style      atomic
+atom_modify     map yes
 
 neighbor        2.0 bin
 neigh_modify    every 10 delay 0 check no

source/api_cc/tests/test_deeppot_ptexpt.cc

@@ -12,7 +12,7 @@
 #include "DeepPot.h"
 #include "DeepPotPTExpt.h"
 #if defined(BUILD_PYTORCH)
-#include "commonPTExpt.h"
+#include "../src/commonPTExpt.h"
 #endif
 #include "neighbor_list.h"
 #include "test_utils.h"

source/lmp/pair_deepmd.cpp

@@ -188,7 +188,8 @@ void PairDeepMD::compute(int eflag, int vflag) {
     }
   }
 
-  // mapping (for DPA-2 JAX)
+  // mapping (for DPA-2/3 .pt2 GNN models that gather ghost features via
+  // the LAMMPS atom-map; harmless for other models).
   std::vector<int> mapping_vec(nall, -1);
   if (comm->nprocs == 1 && atom->map_style != Atom::MAP_NONE) {
     for (size_t ii = 0; ii < nall; ++ii) {

source/tests/pt/model/test_sezm_export.py

@@ -367,6 +367,8 @@ def test_archive_metadata(self) -> None:
             "dim_aparam",
             "dim_chg_spin",
             "mixed_types",
+            "has_message_passing",
+            "has_comm_artifact",
             "has_default_fparam",
             "default_chg_spin",
             "output_keys",
@@ -381,6 +383,8 @@ def test_archive_metadata(self) -> None:
         self.assertEqual(metadata["rcut"], self.params["descriptor"]["rcut"])
         self.assertEqual(list(metadata["sel"]), list(self.params["descriptor"]["sel"]))
         self.assertTrue(metadata["mixed_types"])
+        self.assertTrue(metadata["has_message_passing"])
+        self.assertFalse(metadata["has_comm_artifact"])
         self.assertFalse(metadata["is_spin"])
         self.assertEqual(metadata["dim_fparam"], 0)
         self.assertEqual(metadata["dim_aparam"], 0)
@@ -592,6 +596,16 @@ def test_metadata_records_ntypes_when_type_map_is_empty(self) -> None:
         self.assertEqual(metadata["type_map"], [])
         self.assertEqual(metadata["ntypes"], model.get_descriptor().get_ntypes())
 
+    def test_metadata_records_message_passing_contract(self) -> None:
+        """PTExpt fail-fast depends on the SeZM-specific metadata contract."""
+        model = _build_tiny_sezm_model()
+
+        metadata = _collect_metadata(model, ["energy"])
+
+        self.assertTrue(model.has_message_passing())
+        self.assertTrue(metadata["has_message_passing"])
+        self.assertFalse(metadata["has_comm_artifact"])
+
     def test_charge_spin_export_sample_has_runtime_input_slot(self) -> None:
         """Charge/spin-conditioned exports should not bake defaults into the graph."""
         params = _tiny_sezm_model_params()
@@ -722,6 +736,8 @@ def fake_compile(_exported: torch.export.ExportedProgram, package_path: str):
         self.assertTrue(metadata["is_spin"])
         self.assertEqual(metadata["type_map"], params["type_map"])
         self.assertEqual(metadata["ntypes"], len(params["type_map"]))
+        self.assertTrue(metadata["has_message_passing"])
+        self.assertFalse(metadata["has_comm_artifact"])
         self.assertEqual(metadata["dim_chg_spin"], 0)
         self.assertIsNone(metadata["default_chg_spin"])
         self.assertEqual(metadata["use_spin"], params["spin"]["use_spin"])