fix(pt,dpmodel): drop out-of-rcut neighbors in _format_nlist pad branch

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
 

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
 

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()

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()