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

deepmd/pd/utils/neighbor_stat.py

@@ -85,30 +85,38 @@ def forward(
             1
         ) - coord0.reshape([nframes, -1, 3]).unsqueeze(2)
         assert list(diff.shape) == [nframes, nloc, nall, 3]
-        # remove the diagonal elements
-        mask = paddle.eye(nloc, nall).to(dtype=paddle.bool, device=diff.place)
-        # diff[:, mask] = float("inf")
-        # diff.masked_fill_(
-        #     paddle.broadcast_to(mask.unsqueeze([0, -1]), diff.shape),
-        #     paddle.to_tensor(float("inf")),
-        # )
-        diff[paddle.broadcast_to(mask.unsqueeze([0, -1]), diff.shape)] = float("inf")
+        # Match the generic neighbor-stat operator: virtual atoms cannot be
+        # statistics centers or neighbors, and self pairs are always excluded.
+        self_pair = paddle.eye(nloc, nall).to(dtype=paddle.bool, device=diff.place)
+        real_center = atype >= 0
+        real_neighbor = extend_atype >= 0
+        valid_pair = (
+            ~self_pair.unsqueeze(0)
+            & real_center.unsqueeze(2)
+            & real_neighbor.unsqueeze(1)
+        )
         rr2 = paddle.sum(paddle.square(diff), axis=-1)
+        rr2 = paddle.where(valid_pair, rr2, paddle.full_like(rr2, float("inf")))
         min_rr2 = paddle.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
-            nnei = paddle.zeros([nframes, nloc, self.ntypes], dtype=paddle.int64)
-            for ii in range(self.ntypes):
-                nnei[:, :, ii] = paddle.sum(
-                    mask & ((extend_atype == ii)[:, None, :]), axis=-1
-                )
+            nnei = paddle.stack(
+                [
+                    paddle.sum(
+                        (within_rcut & (extend_atype == ii).unsqueeze(1)).astype(
+                            extend_atype.dtype
+                        ),
+                        axis=-1,
+                    )
+                    for ii in range(self.ntypes)
+                ],
+                axis=-1,
+            )
         else:
-            mask = rr2 < self.rcut**2
-            # virtual types (<0) are not counted
-            nnei = paddle.sum(
-                mask & ((extend_atype >= 0).unsqueeze(1)), axis=-1
-            ).reshape([nframes, nloc, 1])
+            nnei = paddle.sum(within_rcut.astype(extend_atype.dtype), axis=-1).reshape(
+                [nframes, nloc, 1]
+            )
         max_nnei = paddle.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)

source/tests/pd/test_neighbor_stat.py

@@ -4,10 +4,17 @@
 
 import dpdata
 import numpy as np
+import paddle
 
 from deepmd.entrypoints.neighbor_stat import (
     neighbor_stat,
 )
+from deepmd.pd.utils.env import (
+    DEVICE,
+)
+from deepmd.pd.utils.neighbor_stat import (
+    NeighborStatOP,
+)
 
 from ..seed import (
     GLOBAL_SEED,
@@ -67,3 +74,40 @@ def test_neighbor_stat(self):
                     if not mixed_type:
                         ret.append(0)
                     np.testing.assert_array_equal(max_nbor_size, ret)
+
+
+class TestNeighborStatOP(unittest.TestCase):
+    def test_virtual_atoms_do_not_affect_statistics(self) -> None:
+        """Ignore virtual atoms as both statistics centers and neighbors."""
+        # Atom 1 is virtual and overlaps atom 0. Without both masks, it either
+        # sets the minimum distance to zero or inflates the maximum type-0 count.
+        coord = paddle.to_tensor(
+            [
+                [
+                    [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=paddle.float64,
+            place=DEVICE,
+        ).reshape([1, -1])
+        atype = paddle.to_tensor([[0, -1, 0, 1]], dtype=paddle.int64, place=DEVICE)
+        expected_min_rr2 = np.array([[1.0, np.inf, 1.0, 4.0]])
+
+        for cell in (
+            None,
+            10.0 * paddle.eye(3, dtype=paddle.float64).reshape([1, 9]).to(DEVICE),
+        ):
+            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,
+                    )(coord, atype, cell)
+
+                    np.testing.assert_allclose(min_rr2.numpy(), expected_min_rr2)
+                    expected_max_nnei = [[1]] if mixed_types else [[1, 0]]
+                    np.testing.assert_array_equal(max_nnei.numpy(), expected_max_nnei)