Skip to content

test(pt_expt): document compiled neighbor type equivalence #5860

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
njzjz merged 4 commits into from
Aug 2, 2026
Merged

test(pt_expt): document compiled neighbor type equivalence #5860

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7f5b037
test(pt_expt): document compiled neighbor type equivalence
njzjz-bot Jul 17, 2026
c17f783
test(dpmodel): broaden neighbor type layout coverage
njzjz-bot Aug 1, 2026
8dbf0d4
test(nlist): add independent type-aware reference
njzjz-bot Aug 1, 2026
49b9bd9
test(dpmodel): honor global neighbor candidate cap
njzjz-bot Aug 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions deepmd/pt_expt/train/training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
156 changes: 156 additions & 0 deletions source/tests/common/dpmodel/test_nlist.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,48 @@
extend_coord_with_ghosts,
get_multiple_nlist_key,
inter2phys,
nlist_distinguish_types,
)


def _reference_type_neighbor_list(
coord: np.ndarray,
atype: np.ndarray,
nloc: int,
rcut: float,
sel: list[int],
) -> np.ndarray:
"""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])
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)
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):
selected = [
atom for atom in candidates if atype[frame, atom] == type_index
]
selected = selected[: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
Expand Down Expand Up @@ -153,6 +192,123 @@ def test_nlist_lt(self) -> None:
)
np.testing.assert_allclose(self.expected_nlist, nlist1)

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, 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.
"""
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),
},
"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(
[
[
[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,
)
reference = _reference_type_neighbor_list(
case["coord"],
case["atype"],
case["nloc"],
case["rcut"],
case["sel"],
)

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(reference, case["expected"])
if name == "batched_ghost_ties_virtual_padding":
self.assertTrue(np.any(lower_formatted == -1))


dtype = np.float64

Expand Down
Loading