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
         )
@@ -222,9 +245,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(
@@ -248,19 +280,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,
         )
@@ -698,8 +738,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 +923,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
         )