fix(dpmodel): preserve virtual spin atom types (deepmodeling#5854)

Closes deepmodeling#5663.

## Summary

- preserve negative placeholder types when creating dense and
lower-interface spin partners
- use one array-API-safe lookup that maps virtual atom types to zero
instead of applying Python negative indexing
- zero placeholder spin displacements, virial corrections, magnetic
outputs, and magnetic masks
- add dense, lower, end-to-end invariance, and Array API strict
regressions

## Why existing tests missed this

Existing spin consistency tests use only nonnegative real atom types,
including their lower-interface ghost atoms. Negative-type tests cover
base atomic models and neighbor-list builders, where masking happens
before type-dependent work, but not `SpinModel` preprocessing, which
runs earlier. Cross-backend spin tests therefore never supplied a
nonzero padded spin or asserted that both the real placeholder and its
generated spin partner remained negative.

## Validation

- `pytest source/tests/common/dpmodel/test_spin_model_virtual_types.py
source/tests/common/dpmodel/test_finetune_spin.py
source/tests/common/dpmodel/test_spin_model_legacy_routing.py -q` (12
passed)
- `ruff format .` (1664 files unchanged on final pass)
- `ruff check .` (passed)
- `git diff --check` (passed)

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Bug Fixes**
- Improved spin-model handling of padded and virtual atom placeholders.
- Prevented placeholder atoms from receiving displacement, virial,
scaling, or magnetic-output corrections.
- Ensured consistent behavior across standard and lower-level spin
processing paths.
- Improved compatibility across supported array backends and processing
devices.

- **Tests**
- Added regression coverage for virtual placeholders, masking, output
handling, and consistency between model backends.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz-bot

deepmd/dpmodel/model/spin_model.py

@@ -83,6 +83,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]:
@@ -99,9 +113,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
@@ -153,12 +170,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
         )
@@ -224,9 +247,18 @@ def process_spin_output(
         if virtual_scale:
             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))
-        out_real, out_mag = out_tensor[:, :nloc], out_tensor[:, nloc:]
+            # spin_mask is integral; it multiplies out_mag below, and the array
+            # API does not promote across kinds.
+            mask = xp.astype(
+                self._to_xp(self.spin_mask, xp, out_tensor), out_tensor.dtype
+            )
+        atomic_mask = xp.reshape(
+            self._lookup_type_values(mask, atype, out_tensor),
+            (nframes, nloc, 1),
+        )
+        # Trailing ellipsis: the array API does not specify numpy's implicit
+        # expansion of a partial multi-axis index.
+        out_real, out_mag = out_tensor[:, :nloc, ...], out_tensor[:, nloc:, ...]
         if add_mag:
             out_real = out_real + out_mag
         out_mag = xp.reshape(
@@ -250,19 +282,27 @@ def process_spin_output_lower(
         if virtual_scale:
             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))
+            # spin_mask is integral; it multiplies extended_out_mag below, and
+            # the array API does not promote across kinds.
+            mask = xp.astype(
+                self._to_xp(self.spin_mask, xp, extended_out_tensor),
+                extended_out_tensor.dtype,
+            )
+        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],
-                extended_out_tensor[:, nloc + nloc : nloc + nall],
+                extended_out_tensor[:, :nloc, ...],
+                extended_out_tensor[:, nloc + nloc : nloc + nall, ...],
             ],
             axis=1,
         )
         extended_out_mag = xp.concat(
             [
-                extended_out_tensor[:, nloc : nloc + nloc],
-                extended_out_tensor[:, nloc + nall :],
+                extended_out_tensor[:, nloc : nloc + nloc, ...],
+                extended_out_tensor[:, nloc + nall :, ...],
             ],
             axis=1,
         )
@@ -700,8 +740,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
 
@@ -883,8 +925,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
 

deepmd/pt/model/model/spin_model.py

@@ -91,21 +91,29 @@ def _pack_spin_stat_sample(
 
 def _lookup_type_values(values: torch.Tensor, atype: torch.Tensor) -> torch.Tensor:
     """
-    Gather one scalar value per atom type.
-
-    ``values[atype]`` is semantically equivalent, but AOTInductor may lower
-    that advanced-indexing form to a CUDA ``index.Tensor`` shim even for a CPU
-    ``.pt2`` package. ``index_select`` keeps the exported spin graph device
-    stable while preserving the same lookup semantics.
-
-    Padding ghost slots carry ``atype == -1`` (batched extended regions are
-    padded to a uniform ``nall``). Unlike advanced indexing, ``index_select``
-    rejects negative indices, so the padding entries are clamped to row 0; their
-    looked-up value is irrelevant because padding atoms carry zero spin and are
-    dropped from the per-local output downstream.
+    Gather one scalar value per atom type, mapping virtual atom types to zero.
+
+    ``values[atype]`` is semantically equivalent for real atoms, but
+    AOTInductor may lower that advanced-indexing form to a CUDA
+    ``index.Tensor`` shim even for a CPU ``.pt2`` package. ``index_select``
+    keeps the exported spin graph device stable.
+
+    Padding slots carry ``atype == -1``: ``deepmd/utils/data.py`` appends it as
+    the virtual-atom padding for mixed-type systems, and batched extended
+    regions are padded to a uniform ``nall``. Those are placeholders, not
+    Python-style indices from the end of the type table, so they get zero
+    rather than row 0's value — otherwise a padded slot picks up a real spin
+    scale and mask whenever type 0 is magnetic. This matches
+    ``SpinModel._lookup_type_values`` in ``deepmd/dpmodel/model/spin_model.py``.
     """
-    flat_atype = torch.clamp_min(atype.reshape(-1).to(dtype=torch.long), 0)
-    return torch.index_select(values.to(atype.device), 0, flat_atype).view(atype.shape)
+    long_atype = atype.to(dtype=torch.long)
+    real_atom = long_atype >= 0
+    # index_select rejects negative indices, unlike advanced indexing.
+    flat_atype = torch.clamp_min(long_atype.reshape(-1), 0)
+    gathered = torch.index_select(values.to(atype.device), 0, flat_atype).view(
+        atype.shape
+    )
+    return torch.where(real_atom, gathered, torch.zeros_like(gathered))
 
 
 class SpinModel(torch.nn.Module):
@@ -140,7 +148,12 @@ def process_spin_input(
         nframes, nloc = atype.shape
         coord = coord.reshape(nframes, nloc, 3)
         spin = spin.reshape(nframes, nloc, 3)
-        atype_spin = torch.concat([atype, atype + self.ntypes_real], dim=-1)
+        # Keep virtual placeholders at -1 instead of offsetting them into a
+        # real type of the spin half of the type table.
+        virtual_atype = torch.where(
+            atype >= 0, atype + self.ntypes_real, torch.full_like(atype, -1)
+        )
+        atype_spin = torch.concat([atype, virtual_atype], dim=-1)
         # spin_dist = s_i * \mu_i
         spin_dist = spin * _lookup_type_values(
             self.virtual_scale_mask,
@@ -193,7 +206,11 @@ def process_spin_input_lower(
             extended_atype,
         ).reshape([nframes, nall, 1])
         virtual_extended_coord = extended_coord + extended_spin_dist
-        virtual_extended_atype = extended_atype + self.ntypes_real
+        virtual_extended_atype = torch.where(
+            extended_atype >= 0,
+            extended_atype + self.ntypes_real,
+            torch.full_like(extended_atype, -1),
+        )
         extended_coord_updated = concat_switch_virtual(
             extended_coord, virtual_extended_coord, nloc
         )