fix: refresh the ZBL lookup on change_type_map; allow native spin + bridging

Two review items.

**3649295675 -- ZBL element lookup went stale on change_type_map.**
``InterPotentialAtomicModel`` inherited the generic
``BaseAtomicModel.change_type_map``, which rewrites the public map and the
stat/exclusion state only, so ``InterPotential`` kept its ORIGINAL nuclear
charges while ``atype`` values already meant the new elements -- silently
wrong energies after a reorder, ``IndexError`` after adding a type, and a
restart-time jump because serialization records the new map.  The lookup
belongs to ``InterPotential``, so it owns the update: a new
``change_type_map`` rebuilds symbols, count and table together from the
symbols (not by permuting indices), keeping the array's namespace/dtype/
device so a pt_expt buffer -- possibly on CUDA -- is updated in place.  The
atomic model delegates rather than reimplementing.

**3649276109 -- native spin + analytical bridging was rejected.**
The combination needs no new machinery: ``get_standard_model`` already OWNS
assembling the atomic model, bridging composition included, so the guard is
gone and ``get_native_spin_model`` now delegates to it and re-classes
whatever atomic model comes back -- a single learned model or a
``LinearEnergyAtomicModel`` over ``[learned, InterPotential]``.  The learned
child consumes ``spin``; the analytical child accepts and ignores it.  Two
supporting changes: a shared ``learned_descriptor()`` so the descriptor
capability gate reads the learned child either way, and ``NSM.serialize``
now records the backbone's own wire type (``backbone_type``) -- it was
hard-coded to "standard", which fails on a composed backbone whose dict has
a different shape and @Version.  Archives without the field still read as
"standard".

Also fixes the C++ pair-exclusion regression added in ecfb7a8: its
``cpu_lmp_nlist`` case compared an nghost=0 explicit-nlist run (no periodic
images) against the PBC reference; it now uses the gas-phase one, the same
convention test_deepspin_dpa4_graph_ptexpt.cc follows.

Tests: change_type_map reorder / add / drop / serialize-roundtrip in both
dpmodel and pt_expt (the pt_expt twin asserts the rebuilt lookup is still a
torch buffer on the model device); native spin + ZBL construction, eager
energy/force/force_mag with mask_mag, the analytical term checked against an
independent in-test pair loop at 1e-10, serialization, and graph
freeze + DeepPot parity at 1e-10.  The stale
``test_native_spin_with_bridging_fails_fast`` is replaced by those.

Author: Han Wang

deepmd/dpmodel/atomic_model/inter_potential.py

@@ -97,13 +97,64 @@ def __init__(self, type_map: list[str], mode: str = "zbl") -> None:
         self.mode = mode
         self.type_map = list(type_map)
         self.ntypes_real = len(type_map)
+        self.atomic_numbers = self._lookup_from_type_map(type_map)
+
+    @staticmethod
+    def _lookup_from_type_map(type_map: list[str], like: Array | None = None) -> Array:
+        """Build the per-type nuclear-charge lookup from element symbols.
+
+        Parameters
+        ----------
+        type_map : list[str]
+            Element symbols; index corresponds to ``atype`` values.
+        like : Array, optional
+            When given, the result is created in this array's namespace, dtype
+            and device instead of NumPy -- so an in-place rebuild on a wrapped
+            backend (pt_expt buffer, possibly on CUDA) stays where it was.
+
+        Returns
+        -------
+        Array
+            Nuclear charges, shape ``(len(type_map),)``.
+
+        Raises
+        ------
+        ValueError
+            If an element symbol is not in :data:`ELEMENT_TO_Z`.
+        """
         atomic_numbers = []
         for elem in type_map:
             z = ELEMENT_TO_Z.get(elem)
             if z is None:
                 raise ValueError(f"Unknown element symbol: {elem}")
             atomic_numbers.append(z)
-        self.atomic_numbers = np.asarray(atomic_numbers, dtype=np.float64)
+        arr = np.asarray(atomic_numbers, dtype=np.float64)
+        if like is None:
+            return arr
+        xp = array_api_compat.array_namespace(like)
+        return xp.asarray(arr, dtype=like.dtype, device=array_api_compat.device(like))
+
+    def change_type_map(self, type_map: list[str]) -> None:
+        """Rebuild the element lookup for a new type map.
+
+        THIS OWNS the element lookup, so it owns every update of it: the
+        symbols, their count and the nuclear-charge table are one piece of
+        state and are replaced together.  Reordering, adding and dropping
+        elements are all covered -- the table is rebuilt from the symbols
+        rather than permuted, so no index bookkeeping can drift.  The rebuilt
+        array keeps the current one's namespace/dtype/device, so a wrapped
+        backend (pt_expt buffer on CPU or CUDA) is updated in place.
+
+        Parameters
+        ----------
+        type_map : list[str]
+            The new element symbols.
+        """
+        self.atomic_numbers = self._lookup_from_type_map(
+            type_map, like=self.atomic_numbers
+        )
+        self.type_map = list(type_map)
+        self.ntypes_real = len(type_map)
 
     @staticmethod
     def _zbl_pair_energy(xp: Any, r: Array, zi: Array, zj: Array) -> Array:
@@ -259,6 +310,31 @@ def __init__(
         )
         super().init_out_stat()
 
+    def change_type_map(
+        self, type_map: list[str], model_with_new_type_stat: Any | None = None
+    ) -> None:
+        """Change the type related params to new ones, according to `type_map` and the original one in the model.
+        If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types.
+
+        The generic base handles the public map and the stat/exclusion state;
+        the element lookup belongs to :class:`InterPotential`, so the update is
+        delegated there rather than reimplemented here (review 3649295675 --
+        without it the lookup keeps the ORIGINAL elements while ``atype``
+        values mean new ones, and a longer new map raises ``IndexError``).
+
+        Parameters
+        ----------
+        type_map : list[str]
+            The new element symbols.
+        model_with_new_type_stat : optional
+            Model with statistics for the new types (unused: an analytical
+            term has no fitted statistics).
+        """
+        super().change_type_map(
+            type_map, model_with_new_type_stat=model_with_new_type_stat
+        )
+        self.potential.change_type_map(type_map)
+
     def fitting_output_def(self) -> FittingOutputDef:
         """Per-atom analytical energy: reducible and fully differentiable."""
         return FittingOutputDef(

deepmd/dpmodel/model/model.py

@@ -240,6 +240,36 @@ def get_spin_model(data: dict) -> SpinModel:
     return SpinModel(backbone_model=backbone_model, spin=spin)
 
 
+def learned_descriptor(atomic_model: Any) -> Any:
+    """The descriptor of the LEARNED term of an atomic model.
+
+    A plain atomic model owns its descriptor directly; a composition (e.g.
+    analytical bridging, ``LinearEnergyAtomicModel`` over ``[learned,
+    InterPotential]``) has exactly one descriptor-bearing child, since the
+    analytical terms are descriptor-free.  Shared by every backend's builder
+    so descriptor capability gates read the same object regardless of whether
+    the model happens to be composed.
+
+    Parameters
+    ----------
+    atomic_model
+        The atomic model to inspect.
+
+    Returns
+    -------
+    descriptor or None
+        The learned descriptor, or ``None`` if no child has one.
+    """
+    descriptor = getattr(atomic_model, "descriptor", None)
+    if descriptor is not None:
+        return descriptor
+    for child in getattr(atomic_model, "models", []):
+        descriptor = getattr(child, "descriptor", None)
+        if descriptor is not None:
+            return descriptor
+    return None
+
+
 def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     """Get a native (virtual-atom-free) spin model from a dictionary.
 
@@ -250,19 +280,22 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     eligible; the gate is the capability method, not a descriptor-type
     list.
 
+    The non-spin backbone is built by :func:`get_standard_model`, which OWNS
+    everything about assembling the atomic model -- descriptor/fitting,
+    exclusions and the analytical-bridging composition -- so ``spin`` and
+    ``bridging_method`` combine for free: the wrapper re-classes whatever
+    atomic model came back, be it a single learned model or a
+    ``LinearEnergyAtomicModel`` over ``[learned, InterPotential]`` (the
+    analytical child accepts and ignores ``spin``; the learned child consumes
+    it).
+
     Parameters
     ----------
     data : dict
         The data to construct the model.
     """
     data = copy.deepcopy(data)
     spin_cfg = data.pop("spin")
-    if str(data.get("bridging_method", "none")).lower() not in ("none", ""):
-        raise NotImplementedError(
-            "analytical bridging combined with the native spin scheme is a "
-            "follow-up (the bridged model is a linear composition; the "
-            "native-spin factory composes over a single standard model)"
-        )
     # Expand index/symbol forms of ``use_spin`` against ``type_map`` into the
     # per-type boolean list (pure; validates symbols).
     use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"])
@@ -271,10 +304,10 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
         virtual_scale=spin_cfg.get("virtual_scale", 1.0),
         allow_missing_label=spin_cfg.get("allow_missing_label", False),
     )
+    data.setdefault("descriptor", {})
     data["descriptor"]["use_spin"] = use_spin
-    ntypes = len(data["type_map"])
     try:
-        descriptor, fitting, _ = _get_standard_model_components(data, ntypes)
+        backbone_model = get_standard_model(data)
     except TypeError as err:
         if "use_spin" not in str(err):
             # Unrelated construction error (e.g. a bogus fitting kwarg):
@@ -289,19 +322,13 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
             "support (supports_native_spin()); descriptor type "
             f"{data['descriptor'].get('type')!r} does not accept `use_spin`"
         ) from err
-    if not descriptor.supports_native_spin():
+    descriptor = learned_descriptor(backbone_model.atomic_model)
+    if descriptor is None or not descriptor.supports_native_spin():
         raise NotImplementedError(
             "spin scheme 'native' requires a descriptor declaring "
             "supports_native_spin()"
         )
-    return NativeSpinEnergyModel(
-        descriptor=descriptor,
-        fitting=fitting,
-        type_map=data["type_map"],
-        atom_exclude_types=data.get("atom_exclude_types", []),
-        pair_exclude_types=data.get("pair_exclude_types", []),
-        spin=spin,
-    )
+    return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin)
 
 
 def get_model(data: dict) -> BaseModel:

deepmd/dpmodel/model/native_spin_model.py

@@ -227,6 +227,12 @@ def call(
 
         def serialize(self) -> dict:
             data = super().serialize()
+            # The backbone's own wire type would be lost under "native_spin";
+            # keep it so deserialize can rebuild the RIGHT backbone.  It is
+            # not always "standard": with analytical bridging the backbone is
+            # a composition ("linear"), whose dict has a different shape and
+            # @version.
+            data["backbone_type"] = data.get("type", "standard")
             data["type"] = "native_spin"
             data["spin"] = self.spin.serialize()
             return data
@@ -236,11 +242,18 @@ def deserialize(cls, data: dict) -> "NSM":
             data = data.copy()
             data.pop("type", None)
             spin = Spin.deserialize(data.pop("spin"))
-            # make_model flat shape: the remaining dict IS the standard
-            # model (atomic) dict -- its @class/@version belong to the
-            # atomic deserialize and must stay.
-            data["type"] = "standard"
-            backbone = T_Model.deserialize(data)
+            # make_model flat shape: the remaining dict IS the backbone
+            # (atomic) dict -- its @class/@version belong to the backbone's
+            # deserialize and must stay.  Archives written before
+            # ``backbone_type`` existed are all plain standard models.
+            backbone_type = data.pop("backbone_type", "standard")
+            data["type"] = backbone_type
+            backbone_cls = (
+                T_Model
+                if backbone_type == "standard"
+                else T_Model.get_class_by_type(backbone_type)
+            )
+            backbone = backbone_cls.deserialize(data)
             return cls(atomic_model_=backbone.atomic_model, spin=spin)
 
     return NSM

deepmd/pt_expt/model/get_model.py

@@ -271,12 +271,6 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     data = copy.deepcopy(data)
     spin_cfg = data.pop("spin")
     data.setdefault("descriptor", {})
-    if str(data.get("bridging_method", "none")).lower() not in ("none", ""):
-        raise NotImplementedError(
-            "analytical bridging combined with the native spin scheme is a "
-            "follow-up (the bridged model is a linear composition; the "
-            "native-spin factory composes over a single standard model)"
-        )
     # Expand index/symbol forms of ``use_spin`` against ``type_map`` into the
     # per-type boolean list (pure; validates symbols).
     use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"])
@@ -306,8 +300,14 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
             "support (supports_native_spin()); descriptor type "
             f"{data['descriptor'].get('type')!r} does not accept `use_spin`"
         ) from err
-    descriptor = backbone_model.atomic_model.descriptor
-    if not descriptor.supports_native_spin():
+    # A bridged backbone is a LinearEnergyAtomicModel; the capability gate
+    # reads the LEARNED child's descriptor either way (shared helper).
+    from deepmd.dpmodel.model.model import (
+        learned_descriptor,
+    )
+
+    descriptor = learned_descriptor(backbone_model.atomic_model)
+    if descriptor is None or not descriptor.supports_native_spin():
         raise NotImplementedError(
             "spin scheme 'native' requires a descriptor declaring "
             "supports_native_spin()"

source/api_cc/tests/test_deepspin_dpa4_pairexcl_ptexpt.cc

@@ -72,9 +72,18 @@ class TestInferDeepSpinDpa4PairExclPtExpt : public ::testing::Test {
   std::vector<VALUETYPE> expected_f;
   std::vector<VALUETYPE> expected_fm;
   std::vector<VALUETYPE> expected_tot_v;
+  // NoPBC twin: an explicit nghost=0 InputNlist carries no periodic images,
+  // so the LAMMPS-nlist case below is the gas-phase system and must be
+  // compared against the gas-phase reference (same convention as
+  // test_deepspin_dpa4_graph_ptexpt.cc, which runs cpu_lmp_nlist only in its
+  // Nopbc fixture).
+  std::vector<VALUETYPE> expected_e_nopbc;
+  std::vector<VALUETYPE> expected_f_nopbc;
+  std::vector<VALUETYPE> expected_fm_nopbc;
 
   int natoms;
   double expected_tot_e;
+  double expected_tot_e_nopbc;
 
   deepmd::DeepSpin dp_excl;
   deepmd::DeepSpin dp_base;
@@ -98,14 +107,19 @@ class TestInferDeepSpinDpa4PairExclPtExpt : public ::testing::Test {
     expected_f = ref.get<VALUETYPE>("pbc", "expected_f");
     expected_fm = ref.get<VALUETYPE>("pbc", "expected_fm");
     expected_tot_v = ref.get<VALUETYPE>("pbc", "expected_tot_v");
+    expected_e_nopbc = ref.get<VALUETYPE>("nopbc", "expected_e");
+    expected_f_nopbc = ref.get<VALUETYPE>("nopbc", "expected_f");
+    expected_fm_nopbc = ref.get<VALUETYPE>("nopbc", "expected_fm");
 
     natoms = expected_e.size();
     EXPECT_EQ(natoms * 3, expected_f.size());
     EXPECT_EQ(natoms * 3, expected_fm.size());
     EXPECT_EQ(9, expected_tot_v.size());
     expected_tot_e = 0.;
+    expected_tot_e_nopbc = 0.;
     for (int ii = 0; ii < natoms; ++ii) {
       expected_tot_e += expected_e[ii];
+      expected_tot_e_nopbc += expected_e_nopbc[ii];
     }
   };
 
@@ -155,11 +169,13 @@ TYPED_TEST(TestInferDeepSpinDpa4PairExclPtExpt, cpu_lmp_nlist) {
   const std::vector<VALUETYPE>& coord = this->coord;
   const std::vector<VALUETYPE>& spin = this->spin;
   std::vector<int>& atype = this->atype;
-  std::vector<VALUETYPE>& box = this->box;
-  std::vector<VALUETYPE>& expected_f = this->expected_f;
-  std::vector<VALUETYPE>& expected_fm = this->expected_fm;
+  std::vector<VALUETYPE>& expected_f = this->expected_f_nopbc;
+  std::vector<VALUETYPE>& expected_fm = this->expected_fm_nopbc;
   int& natoms = this->natoms;
-  double& expected_tot_e = this->expected_tot_e;
+  double& expected_tot_e = this->expected_tot_e_nopbc;
+  // An nghost=0 InputNlist carries no periodic images, so this is the
+  // gas-phase system: empty box + the NoPBC reference.
+  std::vector<VALUETYPE> box = {};
 
   double ener;
   std::vector<VALUETYPE> force, force_mag, virial;