diff --git a/deepmd/dpmodel/utils/neighbor_stat.py b/deepmd/dpmodel/utils/neighbor_stat.py index 0ccb419d68..d340600496 100644 --- a/deepmd/dpmodel/utils/neighbor_stat.py +++ b/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 diff --git a/source/tests/common/dpmodel/array_api/test_neighbor_stat.py b/source/tests/common/dpmodel/array_api/test_neighbor_stat.py new file mode 100644 index 0000000000..fa20bfde2b --- /dev/null +++ b/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) diff --git a/source/tests/common/dpmodel/test_neighbor_stat.py b/source/tests/common/dpmodel/test_neighbor_stat.py new file mode 100644 index 0000000000..8dcaa42d3c --- /dev/null +++ b/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)