diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 28e1118718..aba6b9fd48 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -652,11 +652,29 @@ def _format_nlist( ret = xp_take_along_axis(ret, ret_mapping, axis=2) ret = xp.where(rr > rcut, -1, ret) ret = ret[..., :nnei] - # not extra_nlist_sort and n_nnei <= nnei: - elif n_nnei == nnei: - ret = nlist else: - pass + # not extra_nlist_sort and n_nnei <= nnei: no reordering is + # needed (these descriptors reduce over neighbors order- + # independently), but we must still drop neighbors beyond rcut. + # The C++/LAMMPS neighbor list is built with rcut+skin and is + # NOT rcut-filtered before forward_lower; without this, out-of- + # rcut neighbors leak into the descriptor whenever the per-atom + # neighbor count <= nnei (this branch), making the result + # order-dependent (see discussion #5438). + if n_nnei == nnei: + ret = nlist + # else (n_nnei < nnei): `ret` is already padded to nnei above. + n_nf, n_nloc, n_pad = ret.shape + m_real_nei = ret >= 0 + coord0 = xp_take_first_n(extended_coord, 1, n_nloc) + index = xp.tile( + xp.where(m_real_nei, ret, 0).reshape(n_nf, n_nloc * n_pad, 1), + (1, 1, 3), + ) + coord1 = xp_take_along_axis(extended_coord, index, axis=1) + coord1 = coord1.reshape(n_nf, n_nloc, n_pad, 3) + rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1) + ret = xp.where(m_real_nei & (rr > rcut), -1, ret) assert ret.shape[-1] == nnei return ret diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index 78705b153c..6f5e347e68 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -501,7 +501,26 @@ def _format_nlist( nlist = torch.where(rr > rcut, -1, nlist) nlist = nlist[..., :nnei] else: # not extra_nlist_sort and n_nnei <= nnei: - pass # great! + # No reordering is needed here (these descriptors reduce over + # neighbors order-independently), but we must still drop + # neighbors beyond rcut. The C++/LAMMPS neighbor list is built + # with rcut+skin and is NOT rcut-filtered before forward_lower; + # without this, out-of-rcut neighbors leak into the descriptor + # whenever the per-atom neighbor count <= nnei (this branch), + # making the result order-dependent (see discussion #5438). + n_nf, n_nloc, n_nnei = nlist.shape + m_real_nei = nlist >= 0 + coord0 = extended_coord[:, :n_nloc, :] + index = ( + torch.where(m_real_nei, nlist, 0) + .view(n_nf, n_nloc * n_nnei, 1) + .expand(-1, -1, 3) + ) + coord1 = torch.gather(extended_coord, 1, index).view( + n_nf, n_nloc, n_nnei, 3 + ) + rr = torch.linalg.norm(coord0[:, :, None, :] - coord1, dim=-1) + nlist = torch.where(m_real_nei & (rr > rcut), -1, nlist) assert nlist.shape[-1] == nnei return nlist diff --git a/source/tests/common/dpmodel/test_format_nlist_overcut.py b/source/tests/common/dpmodel/test_format_nlist_overcut.py new file mode 100644 index 0000000000..530066f59c --- /dev/null +++ b/source/tests/common/dpmodel/test_format_nlist_overcut.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""dpmodel counterpart of the pt ``test_format_nlist_overcut`` regression. + +``_format_nlist`` exists in both backends (``deepmd/pt`` and ``deepmd/dpmodel``; +the latter is shared by ``pt_expt``). The pad branch previously did not drop +out-of-``rcut`` neighbors, so an over-``rcut`` neighbor list (what the C++/LAMMPS +path passes, unfiltered) leaked into the descriptor and made ``call_lower`` +order-dependent. This test exercises the **dpmodel** (numpy) path directly -- +the pt test only covers ``deepmd/pt``. + +An over-``rcut`` nlist (``rcut + 2``, pad branch) is fed to ``call_lower`` both +as-is and reversed; both must match the canonical ``rcut``-bounded evaluation +(reduced + per-atom energy). Without the fix the reversed case diverges. +""" + +import unittest + +import numpy as np + +from deepmd.dpmodel.model.model import ( + get_model, +) +from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +model_se_r = { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "se_e2_r", + "sel": [46, 92, 4], + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [25, 50, 100], + "resnet_dt": False, + "seed": 1, + }, + "fitting_net": {"neuron": [24, 24, 24], "resnet_dt": True, "seed": 1}, + "data_stat_nbatch": 20, +} + + +class TestFormatNlistOvercutDP(unittest.TestCase): + def setUp(self) -> None: + self.model = get_model(model_se_r) + rng = np.random.default_rng(20240131) + self.natoms = 6 + self.cell = 6.0 * np.eye(3) + self.coord = 5.5 * rng.random([self.natoms, 3]) + self.atype = np.array([0, 0, 1, 1, 2, 2], dtype=np.int64) + + def _lower(self, rcut_build, reverse): + ec, ea, mp, nlist = extend_input_and_build_neighbor_list( + self.coord[None], + self.atype[None], + rcut_build, + sum(self.model.get_sel()), + mixed_types=True, + box=self.cell[None], + ) + if reverse: + nlist = nlist[..., ::-1] + out = self.model.call_lower(ec, ea, nlist, mp, do_atomic_virial=False) + return out["energy"], out["atom_energy"] + + def test_overcut_matches_canonical(self) -> None: + rcut = self.model.get_rcut() + er_ref, ea_ref = self._lower(rcut, reverse=False) + for reverse in (False, True): # over-cut nlist, kept order vs reversed + with self.subTest(reverse=reverse): + er, ea = self._lower(rcut + 2.0, reverse=reverse) + np.testing.assert_allclose(er, er_ref, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(ea, ea_ref, rtol=1e-10, atol=1e-10) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_format_nlist_overcut.py b/source/tests/pt/model/test_format_nlist_overcut.py new file mode 100644 index 0000000000..29fbf08d4c --- /dev/null +++ b/source/tests/pt/model/test_format_nlist_overcut.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression test for the ``_format_nlist`` pad-branch rcut filter. + +The C++/LAMMPS neighbor list is built with ``rcut + skin`` and is passed into +``forward_lower`` WITHOUT any rcut filtering (see ``DeepPotPT.cc`` / +``copy_from_nlist``). Its width is the per-atom neighbor count, which can be +``<= nnei`` (``sum(sel)``) on sparse systems -- exactly the case in discussion +#5438 (width 39 < 100). + +In that regime ``_format_nlist`` takes its pad branch (``n_nnei <= nnei`` and +``extra_nlist_sort`` is False). Previously that branch did NOT drop neighbors +beyond ``rcut``, so out-of-``rcut`` neighbors leaked into the descriptor, making +``forward_lower`` order-dependent (reversing the nlist changed the energy by +~1e-4 for se_r, ~4e-6 for se_a). The fix filters ``rr > rcut`` in the pad branch +too. + +This test feeds an over-``rcut`` nlist (``rcut + 2``, pad branch) -- once as-is +and once reversed -- and asserts both match the canonical ``rcut``-bounded +evaluation (energy and force) to machine precision. Without the fix the +over-cut and reversed cases diverge from the canonical reference. +""" + +import copy +import unittest + +import torch + +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + +from ...consistent.common import ( + parameterized, +) +from ...seed import ( + GLOBAL_SEED, +) +from .test_forward_lower import ( + reduce_tensor, +) +from .test_permutation import ( + model_se_e2_a, +) + +dtype = torch.float64 + +model_se_r = { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "se_e2_r", + "sel": [46, 92, 4], + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [25, 50, 100], + "resnet_dt": False, + "seed": 1, + }, + "fitting_net": {"neuron": [24, 24, 24], "resnet_dt": True, "seed": 1}, + "data_stat_nbatch": 20, +} + + +@parameterized( + ( + "se_a", + "se_r", + ), # descriptor flavour (se_a damps over-rcut via direction; se_r does not) + (False, True), # reverse the over-cut nlist before forward_lower +) +class TestFormatNlistOvercut(unittest.TestCase): + def setUp(self) -> None: + flavour, self.reverse = self.param + params = copy.deepcopy(model_se_e2_a if flavour == "se_a" else model_se_r) + self.model = get_model(params).to(env.DEVICE) + + def _make_system(self): + # sparse system: per-atom neighbor count stays below sum(sel), so the + # over-cut nlist exercises the pad branch of _format_nlist. + natoms = 6 + cell = 6.0 * torch.eye(3, dtype=dtype, device=env.DEVICE) + generator = torch.Generator(device=env.DEVICE).manual_seed(GLOBAL_SEED) + coord = 5.5 * torch.rand( + [natoms, 3], dtype=dtype, device=env.DEVICE, generator=generator + ) + atype = torch.tensor([0, 0, 1, 1, 2, 2], dtype=torch.int64, device=env.DEVICE) + return coord, atype, cell + + def _energy_force(self, coord, atype, cell, rcut_build, reverse): + ec, ea, mp, nlist = extend_input_and_build_neighbor_list( + coord.unsqueeze(0), + atype.unsqueeze(0), + rcut_build, + sum(self.model.get_sel()), + mixed_types=True, + box=cell.unsqueeze(0), + ) + if reverse: + nlist = torch.flip(nlist, dims=[-1]) + out = self.model.forward_lower(ec, ea, nlist, mp, do_atomic_virial=False) + force = reduce_tensor(out["extended_force"], mp, coord.shape[0]) + return out["energy"], force + + def test_overcut_matches_canonical(self) -> None: + coord, atype, cell = self._make_system() + rcut = self.model.get_rcut() + + # canonical: rcut-bounded nlist (no out-of-rcut neighbors), not reversed + e_ref, f_ref = self._energy_force(coord, atype, cell, rcut, reverse=False) + # over-cut nlist (rcut + 2 -> pad branch with out-of-rcut neighbors) + e_out, f_out = self._energy_force( + coord, atype, cell, rcut + 2.0, reverse=self.reverse + ) + + torch.testing.assert_close(e_out, e_ref, rtol=1e-10, atol=1e-10) + torch.testing.assert_close(f_out, f_ref, rtol=1e-10, atol=1e-10) + + +if __name__ == "__main__": + unittest.main()