fix(dpmodel): preserve virtual spin atom types

deepmd/dpmodel/model/spin_model.py

@@ -81,6 +81,20 @@ def _to_xp(self, arr: Any, xp: Any, ref_arr: Any) -> Any:
         """Convert a numpy array to the same namespace as ref_arr."""
         return xp.asarray(arr, device=array_api_compat.device(ref_arr))
 
+    def _lookup_type_values(self, values: Any, atype: Array, ref_arr: Array) -> Array:
+        """Gather per-type values while mapping virtual atom types to zero.
+
+        Negative atom types are padding placeholders, not Python-style indices
+        from the end of the type table. Their spin scale and mask must remain
+        zero until the backbone model applies its normal virtual-atom mask.
+        """
+        xp = array_api_compat.array_namespace(ref_arr)
+        values = self._to_xp(values, xp, ref_arr)
+        real_atom = atype >= 0
+        safe_atype = xp.where(real_atom, atype, xp.zeros_like(atype))
+        gathered = values[safe_atype]
+        return xp.where(real_atom, gathered, xp.zeros_like(gathered))
+
     def process_spin_input(
         self, coord: Array, atype: Array, spin: Array
     ) -> tuple[Array, Array, Array]:
@@ -97,9 +111,12 @@ def process_spin_input(
         """
         xp = array_api_compat.array_namespace(coord)
         nframes, nloc = coord.shape[:-1]
-        atype_spin = xp.concat([atype, atype + self.ntypes_real], axis=-1)
-        vsm = self._to_xp(self.virtual_scale_mask, xp, coord)
-        spin_dist = spin * xp.reshape(vsm[atype], (nframes, nloc, 1))
+        virtual_atype = xp.where(atype >= 0, atype + self.ntypes_real, atype)
+        atype_spin = xp.concat([atype, virtual_atype], axis=-1)
+        spin_dist = spin * xp.reshape(
+            self._lookup_type_values(self.virtual_scale_mask, atype, coord),
+            (nframes, nloc, 1),
+        )
         virtual_coord = coord + spin_dist
         coord_spin = xp.concat([coord, virtual_coord], axis=-2)
         # for spin virial correction
@@ -151,12 +168,18 @@ def process_spin_input_lower(
         xp = array_api_compat.array_namespace(extended_coord)
         nframes, nall = extended_coord.shape[:2]
         nloc = nlist.shape[1]
-        vsm = self._to_xp(self.virtual_scale_mask, xp, extended_coord)
         extended_spin_dist = extended_spin * xp.reshape(
-            vsm[extended_atype], (nframes, nall, 1)
+            self._lookup_type_values(
+                self.virtual_scale_mask, extended_atype, extended_coord
+            ),
+            (nframes, nall, 1),
         )
         virtual_extended_coord = extended_coord + extended_spin_dist
-        virtual_extended_atype = extended_atype + self.ntypes_real
+        virtual_extended_atype = xp.where(
+            extended_atype >= 0,
+            extended_atype + self.ntypes_real,
+            extended_atype,
+        )
         extended_coord_updated = self.concat_switch_virtual(
             extended_coord, virtual_extended_coord, nloc
         )
@@ -223,7 +246,10 @@ def process_spin_output(
             mask = self._to_xp(self.virtual_scale_mask, xp, out_tensor)
         else:
             mask = self._to_xp(self.spin_mask, xp, out_tensor)
-        atomic_mask = xp.reshape(mask[atype], (nframes, nloc, 1))
+        atomic_mask = xp.reshape(
+            self._lookup_type_values(mask, atype, out_tensor),
+            (nframes, nloc, 1),
+        )
         out_real, out_mag = out_tensor[:, :nloc], out_tensor[:, nloc:]
         if add_mag:
             out_real = out_real + out_mag
@@ -249,7 +275,10 @@ def process_spin_output_lower(
             mask = self._to_xp(self.virtual_scale_mask, xp, extended_out_tensor)
         else:
             mask = self._to_xp(self.spin_mask, xp, extended_out_tensor)
-        atomic_mask = xp.reshape(mask[extended_atype], (nframes, nall, 1))
+        atomic_mask = xp.reshape(
+            self._lookup_type_values(mask, extended_atype, extended_out_tensor),
+            (nframes, nall, 1),
+        )
         extended_out_real = xp.concat(
             [
                 extended_out_tensor[:, :nloc],
@@ -698,8 +727,10 @@ def call_common(
         if "mask_mag" not in model_ret:
             xp = array_api_compat.array_namespace(atype)
             nframes_m, nloc_m = atype.shape[:2]
-            vsm = self._to_xp(self.virtual_scale_mask, xp, atype)
-            atomic_mask = xp.reshape(vsm[atype], (nframes_m, nloc_m, 1))
+            atomic_mask = xp.reshape(
+                self._lookup_type_values(self.virtual_scale_mask, atype, atype),
+                (nframes_m, nloc_m, 1),
+            )
             model_ret["mask_mag"] = atomic_mask > 0.0
         return model_ret
 
@@ -881,8 +912,12 @@ def call_common_lower(
         if "mask_mag" not in model_ret:
             xp = array_api_compat.array_namespace(extended_atype)
             nall = extended_atype.shape[1]
-            vsm = self._to_xp(self.virtual_scale_mask, xp, extended_atype)
-            atomic_mask = xp.reshape(vsm[extended_atype], (nframes, nall, 1))
+            atomic_mask = xp.reshape(
+                self._lookup_type_values(
+                    self.virtual_scale_mask, extended_atype, extended_atype
+                ),
+                (nframes, nall, 1),
+            )
             model_ret["mask_mag"] = atomic_mask > 0.0
         return model_ret
 

source/tests/common/dpmodel/test_spin_model_virtual_types.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Regression tests for virtual placeholders in dpmodel SpinModel inputs."""
+
+import numpy as np
+import pytest
+
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
+from deepmd.dpmodel.model.model import (
+    get_model,
+)
+
+MODEL_CONFIG = {
+    "type_map": ["A", "B", "C"],
+    "descriptor": {
+        "type": "se_e2_a",
+        "sel": [4, 4, 4],
+        "rcut_smth": 0.5,
+        "rcut": 4.0,
+        "neuron": [3, 6],
+        "axis_neuron": 2,
+        "precision": "float64",
+        "type_one_side": True,
+        "seed": 1,
+    },
+    "fitting_net": {
+        "type": "ener",
+        "neuron": [5, 5],
+        "precision": "float64",
+        "seed": 1,
+    },
+    # Keep the final real type magnetic so an accidental ``mask[-1]`` lookup
+    # is observable instead of being hidden by a zero scale.
+    "spin": {"use_spin": [False, False, True], "virtual_scale": [0.5]},
+}
+
+
+@pytest.fixture
+def model():
+    """Build a small NumPy spin model with a magnetic final real type."""
+    return get_model(MODEL_CONFIG)
+
+
+def test_dense_spin_expansion_preserves_virtual_types(model) -> None:
+    """A dense placeholder and its spin partner must both stay virtual."""
+    coord = np.arange(9, dtype=np.float64).reshape(1, 3, 3)
+    atype = np.array([[0, -1, 2]], dtype=np.int32)
+    spin = np.ones_like(coord)
+
+    coord_updated, atype_updated, coord_corr = model.process_spin_input(
+        coord, atype, spin
+    )
+
+    np.testing.assert_array_equal(atype_updated, [[0, -1, 2, 3, -1, 5]])
+    # The placeholder's virtual partner has neither a displacement nor a
+    # virial correction, even though the final real type is magnetic.
+    np.testing.assert_array_equal(coord_updated[:, 4], coord[:, 1])
+    np.testing.assert_array_equal(coord_corr[:, 4], 0.0)
+
+    _, magnetic_output, magnetic_mask = model.process_spin_output(
+        atype, np.ones((1, 6, 3), dtype=np.float64)
+    )
+    np.testing.assert_array_equal(magnetic_output[:, 1], 0.0)
+    np.testing.assert_array_equal(magnetic_mask[:, 1], False)
+
+    prediction = model(coord, atype, spin, box=None)
+    np.testing.assert_array_equal(prediction["atom_energy"][:, 1], 0.0)
+    np.testing.assert_array_equal(prediction["mask_mag"][:, 1], False)
+
+    changed_coord = coord.copy()
+    changed_spin = spin.copy()
+    changed_coord[:, 1] += 100.0
+    changed_spin[:, 1] += 100.0
+    changed_prediction = model(changed_coord, atype, changed_spin, box=None)
+    np.testing.assert_allclose(changed_prediction["energy"], prediction["energy"])
+    np.testing.assert_allclose(
+        changed_prediction["atom_energy"][:, [0, 2]],
+        prediction["atom_energy"][:, [0, 2]],
+    )
+
+
+def test_lower_spin_expansion_preserves_virtual_types(model) -> None:
+    """Local and ghost placeholders remain negative in the switched layout."""
+    extended_coord = np.arange(12, dtype=np.float64).reshape(1, 4, 3)
+    extended_atype = np.array([[0, -1, 2, -1]], dtype=np.int32)
+    extended_spin = np.ones_like(extended_coord)
+    nlist = np.array([[[2, -1], [0, -1]]], dtype=np.int32)
+    mapping = np.array([[0, 1, 0, 1]], dtype=np.int32)
+
+    (
+        coord_updated,
+        atype_updated,
+        _,
+        _,
+        coord_corr,
+    ) = model.process_spin_input_lower(
+        extended_coord,
+        extended_atype,
+        extended_spin,
+        nlist,
+        mapping=mapping,
+    )
+
+    np.testing.assert_array_equal(atype_updated, [[0, -1, 3, -1, 2, -1, 5, -1]])
+    for real_index, virtual_index in ((1, 3), (3, 7)):
+        np.testing.assert_array_equal(
+            coord_updated[:, virtual_index], extended_coord[:, real_index]
+        )
+        np.testing.assert_array_equal(coord_corr[:, virtual_index], 0.0)
+
+    _, magnetic_output, magnetic_mask = model.process_spin_output_lower(
+        extended_atype,
+        np.ones((1, 8, 3), dtype=np.float64),
+        nloc=2,
+    )
+    np.testing.assert_array_equal(magnetic_output[:, [1, 3]], 0.0)
+    np.testing.assert_array_equal(magnetic_mask[:, [1, 3]], False)
+
+
+def test_virtual_type_lookup_supports_array_api_strict(model) -> None:
+    """The masked lookup must not rely on NumPy negative-index semantics."""
+    xp = pytest.importorskip("array_api_strict")
+    coord = xp.asarray(np.arange(9, dtype=np.float64).reshape(1, 3, 3))
+    atype = xp.asarray(np.array([[0, -1, 2]], dtype=np.int64))
+    spin = xp.ones_like(coord)
+
+    coord_updated, atype_updated, coord_corr = model.process_spin_input(
+        coord, atype, spin
+    )
+    np.testing.assert_array_equal(to_numpy_array(atype_updated), [[0, -1, 2, 3, -1, 5]])
+    np.testing.assert_array_equal(to_numpy_array(coord_updated)[:, 4], [[3, 4, 5]])
+    np.testing.assert_array_equal(to_numpy_array(coord_corr)[:, 4], 0.0)