fix(dpmodel): mask virtual EnvMat centers

deepmd/dpmodel/utils/env_mat.py

@@ -172,7 +172,11 @@ def call(
         Parameters
         ----------
         nlist
-            The neighbor list. shape: nf x nloc x nnei
+            The neighbor list. shape: nf x nloc x nnei.  Entries equal to ``-1``
+            mark empty neighbor slots, and a virtual center (whose atom type is
+            negative) must have an entire neighbor row of ``-1``; the in-tree
+            builder (``deepmd.dpmodel.utils.nlist.build_neighbor_list``) fills
+            the full row of a virtual atom with ``-1`` by construction.
         coord_ext
             The extended coordinates of atoms. shape: nf x (nallx3)
         atype_ext
@@ -200,10 +204,34 @@ def call(
         em, diff, sw = self._call(nlist, coord_ext, radial_only)
         nf, nloc, nnei = nlist.shape
         atype = xp_take_first_n(atype_ext, 1, nloc)
+        center_is_real = atype >= 0
+        # Virtual atoms use a negative type sentinel.  Never pass that sentinel to
+        # ``take``: NumPy treats -1 as the final real type, while stricter array
+        # namespaces may reject it.  Type zero is only a safe placeholder because
+        # the gathered rows are neutralized below.
+        safe_atype = xp.where(center_is_real, atype, xp.zeros_like(atype))
+        center_mask = xp.reshape(center_is_real, (nf, nloc, 1, 1))
+        # ``_make_env_mat`` already zeroes em, diff and sw wherever ``nlist < 0``,
+        # so a virtual center -- whose neighbor row is empty by the neighbor-list
+        # contract -- leaves this function at zero as long as normalization does
+        # not shift it.  Neutralizing the offset and the scale is therefore the
+        # whole fix; masking em/diff/sw again afterwards would make the
+        # descriptor depend on ``atype_ext``, which the compiled pt_expt DPA2
+        # lower miscompiles into wrong forces.
         if davg is not None:
-            em -= xp.reshape(xp.take(davg, xp.reshape(atype, (-1,)), axis=0), em.shape)
+            center_avg = xp.reshape(
+                xp.take(davg, xp.reshape(safe_atype, (-1,)), axis=0), em.shape
+            )
+            center_avg = xp.where(center_mask, center_avg, xp.zeros_like(center_avg))
+            em -= center_avg
         if dstd is not None:
-            em /= xp.reshape(xp.take(dstd, xp.reshape(atype, (-1,)), axis=0), em.shape)
+            center_std = xp.reshape(
+                xp.take(dstd, xp.reshape(safe_atype, (-1,)), axis=0), em.shape
+            )
+            # A neutral scale avoids hidden divide-by-zero/NaN values in the
+            # masked branch, which is important for differentiable backends.
+            center_std = xp.where(center_mask, center_std, xp.ones_like(center_std))
+            em /= center_std
         return em, diff, sw
 
     def _call(

source/tests/common/dpmodel/array_api/test_env_mat.py

@@ -4,6 +4,7 @@
 import array_api_strict as xp
 
 from deepmd.dpmodel.utils.env_mat import (
+    EnvMat,
     compute_smooth_weight,
 )
 
@@ -23,3 +24,53 @@ def test_compute_smooth_weight(self) -> None:
         self.assert_namespace_equal(w, d)
         self.assert_device_equal(w, d)
         self.assert_dtype_equal(w, d)
+
+    def test_virtual_center_uses_safe_normalization_indices(self) -> None:
+        """Strict array indexing must never receive the negative type sentinel."""
+        coord = xp.asarray([[0.0, 0.0, 0.0]], dtype=xp.float64)
+        atype = xp.asarray([[-1]], dtype=xp.int64)
+        nlist = xp.asarray([[[-1]]], dtype=xp.int64)
+        davg = xp.asarray([[[11.0, 13.0, 17.0, 19.0]]], dtype=xp.float64)
+        # A zero placeholder scale verifies that masking happens before division;
+        # otherwise the virtual row can create hidden NaN or infinity values.
+        dstd = xp.zeros_like(davg)
+
+        env_mat, diff, switch = EnvMat(2.0, 0.5).call(coord, atype, nlist, davg, dstd)
+
+        for output in (env_mat, diff, switch):
+            self.assertTrue(bool(xp.all(output == xp.zeros_like(output))))
+            self.assert_namespace_equal(output, coord)
+            self.assert_device_equal(output, coord)
+            self.assert_dtype_equal(output, coord)
+
+    def test_mixed_centers_keep_virtual_rows_at_zero(self) -> None:
+        """A real center is normalized; a virtual one is left untouched at zero."""
+        coord = xp.asarray([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0]], dtype=xp.float64)
+        atype = xp.asarray([[0, -1]], dtype=xp.int64)
+        # The virtual center's neighbor row is empty, per the neighbor-list
+        # contract, so _make_env_mat leaves its outputs at zero. Only the
+        # normalization below could shift them off zero.
+        nlist = xp.asarray([[[1], [-1]]], dtype=xp.int64)
+        # The last row is what an unguarded ``take`` selects for atype -1, so
+        # keep it nonzero: borrowing it would shift the virtual row off zero.
+        davg = xp.asarray(
+            [[[0.0, 0.0, 0.0, 0.0]], [[11.0, 13.0, 17.0, 19.0]]],
+            dtype=xp.float64,
+        )
+        dstd = xp.asarray(
+            [[[1.0, 1.0, 1.0, 1.0]], [[2.0, 2.0, 2.0, 2.0]]],
+            dtype=xp.float64,
+        )
+
+        env_mat, diff, switch = EnvMat(2.0, 0.5).call(coord, atype, nlist, davg, dstd)
+
+        for output in (env_mat, diff, switch):
+            real_output = output[:, :1, ...]
+            virtual_output = output[:, 1:, ...]
+            self.assertTrue(bool(xp.any(real_output != xp.zeros_like(real_output))))
+            self.assertTrue(
+                bool(xp.all(virtual_output == xp.zeros_like(virtual_output)))
+            )
+            self.assert_namespace_equal(output, coord)
+            self.assert_device_equal(output, coord)
+            self.assert_dtype_equal(output, coord)

source/tests/common/dpmodel/test_env_mat.py

@@ -12,6 +12,7 @@
 )
 from .case_single_frame_with_nlist import (
     TestCaseSingleFrameWithNlist,
+    TestCaseSingleFrameWithNlistWithVirtual,
 )
 
 
@@ -38,3 +39,36 @@ def test_self_consistency(
         np.testing.assert_allclose(mm0, mm1)
         np.testing.assert_allclose(diff0, diff1)
         np.testing.assert_allclose(ww0, ww1)
+
+
+class TestEnvMatWithVirtualCenter(
+    unittest.TestCase, TestCaseSingleFrameWithNlistWithVirtual
+):
+    def setUp(self) -> None:
+        TestCaseSingleFrameWithNlistWithVirtual.setUp(self)
+
+    def test_normalization_keeps_virtual_centers_zero(self) -> None:
+        """Virtual centers must not borrow normalization data from a real type."""
+        nf, nloc, nnei = self.nlist.shape
+        virtual_center = self.atype_ext[:, :nloc] < 0
+
+        for radial_only, width in ((False, 4), (True, 1)):
+            with self.subTest(radial_only=radial_only):
+                # Nonzero values make accidental ``-1`` indexing observable: NumPy
+                # would otherwise select the final real-type row silently.
+                davg = np.arange(1, self.nt * nnei * width + 1, dtype=np.float64)
+                davg = davg.reshape(self.nt, nnei, width)
+                dstd = np.full_like(davg, 2.0)
+
+                env_mat, diff, switch = EnvMat(self.rcut, self.rcut_smth).call(
+                    self.coord_ext,
+                    self.atype_ext,
+                    self.nlist,
+                    davg,
+                    dstd,
+                    radial_only=radial_only,
+                )
+
+                np.testing.assert_allclose(env_mat[virtual_center], 0.0)
+                np.testing.assert_allclose(diff[virtual_center], 0.0)
+                np.testing.assert_allclose(switch[virtual_center], 0.0)