From 7f5b0377e7abe3e45cd3997c7ddf9b65b53bc04a Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 08:52:15 +0800 Subject: [PATCH 1/4] test(pt_expt): document compiled neighbor type equivalence Clarify that compiled training intentionally passes merged global candidates because the lower model performs the same non-mixed type layout. Add an adversarial A/A/B regression proving early and lower-layer type splits produce the same final neighbor list under the established contract. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/pt_expt/train/training.py | 7 ++++ source/tests/common/dpmodel/test_nlist.py | 46 +++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 08218835d7..2dfd3aa1ac 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -924,6 +924,13 @@ def forward( nloc, rcut, sel, + # Keep the candidate list merged, matching eager training's + # DefaultNeighborList contract. forward_common_lower always calls + # model.format_nlist, which performs the type split for non-mixed + # descriptors. The shared builder globally truncates to sum(sel) + # before either split, so distinguish_types=True would only move + # the same layout transform earlier without changing the final + # formatted neighbor list consumed by the lower model. distinguish_types=False, # model-level pair exclusion is a nlist-BUILD transform (decision # #18/A4); the compiled dense lower consumes a pre-excluded nlist. diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py index 515c3d48c5..f23473670f 100644 --- a/source/tests/common/dpmodel/test_nlist.py +++ b/source/tests/common/dpmodel/test_nlist.py @@ -21,6 +21,7 @@ extend_coord_with_ghosts, get_multiple_nlist_key, inter2phys, + nlist_distinguish_types, ) @@ -154,6 +155,51 @@ def test_nlist_lt(self) -> None: np.testing.assert_allclose(self.expected_nlist, nlist1) +class TestNeighborTypeLayout(unittest.TestCase): + """Document global candidate selection followed by per-type layout.""" + + def test_lower_type_split_matches_early_type_split(self) -> None: + """An earlier type split must preserve the final formatted nlist.""" + # The two closest candidates are type 0, while the farther type-1 atom + # is still inside rcut. The established contract first keeps the global + # sum(sel) nearest candidates and only then lays them out by type. + coord = np.array( + [ + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [1.5, 0.0, 0.0], + [2.0, 0.0, 0.0], + ] + ], + dtype=np.float64, + ) + atype = np.array([[0, 0, 0, 1]], dtype=np.int64) + sel = [1, 1] + + merged = build_neighbor_list( + coord, + atype, + nloc=1, + rcut=3.0, + sel=sel, + distinguish_types=False, + ) + lower_formatted = nlist_distinguish_types(merged, atype, sel) + early_formatted = build_neighbor_list( + coord, + atype, + nloc=1, + rcut=3.0, + sel=sel, + distinguish_types=True, + ) + + np.testing.assert_array_equal(merged[0, 0], [1, 2]) + np.testing.assert_array_equal(lower_formatted, early_formatted) + np.testing.assert_array_equal(lower_formatted[0, 0], [1, -1]) + + dtype = np.float64 From c17f783829ea05e5fd9a5eeae58e9c080673b6c4 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 1 Aug 2026 22:18:10 +0800 Subject: [PATCH 2/4] test(dpmodel): broaden neighbor type layout coverage Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/common/dpmodel/test_nlist.py | 126 +++++++++++++++------- 1 file changed, 85 insertions(+), 41 deletions(-) diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py index f23473670f..8dddc189cb 100644 --- a/source/tests/common/dpmodel/test_nlist.py +++ b/source/tests/common/dpmodel/test_nlist.py @@ -154,50 +154,94 @@ def test_nlist_lt(self) -> None: ) np.testing.assert_allclose(self.expected_nlist, nlist1) - -class TestNeighborTypeLayout(unittest.TestCase): - """Document global candidate selection followed by per-type layout.""" - def test_lower_type_split_matches_early_type_split(self) -> None: - """An earlier type split must preserve the final formatted nlist.""" - # The two closest candidates are type 0, while the farther type-1 atom - # is still inside rcut. The established contract first keeps the global - # sum(sel) nearest candidates and only then lays them out by type. - coord = np.array( - [ - [ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [1.5, 0.0, 0.0], - [2.0, 0.0, 0.0], - ] - ], - dtype=np.float64, - ) - atype = np.array([[0, 0, 0, 1]], dtype=np.int64) - sel = [1, 1] + """Global candidate selection and early type splits must be equivalent. - merged = build_neighbor_list( - coord, - atype, - nloc=1, - rcut=3.0, - sel=sel, - distinguish_types=False, - ) - lower_formatted = nlist_distinguish_types(merged, atype, sel) - early_formatted = build_neighbor_list( - coord, - atype, - nloc=1, - rcut=3.0, - sel=sel, - distinguish_types=True, - ) + The matrix covers same-type truncation, real batch axes, multiple local + atoms, ghosts, exact distance ties, virtual atoms, and padded buckets. + Keeping it on the shared formatting fixture makes the contract visible + beside the other short/equal/long neighbor-list cases. + """ + cases = { + "same_type_nearest": { + "coord": np.array( + [ + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [1.5, 0.0, 0.0], + [2.5, 0.0, 0.0], + [2.0, 0.0, 0.0], + ] + ], + dtype=np.float64, + ), + "atype": np.array([[0, 0, 0, 0, 1]], dtype=np.int64), + "nloc": 1, + "sel": [2, 1], + "rcut": 3.0, + "expected": np.array([[[1, 2, 4]]], dtype=np.int64), + }, + "batched_ghost_ties_virtual_padding": { + "coord": np.array( + [ + [ + [0.0, 0.0, 0.0], + [0.0, 2.0, 0.0], + [1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.5, 0.5, 0.0], + ], + [ + [0.0, 0.0, 0.0], + [0.0, 2.5, 0.0], + [1.2, 0.0, 0.0], + [-1.2, 0.0, 0.0], + [0.0, 1.1, 0.0], + [0.4, 0.4, 0.0], + ], + ], + dtype=np.float64, + ), + "atype": np.array( + [[0, 1, 0, 0, 1, -1], [0, 1, 0, 0, 1, -1]], + dtype=np.int64, + ), + "nloc": 2, + "sel": [3, 2], + "rcut": 3.0, + "expected": None, + }, + } + + for name, case in cases.items(): + with self.subTest(name=name): + merged = build_neighbor_list( + case["coord"], + case["atype"], + nloc=case["nloc"], + rcut=case["rcut"], + sel=case["sel"], + distinguish_types=False, + ) + lower_formatted = nlist_distinguish_types( + merged, case["atype"], case["sel"] + ) + early_formatted = build_neighbor_list( + case["coord"], + case["atype"], + nloc=case["nloc"], + rcut=case["rcut"], + sel=case["sel"], + distinguish_types=True, + ) - np.testing.assert_array_equal(merged[0, 0], [1, 2]) - np.testing.assert_array_equal(lower_formatted, early_formatted) - np.testing.assert_array_equal(lower_formatted[0, 0], [1, -1]) + np.testing.assert_array_equal(lower_formatted, early_formatted) + if case["expected"] is not None: + np.testing.assert_array_equal(lower_formatted, case["expected"]) + if name == "batched_ghost_ties_virtual_padding": + self.assertTrue(np.any(lower_formatted == -1)) dtype = np.float64 From 8dbf0d4b20d981b6531ebb6b012c508280261ae0 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 00:18:58 +0800 Subject: [PATCH 3/4] test(nlist): add independent type-aware reference Validate both neighbor-list formatting paths against a standalone per-type distance-sort and truncation reference. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/common/dpmodel/test_nlist.py | 46 ++++++++++++++++++++++- 1 file changed, 44 insertions(+), 2 deletions(-) diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py index 8dddc189cb..ede0e4824f 100644 --- a/source/tests/common/dpmodel/test_nlist.py +++ b/source/tests/common/dpmodel/test_nlist.py @@ -25,6 +25,40 @@ ) +def _reference_type_neighbor_list( + coord: np.ndarray, + atype: np.ndarray, + nloc: int, + rcut: float, + sel: list[int], +) -> np.ndarray: + """Build a simple per-type distance-sorted neighbor-list reference.""" + coord = np.asarray(coord).reshape(coord.shape[0], -1, 3) + expected = np.full((coord.shape[0], nloc, sum(sel)), -1, dtype=np.int64) + offsets = np.cumsum([0, *sel]) + for frame in range(coord.shape[0]): + for center in range(nloc): + if atype[frame, center] < 0: + continue + distances = np.linalg.norm(coord[frame] - coord[frame, center], axis=-1) + for type_index, limit in enumerate(sel): + candidates = [ + atom + for atom in range(coord.shape[1]) + if atom != center + and atype[frame, atom] == type_index + and distances[atom] <= rcut + ] + candidates.sort(key=lambda atom: (distances[atom], atom)) + selected = candidates[:limit] + expected[ + frame, + center, + offsets[type_index] : offsets[type_index] + len(selected), + ] = selected + return expected + + class TestDPModelFormatNlist(unittest.TestCase): def setUp(self) -> None: # nloc == 3, nall == 4 @@ -236,10 +270,18 @@ def test_lower_type_split_matches_early_type_split(self) -> None: sel=case["sel"], distinguish_types=True, ) + reference = _reference_type_neighbor_list( + case["coord"], + case["atype"], + case["nloc"], + case["rcut"], + case["sel"], + ) - np.testing.assert_array_equal(lower_formatted, early_formatted) + np.testing.assert_array_equal(lower_formatted, reference) + np.testing.assert_array_equal(early_formatted, reference) if case["expected"] is not None: - np.testing.assert_array_equal(lower_formatted, case["expected"]) + np.testing.assert_array_equal(reference, case["expected"]) if name == "batched_ghost_ties_virtual_padding": self.assertTrue(np.any(lower_formatted == -1)) From 49b9bd9cbcd2c6dc233ddf9018e67eaad80c1fed Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 00:53:08 +0800 Subject: [PATCH 4/4] test(dpmodel): honor global neighbor candidate cap Coding-Agent: Codex\nCodex-Version: codex-cli 0.144.6\nModel: gpt-5.6-sol\nReasoning-Effort: xhigh --- source/tests/common/dpmodel/test_nlist.py | 44 +++++++++++++++++------ 1 file changed, 34 insertions(+), 10 deletions(-) diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py index ede0e4824f..63892b7d1e 100644 --- a/source/tests/common/dpmodel/test_nlist.py +++ b/source/tests/common/dpmodel/test_nlist.py @@ -32,7 +32,7 @@ def _reference_type_neighbor_list( rcut: float, sel: list[int], ) -> np.ndarray: - """Build a simple per-type distance-sorted neighbor-list reference.""" + """Build a distance-sorted reference with the runtime global cap.""" coord = np.asarray(coord).reshape(coord.shape[0], -1, 3) expected = np.full((coord.shape[0], nloc, sum(sel)), -1, dtype=np.int64) offsets = np.cumsum([0, *sel]) @@ -41,16 +41,20 @@ def _reference_type_neighbor_list( if atype[frame, center] < 0: continue distances = np.linalg.norm(coord[frame] - coord[frame, center], axis=-1) + candidates = [ + atom + for atom in range(coord.shape[1]) + if atom != center + and atype[frame, atom] >= 0 + and distances[atom] <= rcut + ] + candidates.sort(key=lambda atom: (distances[atom], atom)) + candidates = candidates[: sum(sel)] for type_index, limit in enumerate(sel): - candidates = [ - atom - for atom in range(coord.shape[1]) - if atom != center - and atype[frame, atom] == type_index - and distances[atom] <= rcut + selected = [ + atom for atom in candidates if atype[frame, atom] == type_index ] - candidates.sort(key=lambda atom: (distances[atom], atom)) - selected = candidates[:limit] + selected = selected[:limit] expected[ frame, center, @@ -192,7 +196,8 @@ def test_lower_type_split_matches_early_type_split(self) -> None: """Global candidate selection and early type splits must be equivalent. The matrix covers same-type truncation, real batch axes, multiple local - atoms, ghosts, exact distance ties, virtual atoms, and padded buckets. + atoms, ghosts, exact distance ties, virtual atoms, the global candidate + cap, and padded buckets. Keeping it on the shared formatting fixture makes the contract visible beside the other short/equal/long neighbor-list cases. """ @@ -216,6 +221,25 @@ def test_lower_type_split_matches_early_type_split(self) -> None: "rcut": 3.0, "expected": np.array([[[1, 2, 4]]], dtype=np.int64), }, + "global_cap_excludes_far_type": { + "coord": np.array( + [ + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [1.5, 0.0, 0.0], + [2.0, 0.0, 0.0], + [2.5, 0.0, 0.0], + ] + ], + dtype=np.float64, + ), + "atype": np.array([[0, 0, 0, 0, 1]], dtype=np.int64), + "nloc": 1, + "sel": [1, 2], + "rcut": 3.0, + "expected": np.array([[[1, -1, -1]]], dtype=np.int64), + }, "batched_ghost_ties_virtual_padding": { "coord": np.array( [