fix(dpmodel): exclude virtual atoms from neighbor stats

deepmd/dpmodel/utils/neighbor_stat.py

@@ -86,23 +86,38 @@ def call(
             - xp.reshape(coord0, (nframes, -1, 3))[:, :, None, :]
         )
         assert list(diff.shape) == [nframes, nloc, nall, 3]
-        # remove the diagonal elements
-        mask = xp.eye(nloc, nall, dtype=xp.bool, device=array_api_compat.device(diff))
-        mask = xp.tile(mask[None, :, :, None], (nframes, 1, 1, 3))
-        diff = xp.where(mask, xp.full_like(diff, xp.inf), diff)
+        # A valid statistics pair must contain two real atoms.  In particular,
+        # virtual centers must not contribute an artificially large neighbor row,
+        # and virtual neighbors must not drive the minimum distance to zero.
+        self_pair = xp.eye(
+            nloc, nall, dtype=xp.bool, device=array_api_compat.device(diff)
+        )
+        real_center = atype >= 0
+        real_neighbor = extend_atype >= 0
+        valid_pair = (
+            ~self_pair[None, :, :] & real_center[:, :, None] & real_neighbor[:, None, :]
+        )
         rr2 = xp.sum(xp.square(diff), axis=-1)
+        rr2 = xp.where(valid_pair, rr2, xp.full_like(rr2, xp.inf))
         min_rr2 = xp.min(rr2, axis=-1)
         # count the number of neighbors
+        within_rcut = valid_pair & (rr2 < self.rcut**2)
         if not self.mixed_types:
-            mask = rr2 < self.rcut**2
             nneis = []
             for ii in range(self.ntypes):
-                nneis.append(xp.sum(mask & (extend_atype == ii)[:, None, :], axis=-1))
+                nneis.append(
+                    xp.sum(
+                        xp.astype(
+                            within_rcut & (extend_atype == ii)[:, None, :],
+                            extend_atype.dtype,
+                        ),
+                        axis=-1,
+                    )
+                )
             nnei = xp.stack(nneis, axis=-1)
         else:
-            mask = rr2 < self.rcut**2
-            # virtual type (<0) are not counted
-            nnei = xp.sum(mask & (extend_atype >= 0)[:, None, :], axis=-1)
+            # Array API reductions accept numeric rather than boolean inputs.
+            nnei = xp.sum(xp.astype(within_rcut, extend_atype.dtype), axis=-1)
             nnei = xp.reshape(nnei, (nframes, nloc, 1))
         max_nnei = xp.max(nnei, axis=1)
         return min_rr2, max_nnei

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

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import array_api_strict as xp
+
+from deepmd.dpmodel.utils.neighbor_stat import (
+    NeighborStatOP,
+)
+
+from .utils import (
+    ArrayAPITest,
+)
+
+
+class TestNeighborStatOP(unittest.TestCase, ArrayAPITest):
+    def test_virtual_atoms_are_masked_before_reductions(self) -> None:
+        """Virtual-pair masking and neighbor reductions follow the Array API."""
+        coord = xp.reshape(
+            xp.asarray(
+                [
+                    [0.0, 0.0, 0.0],
+                    [0.0, 0.0, 0.0],
+                    [1.0, 0.0, 0.0],
+                    [3.0, 0.0, 0.0],
+                ],
+                dtype=xp.float64,
+            ),
+            (1, -1),
+        )
+        atype = xp.asarray([[0, -1, 0, 1]], dtype=xp.int64)
+        expected_min_rr2 = xp.asarray([[1.0, xp.inf, 1.0, 4.0]], dtype=xp.float64)
+
+        for mixed_types in (False, True):
+            with self.subTest(mixed_types=mixed_types):
+                min_rr2, max_nnei = NeighborStatOP(2, 1.1, mixed_types).call(
+                    coord, atype, None
+                )
+                expected_max_nnei = xp.asarray(
+                    [[1]] if mixed_types else [[1, 0]], dtype=xp.int64
+                )
+
+                self.assertTrue(bool(xp.all(min_rr2 == expected_min_rr2)))
+                self.assertTrue(bool(xp.all(max_nnei == expected_max_nnei)))
+                self.assert_namespace_equal(min_rr2, coord)
+                self.assert_namespace_equal(max_nnei, atype)
+                self.assert_device_equal(min_rr2, coord)
+                self.assert_device_equal(max_nnei, atype)
+                self.assert_dtype_equal(min_rr2, coord)
+                self.assert_dtype_equal(max_nnei, atype)
+                self.assertEqual(min_rr2.shape, (1, 4))
+                self.assertEqual(max_nnei.shape, expected_max_nnei.shape)

source/tests/common/dpmodel/test_neighbor_stat.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.dpmodel.utils.neighbor_stat import (
+    NeighborStatOP,
+)
+
+
+class TestNeighborStatOP(unittest.TestCase):
+    def test_virtual_atoms_do_not_affect_statistics(self) -> None:
+        """Ignore virtual atoms as both neighbor-stat centers and neighbors."""
+        # Atom 1 is virtual and overlaps atom 0.  Without a neighbor mask it
+        # drives the minimum distance to zero; without a center mask it sees both
+        # type-0 atoms and inflates their maximum neighbor count from one to two.
+        coord = np.array(
+            [
+                [0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0],
+                [3.0, 0.0, 0.0],
+            ],
+            dtype=np.float64,
+        ).reshape(1, -1)
+        atype = np.array([[0, -1, 0, 1]], dtype=np.int64)
+        expected_min_rr2 = np.array([[1.0, np.inf, 1.0, 4.0]])
+
+        for cell in (None, 10.0 * np.eye(3).reshape(1, 9)):
+            for mixed_types in (False, True):
+                with self.subTest(cell=cell is not None, mixed_types=mixed_types):
+                    min_rr2, max_nnei = NeighborStatOP(
+                        ntypes=2,
+                        rcut=1.1,
+                        mixed_types=mixed_types,
+                    ).call(coord, atype, cell)
+
+                    np.testing.assert_allclose(min_rr2, expected_min_rr2)
+                    expected_max_nnei = [[1]] if mixed_types else [[1, 0]]
+                    np.testing.assert_array_equal(max_nnei, expected_max_nnei)