From 7f5b0377e7abe3e45cd3997c7ddf9b65b53bc04a Mon Sep 17 00:00:00 2001
From: njzjz-bot <njzjz.bot@gmail.com>
Date: Fri, 17 Jul 2026 08:52:15 +0800
Subject: [PATCH] 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