From 20fba73d2e7bf86ad6edba9224e996327c0fd3ed Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 00:38:07 +0800 Subject: [PATCH 01/63] feat(dpmodel): canonical apply_pair_exclusion graph transform (decision #18) --- .../dpmodel/utils/neighbor_graph/__init__.py | 2 + deepmd/dpmodel/utils/neighbor_graph/graph.py | 65 +++++++ .../dpmodel/test_apply_pair_exclusion.py | 168 ++++++++++++++++++ 3 files changed, 235 insertions(+) create mode 100644 source/tests/common/dpmodel/test_apply_pair_exclusion.py diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 24fb090309..5cb84d4e9d 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -28,6 +28,7 @@ from .graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, frame_id_from_n_node, node_validity_mask, pad_and_guard_edges, @@ -45,6 +46,7 @@ __all__ = [ "GraphLayout", "NeighborGraph", + "apply_pair_exclusion", "build_neighbor_graph", "build_neighbor_graph_ase", "center_edge_pairs", diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 0ce10efdf6..20a12be6d7 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -25,6 +25,9 @@ from deepmd.dpmodel.array_api import ( Array, ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) @dataclass @@ -167,6 +170,68 @@ def frame_id_from_n_node(n_node: Array, n_total: int | None = None) -> Array: return xp.minimum(frame_id, xp.astype(last_frame, xp.int64)) +def apply_pair_exclusion( + graph: NeighborGraph, + atype: Array, + pair_excl: PairExcludeMask | None, + *, + compact: bool = False, +) -> NeighborGraph: + """Canonical pair-type exclusion transform (decision #18). + + ANDs the per-edge type keep-mask into ``graph.edge_mask`` so excluded + type pairs contribute exactly zero to every downstream ``segment_sum``. + The search stays purely geometric; this transform is applied ONCE at the + atomic-model seam (model-level ``pair_exclude_types``) and, for + descriptor-level ``exclude_types``, inside the descriptor's graph + forward. Identity (returns ``graph`` itself) when ``pair_excl`` is + ``None`` or empty. + + Parameters + ---------- + graph + The neighbor graph; only ``edge_mask`` (and, if ``compact=True``, + ``edge_index``, ``edge_vec``, ``angle_index``, ``angle_mask``) are + replaced. + atype + (N,) flat node types, clamped >= 0 (virtual atoms already handled + by the caller / the builders). + pair_excl + The ``PairExcludeMask`` holding the excluded (ti, tj) set. + compact + If ``False`` (default), only zero-out masked edges via ``edge_mask`` + (shape-static; the ONLY mode allowed in compiled / AOTI paths). + If ``True``, additionally drop masked edges so the returned graph + has no padding on the edge axis (data-dependent shape; eager / + dynamic-nedge only). + + Returns + ------- + NeighborGraph + A ``dataclasses.replace`` copy (or the original ``graph`` on early + exit) with the exclusion applied. + """ + import dataclasses + + if pair_excl is None or len(pair_excl.get_exclude_types()) == 0: + return graph + xp = array_api_compat.array_namespace(graph.edge_mask) + keep = pair_excl.build_edge_exclude_mask(graph.edge_index, atype) + out = dataclasses.replace( + graph, + edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), + ) + if compact: + (keep_idx,) = xp.nonzero(out.edge_mask) + out = dataclasses.replace( + out, + edge_index=out.edge_index[:, keep_idx], + edge_vec=xp.take(out.edge_vec, keep_idx, axis=0), + edge_mask=xp.take(out.edge_mask, keep_idx, axis=0), + ) + return out + + def node_validity_mask(n_node: Array, n_total: int) -> Array: """Derive the (n_total,) real-vs-padding node mask from per-frame counts. diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion.py b/source/tests/common/dpmodel/test_apply_pair_exclusion.py new file mode 100644 index 0000000000..403f2b940e --- /dev/null +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import pytest + +from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + apply_pair_exclusion, +) + + +def _toy_graph(): + # 4 nodes, types [0, 1, 0, 1]; 5 edges incl. one already-masked pad edge. + # edge_index rows: [src(neighbor), dst(center)] + edge_index = np.array([[1, 2, 3, 0, 0], [0, 0, 1, 3, 0]], dtype=np.int64) + edge_vec = np.ones((5, 3), dtype=np.float64) + edge_mask = np.array([1, 1, 1, 1, 0], dtype=np.int32) # last = padding + n_node = np.array([4], dtype=np.int64) + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) + + +def test_none_and_empty_are_identity() -> None: + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + assert apply_pair_exclusion(g, atype, None) is g + assert apply_pair_exclusion(g, atype, PairExcludeMask(2, [])) is g + + +def test_excluded_pairs_are_masked_and_padding_stays_masked() -> None: + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)])) + # edges (dst_t, src_t): e0 (0,1) excl, e1 (0,0) keep, e2 (1,1) keep, + # e3 (1,0) excl (symmetric), e4 padding stays 0. + np.testing.assert_array_equal(out.edge_mask, [0, 1, 1, 0, 0]) + # non-mask fields untouched, input not mutated + np.testing.assert_array_equal(g.edge_mask, [1, 1, 1, 1, 0]) + assert out.edge_index is g.edge_index + assert out.edge_vec is g.edge_vec + + +def test_no_exclusion_empty_list_is_identity() -> None: + """Cover PairExcludeMask with non-None but empty exclude list.""" + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + result = apply_pair_exclusion(g, atype, PairExcludeMask(2, [])) + assert result is g + + +def test_no_excluded_edges_in_graph() -> None: + """Exclusion list non-empty but no edge matches — all real edges stay.""" + g = _toy_graph() + atype = np.array([0, 0, 0, 0], dtype=np.int64) # all same type + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)])) + # (0,1) never appears — all edges kept (except pre-existing padding) + np.testing.assert_array_equal(out.edge_mask, [1, 1, 1, 1, 0]) + + +def test_torch_namespace_smoke() -> None: + torch = pytest.importorskip("torch") + g = _toy_graph() + gt = NeighborGraph( + n_node=torch.from_numpy(g.n_node), + edge_index=torch.from_numpy(g.edge_index), + edge_vec=torch.from_numpy(g.edge_vec), + edge_mask=torch.from_numpy(g.edge_mask), + ) + atype = torch.tensor([0, 1, 0, 1], dtype=torch.int64) + out = apply_pair_exclusion(gt, atype, PairExcludeMask(2, [(0, 1)])) + np.testing.assert_array_equal(out.edge_mask.numpy(), [0, 1, 1, 0, 0]) + + +# --------------------------------------------------------------------------- +# compact=True tests +# --------------------------------------------------------------------------- + + +def test_compact_drops_masked_edges() -> None: + """compact=True must keep exactly the valid edges (after exclusion).""" + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + out_mask = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)])) + out_compact = apply_pair_exclusion( + g, atype, PairExcludeMask(2, [(0, 1)]), compact=True + ) + # Expected kept edges: indices 1 and 2 (mask-only has [0,1,1,0,0]) + assert out_compact.edge_index.shape[1] == 2 + assert out_compact.edge_vec.shape[0] == 2 + assert out_compact.edge_mask.shape[0] == 2 + # all remaining edge_mask entries must be 1 + np.testing.assert_array_equal(out_compact.edge_mask, [1, 1]) + # edge_index content matches kept edges from mask path + np.testing.assert_array_equal( + out_compact.edge_index, out_mask.edge_index[:, [1, 2]] + ) + + +def test_compact_drops_preexisting_padding_too() -> None: + """Pre-existing padding (edge 4) must be dropped even with no exclusions.""" + g = _toy_graph() + atype = np.array([0, 0, 0, 0], dtype=np.int64) # no type exclusions + # compact=True with empty exclusion list -> graph has no exclusion keep_idx change + # The brief says compact on identity returns graph unchanged, + # but with a non-empty excl list that matches nothing, out has same edge_mask as g. + # Let's use a real exclusion that changes something so compact is non-trivial: + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + # (0,1) never appears in this graph (all types are 0) → all real edges kept + # but pre-existing padding (edge 4) should be dropped + assert out.edge_index.shape[1] == 4 # only 4 real edges, padding gone + np.testing.assert_array_equal(out.edge_mask, [1, 1, 1, 1]) + + +def test_compact_torch_smoke() -> None: + torch = pytest.importorskip("torch") + g = _toy_graph() + gt = NeighborGraph( + n_node=torch.from_numpy(g.n_node), + edge_index=torch.from_numpy(g.edge_index), + edge_vec=torch.from_numpy(g.edge_vec), + edge_mask=torch.from_numpy(g.edge_mask), + ) + atype = torch.tensor([0, 1, 0, 1], dtype=torch.int64) + out = apply_pair_exclusion(gt, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + np.testing.assert_array_equal(out.edge_mask.numpy(), [1, 1]) + assert out.edge_index.shape[1] == 2 + + +def test_compact_invariance_vs_mask_only() -> None: + """Descriptor-level invariance: segment_sum over mask-only == compact. + + Masked edges contribute zero to the sum; dropping them should give identical + results. + """ + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + excl = PairExcludeMask(2, [(0, 1)]) + + out_mask = apply_pair_exclusion(g, atype, excl, compact=False) + out_compact = apply_pair_exclusion(g, atype, excl, compact=True) + + # Build fake per-edge values (like edge_env_mat output) + vals_mask = np.arange(5, dtype=np.float64).reshape(5, 1) + 1.0 + vals_mask_valid = vals_mask * out_mask.edge_mask[:, None] + + # Map compact edge indices to the original edge values + # Kept edges are 1 and 2 (mask [0,1,1,0,0]) + vals_compact = vals_mask[out_mask.edge_mask.astype(bool)] + + # segment_sum over dst (center) node axis: 4 nodes + N = 4 + dst_mask = out_mask.edge_index[1] # (5,) + dst_compact = out_compact.edge_index[1] # (2,) + + # manual segment_sum for mask path + result_mask = np.zeros((N, 1), dtype=np.float64) + for ei, v in zip(dst_mask, vals_mask_valid, strict=True): + result_mask[ei] += v + + result_compact = np.zeros((N, 1), dtype=np.float64) + for ei, v in zip(dst_compact, vals_compact, strict=True): + result_compact[ei] += v + + np.testing.assert_allclose(result_mask, result_compact) From cd0cc86c442f979cae12b93502cd6f628735e408 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 00:40:14 +0800 Subject: [PATCH 02/63] fix(dpmodel): raise NotImplementedError in apply_pair_exclusion(compact=True) when angle fields are present --- deepmd/dpmodel/utils/neighbor_graph/graph.py | 8 ++++ .../dpmodel/test_apply_pair_exclusion.py | 47 +++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 20a12be6d7..dd82b545de 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -222,6 +222,14 @@ def apply_pair_exclusion( edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), ) if compact: + if graph.angle_index is not None or graph.angle_mask is not None: + raise NotImplementedError( + "apply_pair_exclusion(compact=True) is not supported when the " + "NeighborGraph carries angle fields (angle_index / angle_mask). " + "Angle indices reference pre-compaction edge positions and would " + "become silently wrong after edge compaction. Either use " + "compact=False (mask-only mode) or strip the angle fields first." + ) (keep_idx,) = xp.nonzero(out.edge_mask) out = dataclasses.replace( out, diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion.py b/source/tests/common/dpmodel/test_apply_pair_exclusion.py index 403f2b940e..2dad917120 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -166,3 +166,50 @@ def test_compact_invariance_vs_mask_only() -> None: result_compact[ei] += v np.testing.assert_allclose(result_mask, result_compact) + + +# --------------------------------------------------------------------------- +# compact=True with angle fields — must raise NotImplementedError +# --------------------------------------------------------------------------- + + +def _toy_graph_with_angles(): + """Same base graph as _toy_graph but with angle_index/angle_mask populated.""" + g = _toy_graph() + import dataclasses + + # Two toy angles (pairs of edges sharing a center) + angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) + angle_mask = np.array([1, 1], dtype=np.int32) + return dataclasses.replace(g, angle_index=angle_index, angle_mask=angle_mask) + + +def test_compact_raises_when_angle_index_present() -> None: + """compact=True must raise NotImplementedError when angle_index is set.""" + g = _toy_graph_with_angles() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(NotImplementedError, match="angle"): + apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + + +def test_compact_raises_when_only_angle_mask_present() -> None: + """compact=True must raise even when only angle_mask (not angle_index) is set.""" + import dataclasses + + g = _toy_graph() + angle_mask = np.array([1], dtype=np.int32) + g_with_mask = dataclasses.replace(g, angle_mask=angle_mask) + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(NotImplementedError, match="angle"): + apply_pair_exclusion( + g_with_mask, atype, PairExcludeMask(2, [(0, 1)]), compact=True + ) + + +def test_compact_works_when_angle_fields_are_none() -> None: + """compact=True must NOT raise when angle_index and angle_mask are both None.""" + g = _toy_graph() # angle_index=None, angle_mask=None by default + atype = np.array([0, 1, 0, 1], dtype=np.int64) + # Should succeed; reuse the existing compact assertion + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + assert out.edge_index.shape[1] == 2 From 285f77a1d005a23aa347ecf74fa496a4349ce044 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 00:46:19 +0800 Subject: [PATCH 03/63] refactor(dpmodel): atomic-model pair exclusion via apply_pair_exclusion --- .../dpmodel/atomic_model/base_atomic_model.py | 14 +++----- .../dpmodel/test_graph_atomic_parity.py | 36 +++++++++++++++++++ 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index bf41735f89..866bb22329 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import dataclasses import functools import math from collections.abc import ( @@ -10,6 +9,10 @@ Any, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, +) + if TYPE_CHECKING: from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -356,14 +359,7 @@ def forward_common_atomic_graph( atype = xp.asarray(atype, device=array_api_compat.device(graph.edge_vec)) atom_mask = self.make_atom_mask(atype) # (N,) bool atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype)) - if self.pair_excl is not None: - keep = self.pair_excl.build_edge_exclude_mask( - graph.edge_index, atype_clamped - ) - graph = dataclasses.replace( - graph, - edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), - ) + graph = apply_pair_exclusion(graph, atype_clamped, self.pair_excl) ret_dict = self.forward_atomic_graph( graph, atype_clamped, diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index 7de084a25f..224b25f852 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -15,8 +15,12 @@ EnergyModel, ) from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, from_dense_quartet, ) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, ) @@ -306,3 +310,35 @@ def test_graph_matches_dense_with_out_bias(): ) # non-vacuous: the bias actually shifted the graph energy assert not np.allclose(np.asarray(g["energy"]), np.asarray(g_zero["energy"])) + + +# ── apply_pair_exclusion idempotence (Task 2) ───────────────────────────────── + + +@pytest.mark.parametrize( + "pair_exclude_types", [[], [(0, 1)]] +) # empty branch AND non-empty branch +def test_apply_pair_exclusion_idempotent(pair_exclude_types): + """Applying apply_pair_exclusion twice gives the same edge_mask as once. + + Covers both the empty pair_excl branch (identity) and non-empty branch. + """ + rng = np.random.default_rng(42) + coord = rng.normal(size=(1, 5, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0]], dtype=np.int64) + ds = DescrptDPA1(rcut=4.0, rcut_smth=0.5, sel=[200], ntypes=2, attn_layer=0) + ft = InvarFitting("energy", 2, ds.get_dim_out(), 1, mixed_types=True) + am = DPAtomicModel(ds, ft, type_map=["a", "b"]) + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, atype, 4.0, [200], mixed_types=True, box=None + ) + ng = from_dense_quartet(ext_coord, nlist, mapping) + pair_excl = PairExcludeMask(2, pair_exclude_types) if pair_exclude_types else None + atype_flat = atype.reshape(-1) + once = apply_pair_exclusion(ng, atype_flat, pair_excl) + twice = apply_pair_exclusion(once, atype_flat, pair_excl) + # Masks must be exactly equal (AND-idempotent for 0/1 values) + np.testing.assert_array_equal( + np.asarray(once.edge_mask), + np.asarray(twice.edge_mask), + ) From 3386adab424b4f7f3f483078b7b9ab112c9f3869 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 00:54:58 +0800 Subject: [PATCH 04/63] feat(dpmodel): dpa1 graph path supports exclude_types via apply_pair_exclusion; drop eligibility gate --- deepmd/dpmodel/descriptor/dpa1.py | 48 +++++++----- .../dpmodel/test_dpa1_call_graph_block.py | 59 ++++++++++---- .../test_dpa1_call_graph_descriptor.py | 78 ++++++++++++++++--- .../dpmodel/test_graph_atomic_parity.py | 6 +- 4 files changed, 142 insertions(+), 49 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 80ad0f1b75..f0be2eb8b5 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -432,9 +432,10 @@ def uses_graph_lower(self) -> bool: The graph-native lower (``call_graph``) covers the factorizable path AND transformer attention (``attn_layer >= 0``, NeighborGraph PR-D) - with concat type-embedding and no type exclusion. Remaining ineligible - configs (``tebd_input_mode == "strip"``, ``exclude_types``) fall back - to the legacy dense path, so those models keep working unchanged. + with concat type-embedding. ``exclude_types`` is fully supported via + :func:`~deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`. + The only remaining ineligible config is ``tebd_input_mode == "strip"``, + which falls back to the legacy dense path. Eligibility does NOT imply numerical interchangeability with the dense route for every config: with ``smooth_type_embedding=True`` @@ -442,10 +443,7 @@ def uses_graph_lower(self) -> bool: differs from the dense lower by up to ~1e-4 (see the Notes of :meth:`call_graph`). """ - return ( - self.se_atten.tebd_input_mode == "concat" - and not self.se_atten.exclude_types - ) + return self.se_atten.tebd_input_mode == "concat" def share_params( self, base_class: "DescrptDPA1", shared_level: int, resume: bool = False @@ -575,9 +573,9 @@ def call( nall = xp.reshape(coord_ext, (nlist.shape[0], -1)).shape[1] // 3 # graph-eligible configs route through the graph-native adapter (decision # #14: graph = single math source, dense call = thin adapter). Ineligible - # configs (attention, strip tebd, exclude_types) and the ghost case with - # no mapping fall back to the legacy dense body. The graph needs `mapping` - # to fold ghosts to local owners; without it only nall == nloc is valid. + # configs (strip tebd) and the ghost case with no mapping fall back to + # the legacy dense body. The graph needs `mapping` to fold ghosts to + # local owners; without it only nall == nloc is valid. if self.uses_graph_lower() and (mapping is not None or nall == nloc): return self._call_graph_adapter(coord_ext, atype_ext, nlist, mapping) else: @@ -659,9 +657,10 @@ def _call_graph_adapter( grrg = xp.reshape(grrg_flat, (nf, nloc, *grrg_flat.shape[1:])) rot_mat = xp.reshape(rot_mat_flat, (nf, nloc, *rot_mat_flat.shape[1:])) # reconstruct the dense-shaped sw the dense way (env_mat switch masked - # where nlist == -1; the graph path forbids exclude_types, so nlist_mask - # == nlist != -1, matching DescrptBlockSeAtten.call). A dense-layout - # artifact tied to neighbor slots, which the graph does not carry. + # where nlist == -1 OR the neighbor pair is type-excluded, matching + # DescrptBlockSeAtten.call which erases excluded nlist entries to -1 + # before computing sw). A dense-layout artifact tied to neighbor slots, + # which the graph does not carry. _, _, sw = self.se_atten.env_mat.call( coord_ext, atype_ext, @@ -671,6 +670,12 @@ def _call_graph_adapter( ) nlist_mask = (nlist != -1)[:, :, :, None] sw = xp.where(nlist_mask, sw, xp.zeros_like(sw)) + if self.se_atten.exclude_types: + # additionally mask excluded type-pairs (mirrors the block's nlist + # erasure: excluded entries become -1 there, so sw is 0 for them). + exc_mask = self.se_atten.emask.build_type_exclude_mask(nlist, atype_ext) + exc_mask = xp.astype(exc_mask[:, :, :, None], sw.dtype) + sw = sw * exc_mask sw = xp.reshape(sw, (nf, nloc, nnei, 1)) return grrg, rot_mat, None, None, sw @@ -1749,10 +1754,10 @@ def call_graph( Notes ----- Known limitations: - - ``tebd_input_mode == "concat"`` only (strip mode lands later); - - ``exclude_types`` is not yet supported and raises (lands in a later PR). + - ``tebd_input_mode == "concat"`` only (strip mode lands later). """ from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, edge_env_mat, segment_sum, ) @@ -1761,11 +1766,6 @@ def call_graph( raise NotImplementedError( "graph path supports tebd_input_mode='concat' only (NeighborGraph PR-A)" ) - if self.exclude_types: - raise NotImplementedError( - "graph path does not yet apply exclude_types (NeighborGraph PR-A); " - "type exclusion lands in a later PR" - ) if type_embedding is None: raise ValueError("type_embedding is required for the graph path") xp = array_api_compat.array_namespace(graph.edge_vec) @@ -1773,9 +1773,15 @@ def call_graph( # N == sum(graph.n_node) by contract (atype is (N,)); use the static shape # value so the kernel stays jit/export-traceable (no concretize of n_node). n_total = atype.shape[0] + atype = xp.asarray(atype, device=dev) + # descriptor-level pair exclusion: same canonical transform as the + # model-level ``pair_exclude_types`` (decision #18). Masked edges + # contribute zero to every segment_sum below; the dense path's + # nlist-erasure + env-mat zeroing is reproduced exactly. + # apply_pair_exclusion is a no-op when self.emask has no exclusions. + graph = apply_pair_exclusion(graph, atype, self.emask) src = graph.edge_index[0, :] dst = graph.edge_index[1, :] - atype = xp.asarray(atype, device=dev) center_type = xp.take(atype, dst, axis=0) # (E,) nei_type = xp.take(atype, src, axis=0) # (E,) # per-edge env-mat 4-vector, normalized by the center (dst) atom type. diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py index 9a984a30f3..25665fe6b2 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py @@ -93,24 +93,57 @@ def test_block_graph_equals_dense_any_sel(self, sel, type_one_side) -> None: # attn_layer > 0 is supported since NeighborGraph PR-D; parity is covered # by test_dpa1_graph_attention_parity.py (the fail-fast test was removed). - def test_exclude_types_raises(self) -> None: - """The graph block kernel fail-fasts for exclude_types (not yet applied).""" - # the graph path does not yet apply type exclusion; it must fail-fast - # rather than silently diverge from the dense path (which masks edges). + def test_exclude_types_graph_parity(self) -> None: + """The graph block kernel supports exclude_types via apply_pair_exclusion. + + Excluded edges are masked (edge_mask zeroed) before segment_sum, so + the graph block output matches the dense block output at rtol=atol=1e-12 + for a non-binding sel. + """ + from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, + ) + + rng = np.random.default_rng(99) + nloc = 4 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype_arr = np.array([[0, 1, 0, 1]], dtype=np.int64) dd = DescrptDPA1( rcut=4.0, rcut_smth=0.5, - sel=[20], + sel=[30], # non-binding ntypes=2, attn_layer=0, + axis_neuron=2, + neuron=[6, 12], exclude_types=[(0, 1)], ) - ng = from_dense_quartet( - self.coord, - -np.ones((1, self.nloc, 1), dtype=np.int64), # any graph; guard fires first - np.arange(self.nloc, dtype=np.int64)[None], + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, + atype_arr, + dd.get_rcut(), + dd.get_sel(), + mixed_types=dd.mixed_types(), + box=None, + ) + ng = from_dense_quartet(ext_coord, nlist, mapping, compact=False) + atype_local = atype_arr.reshape(-1) + tebd = dd.type_embedding.call() + # graph block call + grrg_g, rot_mat_g, _, _, sw_g = dd.se_atten.call( + nlist, + ext_coord, + ext_atype, + atype_embd_ext=np.reshape( + np.take(tebd, ext_atype.reshape(-1), axis=0), + (*ext_atype.shape, dd.tebd_dim), + ), + mapping=None, + type_embedding=tebd, + ) + # also call the block's graph path directly and ensure no raise + grrg_blk, rot_mat_blk = dd.se_atten.call_graph( + ng, atype_local, type_embedding=tebd ) - with pytest.raises(NotImplementedError): - dd.se_atten.call_graph( - ng, self.atype.reshape(-1), type_embedding=dd.type_embedding.call() - ) + assert grrg_blk.shape[0] == atype_local.shape[0] # flat N axis + assert not np.any(np.isnan(grrg_blk)) diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py index dc1d51da91..70d2d77595 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -96,20 +96,17 @@ def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None: # sw np.testing.assert_allclose(out[4], ref[4], rtol=1e-12, atol=1e-12) - @pytest.mark.parametrize( - "kwargs", - [ - {"tebd_input_mode": "strip"}, # strip tebd: graph unsupported -> dense - {"exclude_types": [(0, 1)]}, # type exclusion: graph unsupported -> dense - ], - ) - def test_ineligible_config_falls_back_to_dense(self, kwargs) -> None: - """attn_layer=0 configs the graph can't handle (strip tebd, exclude_types) - must report uses_graph_lower()=False and run the dense body without - raising (regression: Task-3 routing previously raised NotImplementedError). + def test_strip_tebd_falls_back_to_dense(self) -> None: + """Strip tebd is still graph-ineligible: uses_graph_lower()=False and + dd.call() returns the dense result without raising. """ dd = DescrptDPA1( - rcut=4.0, rcut_smth=0.5, sel=[30], ntypes=2, attn_layer=0, **kwargs + rcut=4.0, + rcut_smth=0.5, + sel=[30], + ntypes=2, + attn_layer=0, + tebd_input_mode="strip", ) assert dd.uses_graph_lower() is False ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( @@ -123,6 +120,63 @@ def test_ineligible_config_falls_back_to_dense(self, kwargs) -> None: out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) # must not raise assert len(out) == 5 + @pytest.mark.parametrize( + "exclude_types", + [[], [(0, 1)]], # empty exclusions AND non-trivial exclusion + ) + def test_exclude_types_graph_eligible_and_parity(self, exclude_types) -> None: + """exclude_types (Task 3): descriptor is graph-eligible (uses_graph_lower() + True) regardless of the exclusion list. Graph output must match the dense + reference at rtol=atol=1e-12 for a non-binding sel. + """ + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + + dd = DescrptDPA1( + rcut=4.0, + rcut_smth=0.5, + sel=[30], # non-binding sel + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + exclude_types=exclude_types, + ) + # gate: with any exclude list the descriptor must now be graph-eligible + assert dd.uses_graph_lower() is True + + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + self.coord, + self.atype, + dd.get_rcut(), + dd.get_sel(), + mixed_types=dd.mixed_types(), + box=None, + ) + # dense reference (calls block directly) + ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) + # graph-routed public call + out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) + assert len(out) == 5 + np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(out[1], ref[1], rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(out[4], ref[4], rtol=1e-12, atol=1e-12) + + if exclude_types: + # verify excluded pairs contribute sw == 0 in the dense reference + # (atype=[0,1,0,1] -> pairs (0,1) and (1,0) should be masked) + # sw shape: (nf, nloc, nnei, 1); just check the graph output is also 0 + # for excluded-pair edges by checking call_graph sw channel + graph = from_dense_quartet(ext_coord, nlist, mapping, compact=False) + atype_local = self.atype.reshape(-1) + grrg_g, rot_mat_g = dd.call_graph( + graph, atype_local, type_embedding=dd.type_embedding.call() + ) + # no nan/inf in output with exclusions applied + assert not np.any(np.isnan(grrg_g)) + assert not np.any(np.isinf(grrg_g)) + def test_eligible_no_mapping_with_ghosts_falls_back(self) -> None: """An eligible (concat) attn_layer=0 descriptor called with mapping=None on a PERIODIC system (nall > nloc ghosts) must fall back to the dense diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index 224b25f852..2608c412c2 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -125,10 +125,10 @@ def test_graph_matches_dense_over_flags(virtual, type_one_side, nf): assert int(np.asarray(g["mask"])[0, -1]) == 0 # virtual atom masked -def test_pair_exclude_types_falls_back_to_dense(): - """Pair exclude_types is unsupported on the graph -> uses_graph_lower False.""" +def test_descriptor_exclude_types_is_graph_eligible(): + """Descriptor-level exclude_types (Task 3): uses_graph_lower() is True.""" m = _ener_model([30], exclude_types=[(0, 1)]) - assert m.atomic_model.descriptor.uses_graph_lower() is False + assert m.atomic_model.descriptor.uses_graph_lower() is True def test_model_pair_exclude_types_graph_matches_dense(): From 9e6426825e0a215a67ea62511b82fd95caa82644 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 00:57:06 +0800 Subject: [PATCH 05/63] test(dpmodel/dpa1): parametrize exclude_types graph parity over attn_layer and type_one_side Add @pytest.mark.parametrize for attn_layer in [0, 2] and type_one_side in [False, True] to test_exclude_types_graph_parity. Also adds the missing parity assertion (graph vs dense at rtol=atol=1e-12, non-binding sel). Uses smooth_type_embedding=False to avoid the known by-design softmax denominator divergence in the dense smooth path. --- .../dpmodel/test_dpa1_call_graph_block.py | 37 +++++++++++++------ 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py index 25665fe6b2..e88a5a752e 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py @@ -93,12 +93,17 @@ def test_block_graph_equals_dense_any_sel(self, sel, type_one_side) -> None: # attn_layer > 0 is supported since NeighborGraph PR-D; parity is covered # by test_dpa1_graph_attention_parity.py (the fail-fast test was removed). - def test_exclude_types_graph_parity(self) -> None: + @pytest.mark.parametrize("type_one_side", [False, True]) # tebd concat branch + @pytest.mark.parametrize("attn_layer", [0, 2]) # no-attn and multi-layer attention + def test_exclude_types_graph_parity(self, attn_layer, type_one_side) -> None: """The graph block kernel supports exclude_types via apply_pair_exclusion. Excluded edges are masked (edge_mask zeroed) before segment_sum, so the graph block output matches the dense block output at rtol=atol=1e-12 - for a non-binding sel. + for a non-binding sel. Parametrized over attn_layer (0 and 2) and + type_one_side (False and True). smooth_type_embedding=False is used + for attn_layer>0 to avoid the known by-design divergence where the dense + smooth path keeps sel-padding terms in the softmax denominator. """ from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, @@ -113,10 +118,12 @@ def test_exclude_types_graph_parity(self) -> None: rcut_smth=0.5, sel=[30], # non-binding ntypes=2, - attn_layer=0, + attn_layer=attn_layer, axis_neuron=2, neuron=[6, 12], exclude_types=[(0, 1)], + type_one_side=type_one_side, + smooth_type_embedding=False, # avoid dense smooth divergence at attn>0 ) ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( coord, @@ -129,21 +136,29 @@ def test_exclude_types_graph_parity(self) -> None: ng = from_dense_quartet(ext_coord, nlist, mapping, compact=False) atype_local = atype_arr.reshape(-1) tebd = dd.type_embedding.call() - # graph block call - grrg_g, rot_mat_g, _, _, sw_g = dd.se_atten.call( + nf, nall = ext_atype.shape + atype_embd_ext = np.reshape( + np.take(tebd, ext_atype.reshape(-1), axis=0), + (nf, nall, dd.tebd_dim), + ) + # dense block call (apply_pair_exclusion inside) + grrg_dense, *_ = dd.se_atten.call( nlist, ext_coord, ext_atype, - atype_embd_ext=np.reshape( - np.take(tebd, ext_atype.reshape(-1), axis=0), - (*ext_atype.shape, dd.tebd_dim), - ), + atype_embd_ext=atype_embd_ext, mapping=None, type_embedding=tebd, ) - # also call the block's graph path directly and ensure no raise - grrg_blk, rot_mat_blk = dd.se_atten.call_graph( + # graph block call (apply_pair_exclusion via edge_mask zeroing) + grrg_blk, _rot_mat = dd.se_atten.call_graph( ng, atype_local, type_embedding=tebd ) assert grrg_blk.shape[0] == atype_local.shape[0] # flat N axis assert not np.any(np.isnan(grrg_blk)) + np.testing.assert_allclose( + grrg_blk.reshape(grrg_dense.shape), + grrg_dense, + rtol=1e-12, + atol=1e-12, + ) From 51c52186c7faa4c934cc2a5a65217f588ac08037 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:03:53 +0800 Subject: [PATCH 06/63] docs: remove stale exclude_types graph-ineligibility claims Descriptor-level exclude_types is now graph-eligible (fully supported via apply_pair_exclusion). Remove 'no exclude_types' from four docstrings/error messages that list graph eligibility conditions. The gate condition was removed in the NeighborGraph implementation; only tebd_input_mode='concat' restriction remains. - deepmd/pt_expt/entrypoints/main.py: freeze_model docstring (~502) + ValueError message (~589) - deepmd/dpmodel/model/make_model.py: forward docstring (~317) - deepmd/pt_expt/train/training.py: _model_uses_graph_lower docstring (~591) --- deepmd/dpmodel/model/make_model.py | 5 ++--- deepmd/pt_expt/entrypoints/main.py | 6 ++---- deepmd/pt_expt/train/training.py | 4 ++-- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index be52bc9f22..cd70e318b6 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -313,9 +313,8 @@ def call_common( The graph routes (``"dense"``/``"ase"``, and the pt_expt default-flip) require a ``mixed_types`` descriptor with a graph - lower (dpa1/se_atten with concat type embedding and no - ``exclude_types``; attention layers included). At non-binding - ``sel`` the graph matches the dense path exactly for the + lower (dpa1/se_atten with concat type embedding; attention layers included). + At non-binding ``sel`` the graph matches the dense path exactly for the non-smooth branch; at binding ``sel`` the carry-all graph keeps neighbors the dense path truncates, and for ``smooth_type_embedding=True`` the graph drops the dense diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index eeb97dcd2c..f74c911ff5 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -499,8 +499,7 @@ def freeze( Lower-level export form: ``"nlist"`` (default, dense neighbor-list lower) or ``"graph"`` (NeighborGraph edge-list lower). ``"graph"`` is only valid for graph-eligible models (``mixed_types`` and ``uses_graph_lower``: - dpa1/se_atten with concat type embedding and no ``exclude_types``, - attention layers included) and selects the C++ graph inference path; + dpa1/se_atten with concat type embedding) and selects the C++ graph inference path; the per-atom virial is enabled for it (near-free in the graph path: one extra scatter off the shared single backward). NOTE: for ``smooth_type_embedding=True`` the carry-all graph attention @@ -585,8 +584,7 @@ def freeze( raise ValueError( "lower_kind='graph' requires a graph-eligible model " "(mixed_types and a descriptor exposing uses_graph_lower()==True, " - "currently dpa1 with tebd_input_mode='concat' and no " - "exclude_types). Use lower_kind='nlist' for this model." + "currently dpa1 with tebd_input_mode='concat'). Use lower_kind='nlist' for this model." ) do_atomic_virial = True diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index d2868a4082..818ff9c46a 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -587,8 +587,8 @@ def _model_uses_graph_lower(model: torch.nn.Module) -> bool: :meth:`~deepmd.pt_expt.model.make_model.make_model..CM._resolve_graph_method` for ``neighbor_graph_method is None`` (the training default): a model is graph-eligible iff it is ``mixed_types`` AND its single descriptor reports - ``uses_graph_lower() == True`` (dpa1/se_atten with concat type embedding - and no ``exclude_types``; attention layers included). + ``uses_graph_lower() == True`` (dpa1/se_atten with concat type embedding; + attention layers included). When True the compiled lower must be the GRAPH ``forward_common_lower_graph`` so the compiled path matches eager training (which already default-flips to From 0aae745172655836d728f9a3b2dd41c66a264e69 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:15:05 +0800 Subject: [PATCH 07/63] feat(neighbor_graph): dispatcher-level pair_excl post-process (task 3b) build_neighbor_graph, build_neighbor_graph_ase, build_neighbor_graph_vesin, build_neighbor_graph_nv all gain optional keyword-only pair_excl=None and compact=False; default path = geometric search then apply_pair_exclusion. _call_common_graph in pt_expt make_model wires atomic_model.pair_excl to every builder call so model-level pair_exclude_types is applied at build time (the atomic-model seam backstop stays as idempotent identity). Oracle tests assert set-equality of the valid-edge set between builder(pair_excl=X) and builder() + separate apply_pair_exclusion(X), for dense (2 id + 3 oracle cases) and ase (2 cases); vesin gets 4 new tests (2 identity, 2 oracle, parametrized over periodic). --- .../utils/neighbor_graph/ase_builder.py | 24 ++- .../dpmodel/utils/neighbor_graph/builder.py | 27 ++- deepmd/pt_expt/model/make_model.py | 13 +- deepmd/pt_expt/utils/nv_graph_builder.py | 24 ++- deepmd/pt_expt/utils/vesin_graph_builder.py | 42 ++++- .../dpmodel/test_neighbor_graph_builder.py | 157 ++++++++++++++++++ .../pt_expt/utils/test_vesin_graph_builder.py | 53 ++++++ 7 files changed, 332 insertions(+), 8 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py index 3b00ee6fac..fa163634a4 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py @@ -24,6 +24,9 @@ from .from_ijs import ( neighbor_graph_from_ijs, ) +from .graph import ( + apply_pair_exclusion, +) if TYPE_CHECKING: from deepmd.dpmodel.array_api import ( @@ -33,6 +36,7 @@ from .graph import ( GraphLayout, NeighborGraph, + PairExcludeMask, ) @@ -42,6 +46,9 @@ def build_neighbor_graph_ase( box: Array | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using ASE's O(N) cell-list search. @@ -66,6 +73,14 @@ def build_neighbor_graph_ase( cutoff radius. layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. Returns ------- @@ -136,6 +151,13 @@ def _to_cpu_numpy(x: Any) -> np.ndarray: i_all, j_all = i_all[keep], j_all[keep] S_all, nframe_all = S_all[keep], nframe_all[keep] - return neighbor_graph_from_ijs( + graph = neighbor_graph_from_ijs( i_all, j_all, S_all, coord, box, nframe_all, nloc, layout=layout ) + if pair_excl is not None: + import array_api_compat + + xp = array_api_compat.array_namespace(coord) + atype_flat = xp.reshape(xp.asarray(atype), (-1,)) + graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) + return graph diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index 71ca699e1b..0b39c1ddcb 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -42,6 +42,7 @@ from .graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, pad_and_guard_edges, ) @@ -50,6 +51,10 @@ Array, ) + from .graph import ( + PairExcludeMask, + ) + def from_dense_quartet( extended_coord: Array, @@ -203,6 +208,9 @@ def build_neighbor_graph( box: Array | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph DIRECTLY from coordinates (``dense`` search). @@ -221,6 +229,11 @@ def build_neighbor_graph( ``method`` key. Edges map every neighbor to its LOCAL owner (``src = mapping[neighbor]``), so the graph is ghost-free. + When ``pair_excl`` is given, :func:`apply_pair_exclusion` is called as a + post-process after the geometric search (the default path). A builder MAY + natively fuse the exclusion into its search in a future PR; the contract is + set-equality of valid-edge sets with the default post-process path. + Parameters ---------- coord @@ -236,6 +249,14 @@ def build_neighbor_graph( at non-binding ``sel``). layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. """ from deepmd.dpmodel.utils.nlist import ( extend_coord_with_ghosts, @@ -298,9 +319,13 @@ def build_neighbor_graph( edge_index, edge_vec, layout.edge_capacity, layout.min_edges ) n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) - return NeighborGraph( + graph = NeighborGraph( n_node=n_node, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, ) + if pair_excl is not None: + atype_flat = xp.reshape(atype, (-1,)) + graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) + return graph diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index ae2e83eada..ee89b7bd78 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -468,22 +468,27 @@ def _call_common_graph( "graph lower (e.g. dpa1 attn_layer=0)" ) rcut = self.get_rcut() + # Model-level pair_exclude_types — apply at build time so the seam + # backstop in forward_atomic_graph acts as an idempotent identity. + pair_excl = getattr(self.atomic_model, "pair_excl", None) if method == "dense": - ng = build_neighbor_graph(cc, atype, bb, rcut) + ng = build_neighbor_graph(cc, atype, bb, rcut, pair_excl=pair_excl) elif method == "ase": - ng = build_neighbor_graph_ase(cc, atype, bb, rcut) + ng = build_neighbor_graph_ase(cc, atype, bb, rcut, pair_excl=pair_excl) elif method == "vesin": from deepmd.pt_expt.utils.vesin_graph_builder import ( build_neighbor_graph_vesin, ) - ng = build_neighbor_graph_vesin(cc, atype, bb, rcut) + ng = build_neighbor_graph_vesin( + cc, atype, bb, rcut, pair_excl=pair_excl + ) elif method == "nv": from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, ) - ng = build_neighbor_graph_nv(cc, atype, bb, rcut) + ng = build_neighbor_graph_nv(cc, atype, bb, rcut, pair_excl=pair_excl) else: raise ValueError( f"unknown neighbor_graph_method {method!r}; " diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index 06a30c7dd4..65db788d13 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -22,14 +22,21 @@ ) from typing import ( + TYPE_CHECKING, Any, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + import torch from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, neighbor_graph_from_ijs, ) from deepmd.pt.utils.nv_nlist import ( @@ -204,6 +211,9 @@ def build_neighbor_graph_nv( box: Any | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using nvalchemiops' GPU cell list. @@ -219,6 +229,14 @@ def build_neighbor_graph_nv( cutoff radius. layout edge-axis length policy; ``None`` => dynamic with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. Returns ------- @@ -275,6 +293,10 @@ def build_neighbor_graph_nv( center_local, src_local = center_local[keep], src_local[keep] shift, frame_idx = shift[keep], frame_idx[keep] - return neighbor_graph_from_ijs( + graph = neighbor_graph_from_ijs( center_local, src_local, shift, coord, box_out, frame_idx, nloc, layout=layout ) + if pair_excl is not None: + at_flat = torch.as_tensor(atype, device=device).reshape(-1) + graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) + return graph diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py index 3f86fc7b75..414b356f68 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -19,15 +19,22 @@ ) from typing import ( + TYPE_CHECKING, Any, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + import array_api_compat import torch from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, neighbor_graph_from_ijs, ) from deepmd.pt_expt.utils.vesin_neighbor_list import ( @@ -89,11 +96,40 @@ def build_neighbor_graph_vesin( box: Any | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using vesin.torch's O(N) cell list. Mirrors :func:`deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph_ase` but runs on the input tensor's device via ``vesin.torch``. + + Parameters + ---------- + coord + (nf, nloc, 3) or (nf, nloc*3) local coordinates (torch tensor). + atype + (nf, nloc) local atom types; ``type < 0`` marks a virtual atom. + box + (nf, 3, 3) simulation cell, or ``None`` for non-periodic. + rcut + cutoff radius. + layout + edge-axis length policy; ``None`` => dynamic with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. + + Returns + ------- + graph + The carry-all :class:`NeighborGraph` over the LOCAL atoms. """ if not is_vesin_torch_available(): raise ImportError( @@ -162,6 +198,10 @@ def build_neighbor_graph_vesin( # (grad-carrying). Unlike the nv builder, vesin's cell list handles # out-of-cell (unwrapped) positions natively, so no normalize_coord is # needed and S is consistent with the original coords as searched. - return neighbor_graph_from_ijs( + graph = neighbor_graph_from_ijs( i_all, j_all, S_all, coord, box, nf_all, nloc, layout=layout ) + if pair_excl is not None: + at_flat = torch.as_tensor(atype, device=dev).reshape(-1) + graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) + return graph diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 9ba25c0ccb..8325408ff9 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -15,8 +15,12 @@ import numpy as np +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, + apply_pair_exclusion, build_neighbor_graph, from_dense_quartet, ) @@ -313,5 +317,158 @@ def test_adapter_maps_ghost_to_local_owner(self) -> None: np.testing.assert_allclose(ev[0], np.array([3.0, 0.0, 0.0])) +def valid_edge_set(ng): + """Return the set of (src, dst, rounded edge_vec) for all real edges.""" + ei = ng.edge_index[:, ng.edge_mask] + ev = ng.edge_vec[ng.edge_mask] + return { + (int(ei[0, k]), int(ei[1, k]), tuple(np.round(ev[k], 6))) + for k in range(ei.shape[1]) + } + + +class TestBuildNeighborGraphPairExclOracle(unittest.TestCase): + """Oracle harness: builder(pair_excl=X) == builder() + apply_pair_exclusion(X). + + Covers both ``pair_excl=None`` (identity; no exclusion applied) and a + non-empty exclusion set, for the ``dense`` backend. The oracle asserts + SET-EQUALITY of the valid-edge set, matching the Task 3b contract. + """ + + def setUp(self) -> None: + self.rcut = 4.0 + # 4 atoms, 2 types (0 and 1); atom 2 offset avoids degenerate rcut alignment. + self.coord = np.array( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.3, 0.0], [3.5, 0.0, 0.0]], + dtype=np.float64, + ).reshape(1, 4, 3) + # type sequence: 0,1,0,1 -- pairs (0,1) and (1,0) are heterogeneous + self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64) + self.ntypes = 2 + + def _pair_excl(self, exclude_pairs): + """Build a PairExcludeMask from a list of (ti, tj) tuples.""" + return PairExcludeMask(self.ntypes, exclude_pairs) + + def test_pair_excl_none_identity_dense(self) -> None: + """pair_excl=None: builder output unchanged (identity).""" + ng_ref = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + ng_excl = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, pair_excl=None + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_excl)) + + def test_pair_excl_empty_list_identity_dense(self) -> None: + """pair_excl with empty exclude set: builder output unchanged.""" + pe = self._pair_excl([]) + ng_ref = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + ng_excl = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_excl)) + + def test_oracle_set_equality_dense_nonperiodic(self) -> None: + """Builder with pair_excl==(0,1) == builder() + apply_pair_exclusion.""" + pe = self._pair_excl([(0, 1), (1, 0)]) + # reference: build without exclusion then apply separately + ng_base = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + atype_flat = self.atype.reshape(-1) + ng_post = apply_pair_exclusion(ng_base, atype_flat, pe) + # under test: builder applies exclusion internally + ng_fused = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_post), valid_edge_set(ng_fused)) + # sanity: exclusion actually REMOVED some edges + self.assertLess(int(ng_fused.edge_mask.sum()), int(ng_base.edge_mask.sum())) + + def test_oracle_set_equality_dense_periodic(self) -> None: + """Periodic PBC: builder with pair_excl==(0,0) == builder() + apply.""" + pe = self._pair_excl([(0, 0)]) + box = np.eye(3, dtype=np.float64)[None] * 6.0 + ng_base = build_neighbor_graph(self.coord, self.atype, box, self.rcut) + atype_flat = self.atype.reshape(-1) + ng_post = apply_pair_exclusion(ng_base, atype_flat, pe) + ng_fused = build_neighbor_graph( + self.coord, self.atype, box, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_post), valid_edge_set(ng_fused)) + # type-0 centers: atoms 0,2; type-0 neighbors excluded; fewer edges expected + self.assertLess(int(ng_fused.edge_mask.sum()), int(ng_base.edge_mask.sum())) + + def test_oracle_set_equality_dense_multiframe(self) -> None: + """Multi-frame: set-equality holds per frame.""" + pe = self._pair_excl([(0, 1), (1, 0)]) + coord2 = np.concatenate([self.coord, self.coord + 0.5], axis=0) + atype2 = np.concatenate([self.atype, self.atype], axis=0) + ng_base = build_neighbor_graph(coord2, atype2, None, self.rcut) + atype_flat = atype2.reshape(-1) + ng_post = apply_pair_exclusion(ng_base, atype_flat, pe) + ng_fused = build_neighbor_graph(coord2, atype2, None, self.rcut, pair_excl=pe) + self.assertEqual(valid_edge_set(ng_post), valid_edge_set(ng_fused)) + + +class TestBuildNeighborGraphAseOracle(unittest.TestCase): + """Oracle harness for the ASE builder pair_excl parameter. + + Skipped when ``ase`` is not installed. Asserts set-equality of the + valid-edge set between the ASE builder called with ``pair_excl`` and + the dense reference builder + separate :func:`apply_pair_exclusion`. + """ + + @classmethod + def setUpClass(cls) -> None: + try: + import ase # noqa: F401 + except ImportError as e: + import unittest + + raise unittest.SkipTest("ase not installed") from e + + def setUp(self) -> None: + self.rcut = 4.0 + self.coord = np.array( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.3, 0.0], [3.5, 0.0, 0.0]], + dtype=np.float64, + ).reshape(1, 4, 3) + self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64) + self.ntypes = 2 + + def _pair_excl(self, exclude_pairs): + return PairExcludeMask(self.ntypes, exclude_pairs) + + def test_ase_pair_excl_none_identity(self) -> None: + """pair_excl=None: ASE builder output unchanged.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph_ase, + ) + + ng_ref = build_neighbor_graph_ase(self.coord, self.atype, None, self.rcut) + ng_excl = build_neighbor_graph_ase( + self.coord, self.atype, None, self.rcut, pair_excl=None + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_excl)) + + def test_ase_oracle_set_equality(self) -> None: + """ASE builder with pair_excl == dense ref + apply_pair_exclusion.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph_ase, + ) + + pe = self._pair_excl([(0, 1), (1, 0)]) + # dense reference + separate post-process + ng_dense = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + atype_flat = self.atype.reshape(-1) + ng_ref = apply_pair_exclusion(ng_dense, atype_flat, pe) + # ASE builder with fused post-process + ng_ase = build_neighbor_graph_ase( + self.coord, self.atype, None, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_ase)) + # exclusion actually removed edges + ng_ase_plain = build_neighbor_graph_ase(self.coord, self.atype, None, self.rcut) + self.assertLess(int(ng_ase.edge_mask.sum()), int(ng_ase_plain.edge_mask.sum())) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/utils/test_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py index dea25adab3..ce9ac567f9 100644 --- a/source/tests/pt_expt/utils/test_vesin_graph_builder.py +++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py @@ -3,7 +3,11 @@ import pytest import torch +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, build_neighbor_graph, ) @@ -105,3 +109,52 @@ def test_vesin_excludes_virtual_atoms_like_dense(): ei = np.asarray(ng.edge_index)[:, np.asarray(ng.edge_mask)] at = atype.reshape(-1).numpy() assert np.all(at[ei[0]] >= 0) and np.all(at[ei[1]] >= 0) + + +def _valid_edge_set(ng): + """Return the set of (src, dst, rounded edge_vec) for all real edges.""" + ei = np.asarray(ng.edge_index) + ev = np.asarray(ng.edge_vec) + em = np.asarray(ng.edge_mask) + return { + (int(ei[0, k]), int(ei[1, k]), tuple(np.round(ev[k], 6))) + for k in range(ei.shape[1]) + if em[k] + } + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC and PBC +def test_vesin_pair_excl_none_identity(periodic): + """pair_excl=None: vesin builder output is unchanged (identity).""" + coord, atype, box = _system(periodic) + coord = coord.reshape(1, 4, 3) + box_3d = None if box is None else box.reshape(1, 3, 3) + ng_ref = vesin_builder.build_neighbor_graph_vesin(coord, atype, box_3d, 2.0) + ng_excl = vesin_builder.build_neighbor_graph_vesin( + coord, atype, box_3d, 2.0, pair_excl=None + ) + assert _valid_edge_set(ng_ref) == _valid_edge_set(ng_excl) + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC and PBC +def test_vesin_pair_excl_oracle_set_equality(periodic): + """Vesin builder(pair_excl=X) == dense ref + apply_pair_exclusion(X).""" + coord, atype, box = _system(periodic) + coord = coord.reshape(1, 4, 3) + box_3d = None if box is None else box.reshape(1, 3, 3) + rcut = 2.0 + pe = PairExcludeMask(2, [(0, 1), (1, 0)]) + # dense reference + separate post-process + ng_dense = build_neighbor_graph(coord, atype, box_3d, rcut) + atype_flat = atype.reshape(-1) + ng_ref = apply_pair_exclusion(ng_dense, atype_flat, pe) + # vesin builder with fused post-process + ng_vesin = vesin_builder.build_neighbor_graph_vesin( + coord, atype, box_3d, rcut, pair_excl=pe + ) + assert _valid_edge_set(ng_ref) == _valid_edge_set(ng_vesin) + # exclusion actually removed edges + ng_plain = vesin_builder.build_neighbor_graph_vesin(coord, atype, box_3d, rcut) + assert int(np.asarray(ng_vesin.edge_mask).sum()) < int( + np.asarray(ng_plain.edge_mask).sum() + ) From c2ab950a1f2e42e03db5dfbd645442d3625bd0ce Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:19:00 +0800 Subject: [PATCH 08/63] test(pt_expt): pair_exclude_types graph-vs-legacy parity + vacuity check --- .../pt_expt/model/test_dpa1_graph_lower.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index 11c49cdca2..ae5b4ca88c 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -330,6 +330,69 @@ def test_smooth_attention_divergence_pinned(self) -> None: assert e_diff < 1e-3, f"smooth divergence too large: {e_diff:.3e}" assert f_diff < 1e-3, f"smooth force divergence too large: {f_diff:.3e}" + def test_pair_exclude_types_graph_vs_legacy(self) -> None: + """Model-level pair_exclude_types: graph route and legacy dense agree + bit-tight (fp64, 1e-12), AND the excluded model output differs from the + no-exclude baseline (exclusion is not vacuous). + + Strategy: build the no-exclude model, serialize it, inject + ``pair_exclude_types=[[0,1]]`` into the serialized dict, deserialize + to get an exclude model with IDENTICAL weights, then run both routes. + """ + import copy + + # 1. build the reference (no-exclude) model + model_ref = self._make_model(attn_layer=0) + model_ref.eval() + + # 2. derive the exclude model by patching the serialized dict + data = copy.deepcopy(model_ref.serialize()) + data["pair_exclude_types"] = [[0, 1]] + model_excl = EnergyModel.deserialize(data).to(self.device) + model_excl.eval() + + tol = ( + {"rtol": 1e-12, "atol": 1e-12} + if self.device.type == "cpu" + else {"rtol": 1e-10, "atol": 1e-10} + ) + box = self.cell.reshape(1, 9) + + # 3. graph route (build-time pair exclusion) + graph_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + # 4. legacy dense route (seam backstop in forward_atomic_graph) + legacy_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + # parity: graph == legacy + torch.testing.assert_close( + graph_out["energy_redu"], legacy_out["energy_redu"], **tol + ) + torch.testing.assert_close( + graph_out["energy_derv_r"], legacy_out["energy_derv_r"], **tol + ) + + # 5. reference (no-exclude) via graph route + ref_out = model_ref.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + # exclusion must have an effect + e_diff = (graph_out["energy_redu"] - ref_out["energy_redu"]).abs().max().item() + assert e_diff > 1e-10, ( + f"pair_exclude_types had no effect on energy; diff={e_diff:.3e}" + ) + @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention def test_graph_route_float32(self, attn_layer) -> None: """A float32 model runs the graph route and matches the dense route. From d0e4ab6a43d493a45515b010c41853df77851534 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:24:09 +0800 Subject: [PATCH 09/63] docs: add notes on nv_graph_builder pair_excl oracle gap --- deepmd/pt_expt/utils/nv_graph_builder.py | 6 ++++++ source/tests/common/dpmodel/test_neighbor_graph_builder.py | 5 +++++ 2 files changed, 11 insertions(+) diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index 65db788d13..fe588c94e7 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -248,6 +248,12 @@ def build_neighbor_graph_nv( ------ ImportError if ``nvalchemi-toolkit-ops`` (CUDA) is not installed. + + Notes + ----- + The ``pair_excl`` path of this builder has no local oracle set-equality test + because nvalchemiops requires CUDA; the set-equality contract must be + validated on a GPU box (same pattern as :class:`~deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph_ase`). """ if not is_nv_available(): raise ImportError( diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 8325408ff9..30e1516a75 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -470,5 +470,10 @@ def test_ase_oracle_set_equality(self) -> None: self.assertLess(int(ng_ase.edge_mask.sum()), int(ng_ase_plain.edge_mask.sum())) +# NOTE: nvalchemiops builder has no local oracle set-equality test for pair_excl +# because it requires CUDA; validation is deferred to GPU box tests (PR-C/nv-gtest). +# See deepmd.pt_expt.utils.nv_graph_builder.build_neighbor_graph_nv docstring. + + if __name__ == "__main__": unittest.main() From 33948328436a245e0768fa226738d37dbeb734f5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:29:39 +0800 Subject: [PATCH 10/63] test(pt_expt): graph-route exclude_types coverage (parity + make_fx) --- source/tests/pt_expt/descriptor/test_dpa1.py | 19 +++++++++++--- .../pt_expt/model/test_dpa1_graph_lower.py | 25 ++++++++++++++++--- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index cddd22419f..053263009a 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -252,14 +252,19 @@ def fn(coord_ext, atype_ext, nlist): atol=atol, ) + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion @pytest.mark.parametrize("prec", ["float64"]) # precision - def test_make_fx_graph(self, prec) -> None: + def test_make_fx_graph(self, prec, excl_types) -> None: """make_fx (export-readiness) of the attn_layer=0 GRAPH forward. For ``attn_layer == 0`` the dense ``forward`` routes through the graph-native path (``from_dense_quartet -> call_graph``). This proves that graph forward + ``autograd.grad`` is fx-traceable (full .pt2 - export is PR-B). + export is PR-B). Parametrized over ``excl_types``: with non-empty + exclusion the ``build_edge_exclude_mask`` (mask-only, shape-static) + path is exercised — a tracing failure here would be a bug. """ rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape @@ -276,6 +281,7 @@ def test_make_fx_graph(self, prec) -> None: self.nt, attn_layer=0, precision=prec, + exclude_types=excl_types, seed=GLOBAL_SEED, ).to(self.device) dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device) @@ -311,9 +317,12 @@ def fn(coord_ext, atype_ext, nlist, mapping): atol=atol, ) + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion @pytest.mark.parametrize("smooth", [False, True]) # smooth attention branch @pytest.mark.parametrize("prec", ["float64"]) # precision - def test_make_fx_graph_attn(self, prec, smooth) -> None: + def test_make_fx_graph_attn(self, prec, smooth, excl_types) -> None: """make_fx (export-readiness) of the GRAPH forward with attention. MERGE BLOCKER (NeighborGraph PR-D): pt_expt compiled training routes @@ -321,6 +330,9 @@ def test_make_fx_graph_attn(self, prec, smooth) -> None: (``attn_layer > 0``) must be fx-traceable — the shape-static ``center_edge_pairs`` form keeps the pair enumeration ``nonzero``-free. Covers both the smooth and non-smooth attention branches. + Parametrized over ``excl_types``: with non-empty exclusion the + ``build_edge_exclude_mask`` (mask-only, shape-static) path is exercised + concurrently with attention — a tracing failure here would be a bug. """ rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape @@ -338,6 +350,7 @@ def test_make_fx_graph_attn(self, prec, smooth) -> None: attn_dotr=True, smooth_type_embedding=smooth, precision=prec, + exclude_types=excl_types, seed=GLOBAL_SEED, ).to(self.device) dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device) diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index ae5b4ca88c..33e5223f4b 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -91,7 +91,12 @@ def setup_method(self) -> None: [[0, 0, 0, 1, 1]], dtype=torch.int64, device=self.device ) - def _make_model(self, attn_layer: int = 0, smooth: bool = False) -> EnergyModel: + def _make_model( + self, + attn_layer: int = 0, + smooth: bool = False, + pair_excl_types: list | None = None, + ) -> EnergyModel: ds = DescrptDPA1( self.rcut, self.rcut_smth, @@ -122,7 +127,12 @@ def _make_model(self, attn_layer: int = 0, smooth: bool = False) -> EnergyModel: precision="float64", seed=GLOBAL_SEED, ).to(self.device) - return EnergyModel(ds, ft, type_map=self.type_map).to(self.device) + return EnergyModel( + ds, + ft, + type_map=self.type_map, + pair_exclude_types=pair_excl_types or [], + ).to(self.device) def _prepare_lower_inputs(self, periodic: bool): """Build extended coords, atype, nlist, mapping as torch tensors.""" @@ -172,13 +182,20 @@ def _prepare_lower_inputs(self, periodic: bool): @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention @pytest.mark.parametrize("periodic", [True, False]) # PBC vs non-PBC @pytest.mark.parametrize("do_av", [False, True]) # atom-virial off / on - def test_force_virial_parity_vs_legacy(self, periodic, do_av, attn_layer) -> None: + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion + def test_force_virial_parity_vs_legacy( + self, periodic, do_av, attn_layer, excl_types + ) -> None: """Graph lower energy/force/virial/atom_virial == legacy dense lower on the SAME neighbor set (regime-1 graph from from_dense_quartet). attn_layer=2 exercises graph attention through model-level autograd (smooth=False: exact carry-all parity regime, NeighborGraph PR-D). + Parametrized over exclude_types: empty list (no exclusion) and + [(0,1)] (model-level pair exclusion applied identically on both routes). """ - model = self._make_model(attn_layer=attn_layer) + model = self._make_model(attn_layer=attn_layer, pair_excl_types=excl_types) model.eval() tol = ( {"rtol": 1e-12, "atol": 1e-12} From ab1af0d9f430b70f6c1783b508b1a7ea2ad01ef4 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:33:07 +0800 Subject: [PATCH 11/63] test(pt_expt): descriptor-level exclude_types export + graph-vs-legacy parity --- source/tests/pt_expt/descriptor/test_dpa1.py | 6 +- .../pt_expt/model/test_dpa1_graph_lower.py | 61 +++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index 053263009a..25b3c47a44 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -114,7 +114,10 @@ def test_consistency(self, idt, sm, to, tm, prec, ect) -> None: @pytest.mark.parametrize("idt", [False, True]) # resnet_dt @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision - def test_exportable(self, idt, prec) -> None: + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion + def test_exportable(self, idt, prec, excl_types) -> None: rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -130,6 +133,7 @@ def test_exportable(self, idt, prec) -> None: attn_layer=2, precision=prec, resnet_dt=idt, + exclude_types=excl_types, seed=GLOBAL_SEED, ).to(self.device) dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device) diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index 33e5223f4b..215ce67120 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -96,6 +96,7 @@ def _make_model( attn_layer: int = 0, smooth: bool = False, pair_excl_types: list | None = None, + descr_excl_types: list | None = None, ) -> EnergyModel: ds = DescrptDPA1( self.rcut, @@ -116,6 +117,7 @@ def _make_model( set_davg_zero=False, type_one_side=True, precision="float64", + exclude_types=descr_excl_types or [], seed=GLOBAL_SEED, ).to(self.device) ft = InvarFitting( @@ -410,6 +412,65 @@ def test_pair_exclude_types_graph_vs_legacy(self) -> None: f"pair_exclude_types had no effect on energy; diff={e_diff:.3e}" ) + @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention + def test_descriptor_exclude_types_graph_vs_legacy(self, attn_layer) -> None: + """Descriptor-level exclude_types: graph route and legacy dense agree + bit-tight (fp64, 1e-12) when exclusion is on the DESCRIPTOR (not the + model pair_exclude_types). Uses identical weights across both routes; + also checks exclusion is non-vacuous vs a no-exclude baseline. + """ + import copy + + # 1. no-exclude model (graph route = reference) + model_ref = self._make_model(attn_layer=attn_layer) + model_ref.eval() + + # 2. exclude model: inject exclude_types into the serialized dict + data = copy.deepcopy(model_ref.serialize()) + data["descriptor"]["exclude_types"] = [[0, 1]] + model_excl = EnergyModel.deserialize(data).to(self.device) + model_excl.eval() + + tol = ( + {"rtol": 1e-12, "atol": 1e-12} + if self.device.type == "cpu" + else {"rtol": 1e-10, "atol": 1e-10} + ) + box = self.cell.reshape(1, 9) + + # 3. graph route + graph_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + # 4. legacy dense route + legacy_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + torch.testing.assert_close( + graph_out["energy_redu"], legacy_out["energy_redu"], **tol + ) + torch.testing.assert_close( + graph_out["energy_derv_r"], legacy_out["energy_derv_r"], **tol + ) + + # 5. exclusion must be non-vacuous + ref_out = model_ref.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + e_diff = (graph_out["energy_redu"] - ref_out["energy_redu"]).abs().max().item() + assert e_diff > 1e-10, ( + f"descriptor exclude_types had no effect on energy; diff={e_diff:.3e}" + ) + @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention def test_graph_route_float32(self, attn_layer) -> None: """A float32 model runs the graph route and matches the dense route. From b855b1954b3bebd3c05ac16174a07106eec38f05 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:37:42 +0800 Subject: [PATCH 12/63] fix: add reduced-virial parity assertion in descriptor exclude_types test --- source/tests/pt_expt/model/test_dpa1_graph_lower.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index 215ce67120..8189384c62 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -15,6 +15,8 @@ reduced (per-frame) virial are frame/local quantities and compare directly. """ +import copy + import numpy as np import pytest import torch @@ -419,8 +421,6 @@ def test_descriptor_exclude_types_graph_vs_legacy(self, attn_layer) -> None: model pair_exclude_types). Uses identical weights across both routes; also checks exclusion is non-vacuous vs a no-exclude baseline. """ - import copy - # 1. no-exclude model (graph route = reference) model_ref = self._make_model(attn_layer=attn_layer) model_ref.eval() @@ -458,6 +458,9 @@ def test_descriptor_exclude_types_graph_vs_legacy(self, attn_layer) -> None: torch.testing.assert_close( graph_out["energy_derv_r"], legacy_out["energy_derv_r"], **tol ) + torch.testing.assert_close( + graph_out["energy_derv_c_redu"], legacy_out["energy_derv_c_redu"], **tol + ) # 5. exclusion must be non-vacuous ref_out = model_ref.call_common( From df9cfbbaf2bf59e6196f91b8d4e71e69fc42d13a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:43:31 +0800 Subject: [PATCH 13/63] fix(neighbor_graph): use logical_and+bool cast in apply_pair_exclusion for array_api_strict compat --- deepmd/dpmodel/utils/neighbor_graph/graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index dd82b545de..bef853682b 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -219,7 +219,7 @@ def apply_pair_exclusion( keep = pair_excl.build_edge_exclude_mask(graph.edge_index, atype) out = dataclasses.replace( graph, - edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), + edge_mask=xp.logical_and(graph.edge_mask, xp.astype(keep, xp.bool)), ) if compact: if graph.angle_index is not None or graph.angle_mask is not None: From facb5983865328c9b8509bf90f0f5c32299d6375 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:58:53 +0800 Subject: [PATCH 14/63] feat(dpmodel): apply_pair_exclusion_nlist helper + pair_excl on build_neighbor_list/strategies (A4) Extract the inline pair-exclusion from base_atomic_model.forward_common_atomic into apply_pair_exclusion_nlist(nlist, atype_ext, pair_excl) in nlist.py. The seam is refactored to call the named helper (idempotent backstop remains). Add pair_excl=None to: - build_neighbor_list (dpmodel, nlist.py) - DefaultNeighborList.build - VesinNeighborList.build (pt_expt) - NvNeighborList.build (pt; CUDA-only, API parity) - NeighborList base class signature 12 new unit tests covering: None/empty identity, excluded pairs -> -1, -1 slot preservation, ghost-atom types, idempotence, torch namespace smoke, build_neighbor_list oracle equivalence, DefaultNeighborList oracle, VesinNeighborList oracle. NvNeighborList CUDA-only (not validated locally). --- .../dpmodel/atomic_model/base_atomic_model.py | 8 +- deepmd/dpmodel/utils/__init__.py | 2 + deepmd/dpmodel/utils/default_neighbor_list.py | 40 ++- deepmd/dpmodel/utils/neighbor_list.py | 11 + deepmd/dpmodel/utils/nlist.py | 54 +++- deepmd/pt/utils/nv_nlist.py | 19 ++ deepmd/pt_expt/utils/vesin_neighbor_list.py | 22 ++ .../test_apply_pair_exclusion_nlist.py | 268 ++++++++++++++++++ 8 files changed, 416 insertions(+), 8 deletions(-) create mode 100644 source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 866bb22329..cbe0bdf5d4 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -36,6 +36,7 @@ from deepmd.dpmodel.utils import ( AtomExcludeMask, PairExcludeMask, + apply_pair_exclusion_nlist, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -297,10 +298,9 @@ def forward_common_atomic( xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist) _, nloc, _ = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) - if self.pair_excl is not None: - pair_mask = self.pair_excl.build_type_exclude_mask(nlist, extended_atype) - # exclude neighbors in the nlist - nlist = xp.where(pair_mask == 1, nlist, -1) + # idempotent backstop: externally-supplied nlists (C++/LAMMPS, call_lower + # users) bypass the in-tree builders and land here still unfiltered. + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, self.pair_excl) ext_atom_mask = self.make_atom_mask(extended_atype) ret_dict = self.forward_atomic( diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 3593af5c16..3e439f173e 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -48,6 +48,7 @@ make_multilayer_network, ) from .nlist import ( + apply_pair_exclusion_nlist, build_multiple_neighbor_list, build_neighbor_list, extend_coord_with_ghosts, @@ -94,6 +95,7 @@ "PairExcludeMask", "SameNlocBatchSampler", "aggregate", + "apply_pair_exclusion_nlist", "build_multiple_neighbor_list", "build_neighbor_graph", "build_neighbor_graph_ase", diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index 3628664c5a..c759e2f21d 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Default all-pairs neighbor-list builder (historical deepmd behavior).""" +from typing import ( + TYPE_CHECKING, +) + import array_api_compat from deepmd.dpmodel.array_api import ( @@ -21,6 +25,9 @@ normalize_coord, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + class DefaultNeighborList(NeighborList): """All-pairs builder: replicate the cell into periodic images and rank by @@ -37,7 +44,36 @@ def build( rcut: float, sel: list[int], return_mode: str = "extended", + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Array, Array, Array, Array] | EdgeNeighborList: + """Build extended coordinates and a candidate neighbor list. + + Parameters + ---------- + coord : Array + Local coordinates, shape ``(nf, nloc, 3)`` or ``(nf, nloc*3)``. + atype : Array + Local atom types, shape ``(nf, nloc)``. + box : Array or None + Simulation cell, shape ``(nf, 3, 3)`` or ``(nf, 9)``; ``None`` + for non-periodic systems. + rcut : float + Cutoff radius. + sel : list[int] + Number of selected neighbors per type. + return_mode : str + Must be ``"extended"`` (the only mode this builder supports). + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list immediately after the geometric search by + :func:`~deepmd.dpmodel.utils.nlist.build_neighbor_list`. + + Returns + ------- + tuple[Array, Array, Array, Array] + ``(extended_coord, extended_atype, nlist, mapping)`` as documented + in :meth:`~deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. + """ if return_mode != "extended": raise NotImplementedError( "DefaultNeighborList only supports the extended-coordinate contract." @@ -54,7 +90,8 @@ def build( extended_coord, extended_atype, mapping = extend_coord_with_ghosts( coord_normalized, atype, box, rcut ) - # types are distinguished in the lower interface, so keep them merged here + # types are distinguished in the lower interface, so keep them merged here; + # pair_excl is forwarded so exclusion is applied at build time. nlist = build_neighbor_list( extended_coord, extended_atype, @@ -62,6 +99,7 @@ def build( rcut, sel, distinguish_types=False, + pair_excl=pair_excl, ) extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) return extended_coord, extended_atype, nlist, mapping diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py index e63f5099dd..37c0f63309 100644 --- a/deepmd/dpmodel/utils/neighbor_list.py +++ b/deepmd/dpmodel/utils/neighbor_list.py @@ -14,6 +14,7 @@ dataclass, ) from typing import ( + TYPE_CHECKING, Literal, ) @@ -21,6 +22,9 @@ Array, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + @dataclass class EdgeNeighborList: @@ -69,6 +73,7 @@ def build( rcut: float, sel: list[int], return_mode: Literal["extended", "edges"] = "extended", + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Array, Array, Array, Array] | EdgeNeighborList: """Build the extended system and a candidate neighbor list. @@ -88,6 +93,12 @@ def build( ``"extended"`` returns the historical extended-coordinate quartet. ``"edges"`` returns :class:`EdgeNeighborList`, where ``edge_vec`` is the only geometric displacement consumed by the model. + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) by calling + :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. + Implementations that do not override this parameter fall back to + the default post-build application in the base interface. Returns ------- diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 59e68a64a0..9d6a95cb48 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + TYPE_CHECKING, Any, ) @@ -17,6 +18,9 @@ to_face_distance, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + def _is_ndtensorflow_namespace(xp: Any) -> bool: return getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow" @@ -71,6 +75,44 @@ def extend_input_and_build_neighbor_list( return extended_coord, extended_atype, mapping, nlist +def apply_pair_exclusion_nlist( + nlist: Array, + atype_ext: Array, + pair_excl: "PairExcludeMask | None", +) -> Array: + """Apply model-level pair-type exclusion to a dense neighbor list. + + Replaces excluded neighbor entries with ``-1`` so that downstream + descriptors see them as empty slots. Identity (returns ``nlist`` + unchanged) when *pair_excl* is ``None`` or its exclude-types list is + empty. + + This is the nlist-representation counterpart of + :func:`deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`. + + Parameters + ---------- + nlist : Array + Dense neighbor list of shape ``(nf, nloc, nnei)``. Entries equal + to ``-1`` indicate empty / padding slots. + atype_ext : Array + Extended atom types of shape ``(nf, nall)``. + pair_excl : PairExcludeMask or None + Exclusion mask object, or ``None`` / empty to skip. + + Returns + ------- + Array + Neighbor list of the same shape with excluded entries set to ``-1``. + Erasing ``-1`` entries a second time is a no-op (idempotent). + """ + if pair_excl is None or len(pair_excl.exclude_types) == 0: + return nlist + xp = array_api_compat.array_namespace(nlist, atype_ext) + pair_mask = pair_excl.build_type_exclude_mask(nlist, atype_ext) + return xp.where(pair_mask == 1, nlist, xp.full_like(nlist, -1)) + + ## translated from torch implementation by chatgpt def build_neighbor_list( coord: Array, @@ -79,6 +121,7 @@ def build_neighbor_list( rcut: float, sel: int | list[int], distinguish_types: bool = True, + pair_excl: "PairExcludeMask | None" = None, ) -> Array: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -100,6 +143,12 @@ def build_neighbor_list( types. distinguish_types : bool distinguish different types. + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) immediately after the + geometric search. This is a convenience shortcut for calling + :func:`apply_pair_exclusion_nlist` separately. ``None`` (default) + leaves the list unchanged. Returns ------- @@ -195,9 +244,8 @@ def build_neighbor_list( ) if distinguish_types: - return nlist_distinguish_types(nlist, atype, sel) - else: - return nlist + nlist = nlist_distinguish_types(nlist, atype, sel) + return apply_pair_exclusion_nlist(nlist, atype, pair_excl) def nlist_distinguish_types( diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index 08fc308f72..80d6c85cde 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -51,6 +51,8 @@ Iterator, ) + from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + @contextlib.contextmanager def _suppress_native_stderr() -> Iterator[None]: @@ -155,11 +157,22 @@ def build( rcut: float, sel: list[int], return_mode: str = "extended", + pair_excl: PairExcludeMask | None = None, ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: """Build the extended system and neighbor list. See :meth:`deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. The returned ``nlist`` is distance-sorted and truncated to ``sum(sel)``. + + Parameters + ---------- + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) by + :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. + ``NvNeighborList`` is CUDA-only; the ``pair_excl`` parameter is + accepted for API parity with the other strategies but cannot be + validated on a CPU-only machine. """ device = coord.device nf, nloc = atype.shape[:2] @@ -207,6 +220,12 @@ def build( nlist = _truncate_to_sel_compiled( extended_coord, nlist, target_neighbors, float(rcut) ) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) return extended_coord, extended_atype, nlist, mapping diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 39a0dff883..1b141ebad1 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -19,6 +19,7 @@ """ from typing import ( + TYPE_CHECKING, Any, ) @@ -28,6 +29,9 @@ EdgeNeighborList, NeighborList, ) + +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_ij_shifts, merge_frame_edge_schemas, @@ -60,6 +64,7 @@ def build( rcut: float, sel: list[int], return_mode: str = "extended", + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: """Build the extended system + candidate neighbor list with vesin. @@ -67,6 +72,16 @@ def build( returned ``nlist`` is distance-sorted and truncated to ``sum(sel)`` (matching the default builder); the lower interface still re-formats / type-splits it. + + Parameters + ---------- + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) by + :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. + The atomic-model seam applies the same filter as an idempotent + backstop, so passing ``pair_excl`` here is a build-time + optimization that avoids re-scanning per forward call. """ is_numpy = not isinstance(coord, torch.Tensor) # vesin runs on the device of the inputs: numpy (the dpmodel backend) is @@ -146,6 +161,13 @@ def build( nlist = torch.stack(nlists, dim=0) mapping = torch.stack(mappings, dim=0) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) + if is_numpy: return ( extended_coord.detach().cpu().numpy(), diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py new file mode 100644 index 0000000000..1ecdf8dfb5 --- /dev/null +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for :func:`apply_pair_exclusion_nlist` and the +``pair_excl`` parameter of :func:`build_neighbor_list` / +``DefaultNeighborList.build``. +""" + +import numpy as np +import pytest + +from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask +from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + build_neighbor_list, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _make_nlist_atype(): + """2-frame, 3-local, 4-neighbor test fixture. + + nlist[f, i, k] = index of the k-th neighbor of atom i in frame f. + -1 = empty slot. + + Frame 0 - atoms 0,1,2 (types 0,1,0); ghost atom 3 (type 1). + Frame 1 - same layout, different neighbor assignment. + """ + # shape (nf=2, nloc=3, nnei=4) + nlist = np.array( + [ + # frame 0 + [[1, 2, 3, -1], [0, 3, -1, -1], [0, 1, -1, -1]], + # frame 1 + [[2, 3, -1, -1], [0, -1, -1, -1], [1, 2, 3, -1]], + ], + dtype=np.int64, + ) + # extended atype: local atoms + 1 ghost; shape (nf=2, nall=4) + # types: atom0=0, atom1=1, atom2=0, ghost3=1 + atype_ext = np.array( + [[0, 1, 0, 1], [0, 1, 0, 1]], + dtype=np.int64, + ) + return nlist, atype_ext + + +# --------------------------------------------------------------------------- +# apply_pair_exclusion_nlist — unit tests +# --------------------------------------------------------------------------- + + +def test_none_is_identity() -> None: + nlist, atype_ext = _make_nlist_atype() + result = apply_pair_exclusion_nlist(nlist, atype_ext, None) + assert result is nlist + + +def test_empty_exclude_list_is_identity() -> None: + nlist, atype_ext = _make_nlist_atype() + excl = PairExcludeMask(2, []) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + assert result is nlist + + +def test_excluded_pairs_become_minus_one() -> None: + """Neighbors of excluded type should become -1.""" + nlist, atype_ext = _make_nlist_atype() + # Exclude (type0, type1) pairs (symmetric: also type1,type0). + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + + # Frame 0, atom 0 (type 0): neighbors 1(type1), 2(type0), 3(type1), -1 + # (0,1) and (0,3) are excluded; (0,2) kept; -1 stays -1 + np.testing.assert_array_equal(result[0, 0], [-1, 2, -1, -1]) + # Frame 0, atom 1 (type 1): neighbors 0(type0), 3(type1), -1, -1 + # (1,0) excluded; (1,3) kept (type1,type1 not excluded) + np.testing.assert_array_equal(result[0, 1], [-1, 3, -1, -1]) + # Frame 0, atom 2 (type 0): neighbors 0(type0), 1(type1), -1, -1 + # (2,0) kept (0,0); (2,1) excluded (0,1) + np.testing.assert_array_equal(result[0, 2], [0, -1, -1, -1]) + + +def test_already_empty_slots_preserved() -> None: + """Entries that are already -1 must remain -1 after exclusion.""" + nlist, atype_ext = _make_nlist_atype() + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + # -1 positions in the original nlist must still be -1 + original_minus_one = nlist == -1 + np.testing.assert_array_equal(result[original_minus_one], -1) + + +def test_ghost_atom_type_respected() -> None: + """Ghost atoms (index >= nloc) are indexed by atype_ext; their types count.""" + nlist, atype_ext = _make_nlist_atype() + # atype_ext[*,3] == 1; atom 0 is type 0. (0,1) is excluded. + # atom 0 in frame 0 lists neighbor 3 (ghost, type 1) -> should be excluded. + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + assert result[0, 0, 2] == -1 # was 3 (type 1), now excluded + + +def test_idempotent_double_application() -> None: + """Applying exclusion twice must give the same result as applying once.""" + nlist, atype_ext = _make_nlist_atype() + excl = PairExcludeMask(2, [(0, 1)]) + once = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + twice = apply_pair_exclusion_nlist(once, atype_ext, excl) + np.testing.assert_array_equal(once, twice) + + +def test_no_matching_pairs_leaves_nlist_unchanged() -> None: + """Exclusion list non-empty but no edge matches — nlist unchanged.""" + nlist, atype_ext = _make_nlist_atype() + # All atoms are type 0 or 1; exclude (2,3) which doesn't appear. + excl = PairExcludeMask(4, [(2, 3)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + np.testing.assert_array_equal(result, nlist) + + +# --------------------------------------------------------------------------- +# apply_pair_exclusion_nlist — torch namespace smoke test +# --------------------------------------------------------------------------- + + +def test_torch_namespace_smoke() -> None: + torch = pytest.importorskip("torch") + nlist, atype_ext = _make_nlist_atype() + nlist_t = torch.from_numpy(nlist) + atype_t = torch.from_numpy(atype_ext) + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist_t, atype_t, excl) + np.testing.assert_array_equal( + result.numpy(), apply_pair_exclusion_nlist(nlist, atype_ext, excl) + ) + + +# --------------------------------------------------------------------------- +# build_neighbor_list pair_excl parameter +# --------------------------------------------------------------------------- + + +def _simple_extended_system(): + """2-atom local system with 1 ghost to give a nontrivial nlist. + + Atoms: local 0 (type 0) at (0,0,0), local 1 (type 1) at (1,0,0). + Ghost 2 (type 1) at (-1,0,0) (periodic image of atom 1 across pbc). + + rcut=1.5: atom 0 sees neighbors 1, 2; atom 1 sees neighbor 0. + """ + # nf=1, nall=3 + coord_ext = np.array( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]], + dtype=np.float64, + ) + atype_ext = np.array([[0, 1, 1]], dtype=np.int64) + return coord_ext, atype_ext + + +def test_build_neighbor_list_pair_excl_equals_post_application() -> None: + """build_neighbor_list(pair_excl=excl) must equal apply_pair_exclusion_nlist + applied to the result without pair_excl (oracle equivalence). + """ + coord_ext, atype_ext = _simple_extended_system() + nloc = 2 + rcut = 1.5 + sel = [4] + excl = PairExcludeMask(2, [(0, 1)]) + + nlist_plain = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False + ) + nlist_excl_builtin = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False, pair_excl=excl + ) + nlist_excl_manual = apply_pair_exclusion_nlist(nlist_plain, atype_ext, excl) + + np.testing.assert_array_equal(nlist_excl_builtin, nlist_excl_manual) + + +def test_build_neighbor_list_none_excl_unchanged() -> None: + """pair_excl=None must not change the nlist.""" + coord_ext, atype_ext = _simple_extended_system() + nloc = 2 + rcut = 1.5 + sel = [4] + + nlist_plain = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False + ) + nlist_with_none = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False, pair_excl=None + ) + np.testing.assert_array_equal(nlist_plain, nlist_with_none) + + +# --------------------------------------------------------------------------- +# DefaultNeighborList.build pair_excl parameter +# --------------------------------------------------------------------------- + + +def _local_system(): + """Return a simple local (non-extended) 2-atom system for DefaultNeighborList.""" + # shape (nf=1, nloc=2, 3) + coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]], dtype=np.float64) + atype = np.array([[0, 1]], dtype=np.int64) + return coord, atype + + +def test_default_neighbor_list_pair_excl_equals_seam() -> None: + """DefaultNeighborList(pair_excl=excl) nlist equals build-then-apply.""" + from deepmd.dpmodel.utils.default_neighbor_list import DefaultNeighborList + + coord, atype = _local_system() + rcut = 1.5 + sel = [4] + excl = PairExcludeMask(2, [(0, 1)]) + + builder = DefaultNeighborList() + # Build without exclusion + ext_coord, ext_atype, nlist_plain, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel + ) + nlist_manual = apply_pair_exclusion_nlist(nlist_plain, ext_atype, excl) + + # Build with exclusion at builder level + _, ext_atype2, nlist_builtin, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel, pair_excl=excl + ) + np.testing.assert_array_equal(ext_atype, ext_atype2) + np.testing.assert_array_equal(nlist_builtin, nlist_manual) + + +# --------------------------------------------------------------------------- +# VesinNeighborList.build pair_excl (torch only) +# --------------------------------------------------------------------------- + + +def test_vesin_neighbor_list_pair_excl_equals_seam() -> None: + """VesinNeighborList(pair_excl=excl) nlist equals build-then-apply.""" + torch = pytest.importorskip("torch") + from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, + ) + + if not is_vesin_torch_available(): + pytest.skip("vesin.torch not installed") + + coord = torch.tensor([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]], dtype=torch.float64) + atype = torch.tensor([[0, 1]], dtype=torch.int64) + rcut = 1.5 + sel = [4] + excl = PairExcludeMask(2, [(0, 1)]) + + builder = VesinNeighborList() + ext_coord, ext_atype, nlist_plain, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel + ) + nlist_manual = apply_pair_exclusion_nlist(nlist_plain, ext_atype, excl) + + _, ext_atype2, nlist_builtin, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel, pair_excl=excl + ) + np.testing.assert_array_equal(nlist_builtin.numpy(), nlist_manual.numpy()) From c960e0eb5bba0d3f038e7b1af165f9f712a10090 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 02:06:17 +0800 Subject: [PATCH 15/63] fix(dpmodel): guard edges+pair_excl, fix base docstring (A4 review) --- deepmd/dpmodel/utils/neighbor_list.py | 3 +-- deepmd/pt/utils/nv_nlist.py | 7 ++++++ deepmd/pt_expt/utils/vesin_neighbor_list.py | 7 ++++++ .../test_apply_pair_exclusion_nlist.py | 25 +++++++++++++++++++ .../pt_expt/utils/test_vesin_graph_builder.py | 12 +++++++++ 5 files changed, 52 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py index 37c0f63309..947e8d2d33 100644 --- a/deepmd/dpmodel/utils/neighbor_list.py +++ b/deepmd/dpmodel/utils/neighbor_list.py @@ -97,8 +97,7 @@ def build( When provided, excluded type pairs are erased from the returned neighbor list (entries set to ``-1``) by calling :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. - Implementations that do not override this parameter fall back to - the default post-build application in the base interface. + Subclasses are expected to apply this filter before returning. Returns ------- diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index 80d6c85cde..c6cb7c4ef8 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -173,7 +173,14 @@ def build( ``NvNeighborList`` is CUDA-only; the ``pair_excl`` parameter is accepted for API parity with the other strategies but cannot be validated on a CPU-only machine. + ``return_mode='edges'`` does not support ``pair_excl``; a + :class:`NotImplementedError` is raised in that combination. """ + if return_mode == "edges" and pair_excl is not None: + raise NotImplementedError( + "pair_excl is not supported with return_mode='edges'; " + "use apply_pair_exclusion (graph variant) on the returned EdgeNeighborList." + ) device = coord.device nf, nloc = atype.shape[:2] target_neighbors = int(sum(sel)) diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 1b141ebad1..4965dbf77a 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -82,7 +82,14 @@ def build( The atomic-model seam applies the same filter as an idempotent backstop, so passing ``pair_excl`` here is a build-time optimization that avoids re-scanning per forward call. + ``return_mode='edges'`` does not support ``pair_excl``; a + :class:`NotImplementedError` is raised in that combination. """ + if return_mode == "edges" and pair_excl is not None: + raise NotImplementedError( + "pair_excl is not supported with return_mode='edges'; " + "use apply_pair_exclusion (graph variant) on the returned EdgeNeighborList." + ) is_numpy = not isinstance(coord, torch.Tensor) # vesin runs on the device of the inputs: numpy (the dpmodel backend) is # bridged through CPU torch; torch tensors stay on their own device. Pin diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py index 1ecdf8dfb5..fde86e7530 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py @@ -266,3 +266,28 @@ def test_vesin_neighbor_list_pair_excl_equals_seam() -> None: coord, atype, box=None, rcut=rcut, sel=sel, pair_excl=excl ) np.testing.assert_array_equal(nlist_builtin.numpy(), nlist_manual.numpy()) + + +# --------------------------------------------------------------------------- +# NvNeighborList.build: edges + pair_excl raises +# --------------------------------------------------------------------------- + + +def test_nv_nlist_edges_pair_excl_raises(): + """NvNeighborList.build raises NotImplementedError for edges+pair_excl. + + The guard fires before any CUDA search, so this test runs on CPU. + NvNeighborList requires CUDA to produce results, but the early-exit + raise is device-independent. + """ + import torch + + from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + from deepmd.pt.utils.nv_nlist import NvNeighborList + + coord = torch.zeros((1, 4, 3), dtype=torch.float64) + atype = torch.zeros((1, 4), dtype=torch.int64) + pe = PairExcludeMask(2, [(0, 1)]) + nl = NvNeighborList() + with pytest.raises(NotImplementedError, match="return_mode='edges'"): + nl.build(coord, atype, None, 2.0, [4], return_mode="edges", pair_excl=pe) diff --git a/source/tests/pt_expt/utils/test_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py index ce9ac567f9..41cf0245d3 100644 --- a/source/tests/pt_expt/utils/test_vesin_graph_builder.py +++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py @@ -158,3 +158,15 @@ def test_vesin_pair_excl_oracle_set_equality(periodic): assert int(np.asarray(ng_vesin.edge_mask).sum()) < int( np.asarray(ng_plain.edge_mask).sum() ) + + +def test_vesin_nlist_edges_pair_excl_raises(): + """VesinNeighborList.build with return_mode='edges' and pair_excl raises NotImplementedError.""" + from deepmd.pt_expt.utils.vesin_neighbor_list import VesinNeighborList + + coord = torch.zeros((1, 4, 3), dtype=torch.float64) + atype = torch.zeros((1, 4), dtype=torch.int64) + pe = PairExcludeMask(2, [(0, 1)]) + nl = VesinNeighborList() + with pytest.raises(NotImplementedError, match="return_mode='edges'"): + nl.build(coord, atype, None, 2.0, [4], return_mode="edges", pair_excl=pe) From 398e6b8c627bb8fbf319b34a756c6fa309ead5a0 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 02:11:41 +0800 Subject: [PATCH 16/63] feat(pt_expt): serialize pair_exclude_types into .pt2 metadata + C++ cross-refs Serialize the model-level pair_exclude_types into metadata.json so the C++ DeepPotPTExpt loader can rebuild the flat (ntypes+1)^2 keep table and re-apply the same pair-exclusion transform at the ingestion seam. Add See Also cross-references between the Python transforms (apply_pair_exclusion, apply_pair_exclusion_nlist) and their forthcoming C++ twins. --- deepmd/dpmodel/utils/neighbor_graph/graph.py | 9 ++++++++ deepmd/dpmodel/utils/nlist.py | 8 +++++++ deepmd/pt_expt/utils/serialization.py | 23 ++++++++++++++++++++ 3 files changed, 40 insertions(+) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index bef853682b..9586461312 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -210,6 +210,15 @@ def apply_pair_exclusion( NeighborGraph A ``dataclasses.replace`` copy (or the original ``graph`` on early exit) with the exclusion applied. + + See Also + -------- + C++ twin ``applyPairExclusion`` in ``source/api_cc/include/commonPT.h`` + The inference-path mirror. Same argument order (edge_index, edge_mask, + atype, ...), same variable names (``type_ij``, ``keep``): it computes + ``type_ij = atype[dst]*(ntypes+1) + atype[src]`` and ANDs the flat + ``(ntypes+1)^2`` table lookup into ``edge_mask`` (mask-only mode; no + compact variant on the compiled path). """ import dataclasses diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 9d6a95cb48..d4330d21a9 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -90,6 +90,14 @@ def apply_pair_exclusion_nlist( This is the nlist-representation counterpart of :func:`deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`. + See Also + -------- + C++ twin ``applyPairExclusionNlist`` in ``source/api_cc/include/commonPT.h`` + The inference-path mirror. Same argument order (nlist, atype_ext, ...), + same variable names (``type_ij``, ``keep``): it computes ``type_ij`` + from the center/neighbor types via the flat ``(ntypes+1)^2`` table and + replaces excluded entries with ``-1``. + Parameters ---------- nlist : Array diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 9317e87f4d..ef7a1d69b6 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -738,6 +738,29 @@ def _probe_has_message_passing(obj: object) -> bool | None: # "graph" → NeighborGraph (atype, n_node, edge_index, edge_vec, edge_mask) # The C++ loader branches on this to build the matching inputs. meta["lower_input_kind"] = "graph" if lower_kind == "graph" else "nlist" + + # Model-level pair-type exclusion (``pair_exclude_types``): a list of + # ``[ti, tj]`` type pairs whose interaction is dropped. The compiled AOTI + # forward ALREADY applies this exclusion internally (the graph seam + # ``apply_pair_exclusion`` / dense seam ``apply_pair_exclusion_nlist`` is + # traced into the exported artifact), so this field is redundant for the + # compiled math. It is serialized so that the C++ loaders + # (``DeepPotPTExpt::init``) can rebuild the flat ``(ntypes+1)^2`` keep table + # and re-apply the SAME transform (``applyPairExclusion`` / + # ``applyPairExclusionNlist`` in ``commonPT.h``) at the ingestion seam as an + # idempotent backstop, keeping the C++ ingestion path side-by-side + # reviewable with the Python transforms. Descriptor-level ``exclude_types`` + # needs NO metadata: it is fully inside the compiled graph. + pair_exclude_types: list[list[int]] = [] + for obj in ( + getattr(model, "atomic_model", None), + model, + ): + pet = getattr(obj, "pair_exclude_types", None) + if pet: + pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] + break + meta["pair_exclude_types"] = pair_exclude_types return meta From aab77c71a763bb782d97bf76b97898b5744f25e5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 02:17:03 +0800 Subject: [PATCH 17/63] feat(api_cc): C++ pair-exclusion twins at the pt_expt ingestion seam Add buildPairExcludeTable / applyPairExclusion (graph) / applyPairExclusionNlist (dense) in commonPT.h, structurally mirroring the Python transforms (apply_pair_exclusion, apply_pair_exclusion_nlist) with the same argument order and variable names (type_ij, keep). DeepPotPTExpt::init rebuilds the flat (ntypes+1)^2 keep table from the pair_exclude_types metadata; the seam applies it before every model call (graph and dense, LAMMPS and standalone) as an idempotent backstop to the exclusion already compiled into the .pt2. --- source/api_cc/include/DeepPotPTExpt.h | 7 ++ source/api_cc/include/commonPT.h | 141 ++++++++++++++++++++++++++ source/api_cc/src/DeepPotPTExpt.cc | 46 +++++++-- 3 files changed, 188 insertions(+), 6 deletions(-) diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index ddaea35646..1c3400cc4b 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -331,6 +331,13 @@ class DeepPotPTExpt : public DeepPotBackend { // continue to work; GNN archives must be regenerated to opt into // the fail-fast guard against the silent-corruption bug. bool has_message_passing_ = false; + // Flat (ntypes+1)^2 model-level pair-type keep table, rebuilt in ``init`` from + // the ``pair_exclude_types`` metadata field (see + // ``deepmd::buildPairExcludeTable``). Empty => no model-level exclusion. + // Applied at the C++ ingestion seam (``applyPairExclusion`` graph / + // ``applyPairExclusionNlist`` dense) as an idempotent backstop; the compiled + // .pt2 graph already applies the same transform internally. + std::vector pair_exclude_table_; std::unique_ptr with_comm_tempfile_; std::unique_ptr with_comm_loader; diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 02c25aa047..c379e9acc9 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -7,8 +7,10 @@ #include #include #include +#include #include #include +#include #include #include "common.h" @@ -462,6 +464,145 @@ inline GraphTensorPack buildGraphTensors( return pack; } +/** + * @brief Build the flat ``(ntypes+1)^2`` pair-type keep table. + * + * Inference-path mirror of the Python ``PairExcludeMask`` constructor + * (``deepmd/dpmodel/utils/exclude_mask.py``). The table is row-major over + * ``[tj][ti]`` (flat index ``tj * (ntypes+1) + ti``); an entry is ``0`` when the + * ordered pair ``(ti, tj)`` is excluded and ``1`` otherwise. Both ``(ti, tj)`` + * and ``(tj, ti)`` are inserted into the exclude set, so the table is + * symmetric. Type ``ntypes`` is the reserved virtual-atom row/column. + * + * Returns an empty vector when ``exclude_types`` is empty, so callers can treat + * an empty table as "no exclusion" (identity) just like the Python + * ``pair_excl is None`` early-exit. + * + * @param ntypes Number of real atom types. + * @param exclude_types List of excluded ``(ti, tj)`` type pairs. + */ +inline std::vector buildPairExcludeTable( + const int ntypes, const std::vector>& exclude_types) { + if (exclude_types.empty()) { + return {}; + } + const int n1 = ntypes + 1; + std::set> excl; + for (const auto& tt : exclude_types) { + excl.insert({tt.first, tt.second}); + excl.insert({tt.second, tt.first}); + } + // type_mask[tj][ti] == 0 iff (ti, tj) is excluded (mirrors the Python + // list comprehension in PairExcludeMask.__init__, reshape(-1)). + std::vector type_mask(static_cast(n1) * n1, 1); + for (int tj = 0; tj < n1; ++tj) { + for (int ti = 0; ti < n1; ++ti) { + if (excl.count({ti, tj})) { + type_mask[static_cast(tj) * n1 + ti] = 0; + } + } + } + return type_mask; +} + +/** + * @brief Graph pair-type exclusion: AND the per-edge keep-mask into + * ``edge_mask``. + * + * Inference-path twin of Python ``apply_pair_exclusion`` (mask-only mode) in + * ``deepmd/dpmodel/utils/neighbor_graph/graph.py`` + + * ``PairExcludeMask.build_edge_exclude_mask``. Kept side-by-side reviewable: + * same argument order (edge_index, edge_mask, atype, ...) and same variable + * names (``type_ij``, ``keep``). + * + * The compiled ``.pt2`` graph already applies this exclusion internally (the + * seam transform is traced into the exported forward), so this call is an + * idempotent backstop at the C++ ingestion seam. + * + * @param edge_index (2, E) int64 ``[src, dst]``; src = neighbor, dst = center. + * @param edge_mask (E,) bool real-vs-padding mask to be ANDed in place. + * @param atype (N,) int64 flat node types (clamped >= 0). + * @param type_mask_table Flat ``(ntypes+1)^2`` table from + * ``buildPairExcludeTable``. Empty => identity (returns ``edge_mask``). + * @param ntypes Number of real atom types. + * @return New ``edge_mask`` with excluded edges cleared. + */ +inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& atype, + const std::vector& type_mask_table, + const int ntypes) { + if (type_mask_table.empty()) { + return edge_mask; + } + const auto device = edge_mask.device(); + const auto src = edge_index.index({0}); // (E,) neighbour + const auto dst = edge_index.index({1}); // (E,) center + const auto src_t = atype.index_select(0, src); + const auto dst_t = atype.index_select(0, dst); + // type_ij = atype[dst] * (ntypes + 1) + atype[src] (matches Python) + const auto type_ij = dst_t * (ntypes + 1) + src_t; + const auto table = + torch::from_blob(const_cast(type_mask_table.data()), + {static_cast(type_mask_table.size())}, + torch::TensorOptions().dtype(torch::kInt32)) + .clone() + .to(device); + const auto keep = table.index_select(0, type_ij).to(torch::kBool); + return torch::logical_and(edge_mask, keep); +} + +/** + * @brief Dense-nlist pair-type exclusion: erase excluded neighbours to ``-1``. + * + * Inference-path twin of Python ``apply_pair_exclusion_nlist`` in + * ``deepmd/dpmodel/utils/nlist.py`` + ``PairExcludeMask.build_type_exclude_mask``. + * Same argument order (nlist, atype_ext, ...) and same variable names + * (``type_ij``, ``keep``). Idempotent: erasing ``-1`` a second time is a no-op. + * + * @param nlist (nf, nloc, nnei) int64 neighbour list; ``-1`` == empty slot. + * @param atype_ext (nf, nall) int64 extended atom types. + * @param type_mask_table Flat ``(ntypes+1)^2`` table from + * ``buildPairExcludeTable``. Empty => identity (returns ``nlist``). + * @param ntypes Number of real atom types. + * @return New neighbour list with excluded entries set to ``-1``. + */ +inline torch::Tensor applyPairExclusionNlist( + const torch::Tensor& nlist, + const torch::Tensor& atype_ext, + const std::vector& type_mask_table, + const int ntypes) { + if (type_mask_table.empty()) { + return nlist; + } + const auto device = nlist.device(); + const std::int64_t nf = nlist.size(0); + const std::int64_t nloc = nlist.size(1); + const std::int64_t nnei = nlist.size(2); + const std::int64_t nall = atype_ext.size(1); + // center types: first nloc extended atoms. type_i = atype * (ntypes + 1). + const auto type_i = atype_ext.slice(1, 0, nloc) * (ntypes + 1); // (nf, nloc) + // append virtual atom of type ntypes; map -1 neighbours to it. + const auto ae = torch::cat( + {atype_ext, torch::full({nf, 1}, ntypes, atype_ext.options())}, 1); + const auto nlist_for_type = + torch::where(nlist == -1, torch::full_like(nlist, nall), nlist); + const auto type_j = torch::gather( + ae.unsqueeze(1).expand({nf, nloc, nall + 1}), 2, nlist_for_type); + // type_ij = type_i * (ntypes + 1) + type_j (matches Python: type_i already + // scaled above; here just add the neighbour type). + const auto type_ij = type_i.unsqueeze(2) + type_j; // (nf, nloc, nnei) + const auto table = + torch::from_blob(const_cast(type_mask_table.data()), + {static_cast(type_mask_table.size())}, + torch::TensorOptions().dtype(torch::kInt32)) + .clone() + .to(device); + const auto keep = + table.index_select(0, type_ij.reshape({-1})).reshape({nf, nloc, nnei}); + return torch::where(keep == 1, nlist, torch::full_like(nlist, -1)); +} + /** * @brief Remap NeighborGraph (graph-schema) public outputs onto the dense * internal-key layout the rest of ``compute`` consumes. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index cff23388b6..ad09ae9365 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -200,6 +200,21 @@ void DeepPotPTExpt::init(const std::string& model, // scripts and carry the explicit value. has_message_passing_ = metadata.obj_val.count("has_message_passing") && metadata["has_message_passing"].as_bool(); + + // Model-level pair-type exclusion table. ``pair_exclude_types`` is a list of + // [ti, tj] pairs; rebuild the flat (ntypes+1)^2 keep table exactly like the + // Python ``PairExcludeMask`` ctor so the ingestion seam can re-apply the same + // exclusion (idempotent backstop; the compiled graph already applies it). + { + std::vector> pair_exclude_types; + if (metadata.obj_val.count("pair_exclude_types")) { + for (const auto& v : metadata["pair_exclude_types"].as_array()) { + pair_exclude_types.emplace_back(v[0].as_int(), v[1].as_int()); + } + } + pair_exclude_table_ = deepmd::buildPairExcludeTable(ntypes, + pair_exclude_types); + } if (has_comm_artifact_) { try { // Extract the nested ``extra/forward_lower_with_comm.pt2`` into a @@ -875,12 +890,22 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, torch::full({1}, n_node_count, int_option).to(device); at::Tensor node_atype = atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); + // Model-level pair exclusion at the ingestion seam (idempotent backstop; + // the compiled graph already applies the same transform internally). + const at::Tensor graph_edge_mask = + deepmd::applyPairExclusion(edge_tensors.edge_index, + edge_tensors.edge_mask, node_atype, + pair_exclude_table_, ntypes); flat_outputs = run_model_graph(node_atype, n_node_tensor, edge_tensors.edge_index, - edge_tensors.edge_vec, edge_tensors.edge_mask, - fparam_tensor, aparam_tensor, charge_spin_tensor); + edge_tensors.edge_vec, graph_edge_mask, fparam_tensor, + aparam_tensor, charge_spin_tensor); } else { - flat_outputs = run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, + // Model-level pair exclusion at the dense ingestion seam (idempotent + // backstop; the compiled dense forward already applies the same erase). + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); + flat_outputs = run_model(coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, fparam_tensor, aparam_tensor, charge_spin_tensor); } @@ -1238,13 +1263,22 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_index_ext, edge_tensors.edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); } else if (lower_input_is_graph_) { + // Model-level pair exclusion at the ingestion seam (idempotent backstop; + // the compiled graph already applies the same transform internally). + const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( + graph_tensors.edge_index, graph_tensors.edge_mask, graph_tensors.atype, + pair_exclude_table_, ntypes); flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.edge_index, - graph_tensors.edge_vec, graph_tensors.edge_mask, fparam_tensor, - aparam_tensor, charge_spin_tensor); + graph_tensors.edge_vec, graph_edge_mask, fparam_tensor, aparam_tensor, + charge_spin_tensor); } else { + // Model-level pair exclusion at the dense ingestion seam (idempotent + // backstop; the compiled dense forward already applies the same erase). + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + nlist_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = - run_model(coord_Tensor, atype_Tensor, nlist_tensor, mapping_tensor, + run_model(coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, fparam_tensor, aparam_tensor, charge_spin_tensor); } From c1e97db2cd6b435f6859441836f50dfa40fc2c2b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 02:32:09 +0800 Subject: [PATCH 18/63] test(api_cc): gtest + gen script for C++ pair-exclusion seam (dpa1) gen_dpa1_pairexcl.py exports graph-route and dense-route DPA1(attn_layer=0) .pt2 models with model-level pair_exclude_types=[[0,1]] plus a no-exclusion baseline; references come from Python DeepEval of each model. test_deeppot_dpa1_pairexcl_ptexpt.cc validates both C++ ingestion routes (applyPairExclusion / applyPairExclusionNlist) against the Python references at 1e-10 (fp64), cross-checks graph==dense, and proves the exclusion is active by comparing against the empty-table baseline. Wired into test_cc_local.sh. 8/8 tests pass locally. --- .../test_deeppot_dpa1_pairexcl_ptexpt.cc | 176 ++++++++++++++++ source/install/test_cc_local.sh | 3 + source/tests/infer/gen_dpa1_pairexcl.py | 188 ++++++++++++++++++ 3 files changed, 367 insertions(+) create mode 100644 source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc create mode 100644 source/tests/infer/gen_dpa1_pairexcl.py diff --git a/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc new file mode 100644 index 0000000000..18a3016b03 --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test the C++ model-level pair-exclusion ingestion seam of the pt_expt +// backend (Task A3/A4). Two DPA1(attn_layer=0) models with identical weights +// and model-level pair_exclude_types=[[0,1]] are exported, one through each C++ +// ingestion route: +// - deeppot_dpa1_pairexcl_graph.pt2 -> applyPairExclusion (graph route) +// - deeppot_dpa1_pairexcl_nlist.pt2 -> applyPairExclusionNlist (dense route) +// A no-exclusion baseline (deeppot_dpa1_pairexcl_none.pt2, empty exclude table) +// exercises the identity/pre-change branch of both helpers. +// +// The compiled .pt2 forward ALREADY applies the exclusion internally (the seam +// transform is traced into the exported artifact), so the C++ helpers are an +// idempotent backstop; the reference values (.expected sidecars) come from the +// Python DeepEval of the SAME .pt2, so a 1e-10 match validates the whole chain +// (pair_exclude_types metadata round-trip + init table build + seam apply + +// compiled math). A separate assertion (excluded energy != baseline energy) +// proves the exclusion is genuinely active and not silently dropped. +#include + +#include +#include + +#include "DeepPot.h" +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "test_utils.h" + +namespace { +constexpr const char* kGraphModel = + "../../tests/infer/deeppot_dpa1_pairexcl_graph.pt2"; +constexpr const char* kNlistModel = + "../../tests/infer/deeppot_dpa1_pairexcl_nlist.pt2"; +constexpr const char* kNoneModel = + "../../tests/infer/deeppot_dpa1_pairexcl_none.pt2"; +constexpr const char* kGraphRef = + "../../tests/infer/deeppot_dpa1_pairexcl_graph.expected"; +constexpr const char* kNlistRef = + "../../tests/infer/deeppot_dpa1_pairexcl_nlist.expected"; +} // namespace + +template +class TestInferDpa1PairExclPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + + // Excluded models (one per ingestion route) + no-exclusion baseline. + static deepmd::DeepPot dp_graph; + static deepmd::DeepPot dp_nlist; + static deepmd::DeepPot dp_none; + + static void SetUpTestSuite() { +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT + dp_graph.init(kGraphModel); + dp_nlist.init(kNlistModel); + dp_none.init(kNoneModel); +#endif + } + + static void TearDownTestSuite() { + dp_graph = deepmd::DeepPot(); + dp_nlist = deepmd::DeepPot(); + dp_none = deepmd::DeepPot(); + } + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + } + + // Load per-atom reference from a .expected sidecar and reduce to totals. + void load_ref(const char* path, + double& tot_e, + std::vector& per_f, + std::vector& tot_v, + int& natoms) { + deepmd_test::ExpectedRef ref; + ref.load(path); + const auto per_e = ref.get("pbc", "expected_e"); + per_f = ref.get("pbc", "expected_f"); + const auto per_v = ref.get("pbc", "expected_v"); + natoms = per_e.size(); + tot_e = 0.; + for (int ii = 0; ii < natoms; ++ii) { + tot_e += per_e[ii]; + } + tot_v.assign(9, 0.); + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 9; ++dd) { + tot_v[dd] += per_v[ii * 9 + dd]; + } + } + } + + // Run one model through the standalone build-nlist path and check it against + // its Python DeepEval reference at EPSILON. + void check_against_ref(deepmd::DeepPot& dp, const char* ref_path) { + double tot_e; + std::vector per_f, tot_v; + int natoms; + load_ref(ref_path, tot_e, per_f, tot_v, natoms); + + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - per_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - tot_v[ii]), EPSILON); + } + } +}; + +template +deepmd::DeepPot TestInferDpa1PairExclPtExpt::dp_graph; +template +deepmd::DeepPot TestInferDpa1PairExclPtExpt::dp_nlist; +template +deepmd::DeepPot TestInferDpa1PairExclPtExpt::dp_none; + +TYPED_TEST_SUITE(TestInferDpa1PairExclPtExpt, ValueTypes); + +// Graph route: applyPairExclusion at the ingestion seam + compiled exclusion. +TYPED_TEST(TestInferDpa1PairExclPtExpt, graph_route_matches_python_ref) { + this->check_against_ref(this->dp_graph, kGraphRef); +} + +// Dense route: applyPairExclusionNlist at the ingestion seam + compiled +// exclusion. +TYPED_TEST(TestInferDpa1PairExclPtExpt, nlist_route_matches_python_ref) { + this->check_against_ref(this->dp_nlist, kNlistRef); +} + +// The two ingestion routes carry the SAME weights and exclusion, so at +// non-binding sel they must agree bit-for-bit (fp64 ~1e-10). +TYPED_TEST(TestInferDpa1PairExclPtExpt, graph_equals_nlist_route) { + using VALUETYPE = TypeParam; + double e_g, e_n; + std::vector f_g, v_g, f_n, v_n; + this->dp_graph.compute(e_g, f_g, v_g, this->coord, this->atype, this->box); + this->dp_nlist.compute(e_n, f_n, v_n, this->coord, this->atype, this->box); + EXPECT_LT(fabs(e_g - e_n), EPSILON); + ASSERT_EQ(f_g.size(), f_n.size()); + for (size_t ii = 0; ii < f_g.size(); ++ii) { + EXPECT_LT(fabs(f_g[ii] - f_n[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(v_g[ii] - v_n[ii]), EPSILON); + } +} + +// The no-exclusion baseline exercises the EMPTY-table (identity) branch of the +// C++ helpers; it must run cleanly and produce an energy that DIFFERS from the +// excluded models (proving pair_exclude_types is genuinely active, not dropped). +TYPED_TEST(TestInferDpa1PairExclPtExpt, exclusion_is_active_vs_baseline) { + using VALUETYPE = TypeParam; + double e_none, e_g, e_n; + std::vector f, v; + this->dp_none.compute(e_none, f, v, this->coord, this->atype, this->box); + this->dp_graph.compute(e_g, f, v, this->coord, this->atype, this->box); + this->dp_nlist.compute(e_n, f, v, this->coord, this->atype, this->box); + + EXPECT_TRUE(std::isfinite(e_none)); + // Excluding all O-H pairs changes the energy well above the fp64 tolerance. + EXPECT_GT(fabs(e_g - e_none), 1e-6); + EXPECT_GT(fabs(e_n - e_none), 1e-6); +} diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 2247dcb46c..2b1c7fddb2 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -96,7 +96,10 @@ else: env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4.py & PID9=$! + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa1_pairexcl.py & + PID10=$! wait $PID9 + wait $PID10 env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_spin.py & PID7=$! diff --git a/source/tests/infer/gen_dpa1_pairexcl.py b/source/tests/infer/gen_dpa1_pairexcl.py new file mode 100644 index 0000000000..609ad91ddd --- /dev/null +++ b/source/tests/infer/gen_dpa1_pairexcl.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate DPA1 test models carrying model-level ``pair_exclude_types``. + +Produces two graph-eligible DPA1(attn_layer=0) models with identical weights, +one exported through each C++ ingestion route, plus a no-exclusion baseline: + + - deeppot_dpa1_pairexcl_graph.pt2 (lower_kind="graph", pair_exclude=[[0,1]]) + - deeppot_dpa1_pairexcl_nlist.pt2 (lower_kind="nlist", pair_exclude=[[0,1]]) + - deeppot_dpa1_pairexcl_none.pt2 (lower_kind="graph", NO exclusion) + +The pair models exercise the C++ pair-exclusion ingestion seam: +``applyPairExclusion`` (graph route) / ``applyPairExclusionNlist`` (dense route) +plus the ``pair_exclude_types`` metadata round-trip in ``DeepPotPTExpt::init``. +The compiled ``.pt2`` forward ALREADY applies the exclusion internally (the seam +transform is traced into the exported artifact), so the C++ seam is an +idempotent backstop; the gtest validates C++ energy/force vs the Python +``DeepEval`` reference at 1e-10 and, by comparing against the ``_none`` baseline, +confirms the exclusion is actually active. + +Reference sidecar files (.expected) consumed by the C++ gtest are written from +the Python ``DeepEval`` evaluation of each pair model (PBC + NoPBC sections). + +exclude_types = [[0, 1]] drops every O-H (cross-type) interaction while keeping +O-O and H-H, so both energy and forces change measurably but stay non-degenerate +(H-H pairs survive). + +This is a SEPARATE script from ``gen_dpa1.py`` so it can be run independently on +boxes where the unrelated attn_layer=2 model in ``gen_dpa1.py`` trips a torch +inductor codegen bug. +""" + +import copy +import os +import sys + +import numpy as np + +# Ensure the source tree is on the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) + +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + + +def main(): + from deepmd.dpmodel.model.model import ( + get_model, + ) + + ensure_inductor_compiler() + + base_dir = os.path.dirname(__file__) + + # Graph-eligible DPA1 (attn_layer=0); same descriptor as gen_dpa1.py + # Section B so the two fixtures stay comparable. + base_config = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut_smth": 2.0, + "rcut": 6.0, + "neuron": [2, 4, 8], + "axis_neuron": 4, + "attn": 5, + "attn_layer": 0, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": True, + "temperature": 1.0, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [5, 5, 5], + "resnet_dt": True, + "seed": 1, + }, + } + + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file as pt_expt_deserialize_to_file, + ) + + load_custom_ops() + + from deepmd.infer import ( + DeepPot, + ) + + coord = np.array( + [ + 12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 0.25, 3.32, 1.68, + 3.36, 3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56, + ], + dtype=np.float64, + ) # fmt: skip + atype = [0, 1, 1, 0, 1, 1] + box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], dtype=np.float64) + + def build_data(exclude_types): + cfg = copy.deepcopy(base_config) + if exclude_types: + cfg["pair_exclude_types"] = exclude_types + model = get_model(copy.deepcopy(cfg)) + return { + "model": copy.deepcopy(model.serialize()), + "model_def_script": cfg, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- No-exclusion baseline (graph route) ---- + none_pt2 = os.path.join(base_dir, "deeppot_dpa1_pairexcl_none.pt2") + print(f"Exporting no-exclusion baseline to {none_pt2} ...") # noqa: T201 + pt_expt_deserialize_to_file( + none_pt2, build_data(None), do_atomic_virial=True, lower_kind="graph" + ) + dp_none = DeepPot(none_pt2) + e_none = dp_none.eval(coord, box, atype, atomic=False)[0][0, 0] + print(f" baseline PBC energy: {e_none:.18e}") # noqa: T201 + + # ---- Pair-exclusion models (graph + dense routes) ---- + exclude_types = [[0, 1]] + data_e = build_data(exclude_types) + for lower_kind, tag in (("graph", "graph"), ("nlist", "nlist")): + pt2_path = os.path.join(base_dir, f"deeppot_dpa1_pairexcl_{tag}.pt2") + print( # noqa: T201 + f"Exporting to {pt2_path} (lower_kind='{lower_kind}', " + f"pair_exclude_types={exclude_types}) ..." + ) + pt_expt_deserialize_to_file( + pt2_path, + copy.deepcopy(data_e), + do_atomic_virial=True, + lower_kind=lower_kind, + ) + dp_e = DeepPot(pt2_path) + e_e1, f_e1, v_e1, ae_e1, av_e1 = dp_e.eval(coord, box, atype, atomic=True) + e_enp, f_enp, v_enp, ae_enp, av_enp = dp_e.eval(coord, None, atype, atomic=True) + + # Confirm the exclusion is ACTIVE: energy must differ from the + # no-exclusion baseline (identical weights minus pair_exclude_types). + e_diff = float(abs(e_e1[0, 0] - e_none)) + print(f" {tag}: |E(excl) - E(no-excl)| = {e_diff:.3e}") # noqa: T201 + if e_diff < 1e-6: + raise RuntimeError( + f"BLOCKED: pair_exclude_types had no effect on the {tag} model " + f"(energy delta {e_diff:.2e} < 1e-6); exclusion may be dropped." + ) + f_max = float(np.max(np.abs(f_e1))) + if f_max < 1e-10: + raise RuntimeError( + f"Pair-exclude {tag} forces are degenerate (max={f_max:.2e})." + ) + + ref_path = os.path.join(base_dir, f"deeppot_dpa1_pairexcl_{tag}.expected") + write_expected_ref( + ref_path, + sections={ + "pbc": { + "expected_e": ae_e1[0, :, 0], + "expected_f": f_e1[0], + "expected_v": av_e1[0], + }, + "nopbc": { + "expected_e": ae_enp[0, :, 0], + "expected_f": f_enp[0], + "expected_v": av_enp[0], + }, + }, + source_script="source/tests/infer/gen_dpa1_pairexcl.py", + ) + print(f" Wrote {ref_path}") # noqa: T201 + + print("\nAll pair-exclude models generated.") # noqa: T201 + print("Done!") # noqa: T201 + + +if __name__ == "__main__": + main() From eced29318c7ca847b3fddebe446c2155146d1526 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 03:02:14 +0800 Subject: [PATCH 19/63] docs: exclude_types is graph-native; fix stale comments in _call_dense and forward_common_atomic_graph --- deepmd/dpmodel/atomic_model/base_atomic_model.py | 2 +- deepmd/dpmodel/descriptor/dpa1.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index cbe0bdf5d4..240d4d6720 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -333,7 +333,7 @@ def forward_common_atomic_graph( ``self.pair_excl is not None``, an edge-keep mask is ANDed into ``graph.edge_mask`` before the descriptor forward, so excluded type-pairs contribute zero to the segment_sum. Descriptor-level ``exclude_types`` is - gated by ``uses_graph_lower()==False``. + handled inside the descriptor's ``call_graph`` (graph-native). Parameters ---------- diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index f0be2eb8b5..9b44eff1d8 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -685,8 +685,8 @@ def _call_dense( atype_ext: Array, nlist: Array, ) -> Array: - """Legacy dense descriptor body (the ineligible ``call`` path: attention, - strip tebd, exclude_types, or the no-mapping ghost case). + """Legacy dense descriptor body (the ineligible ``call`` path: + strip tebd or the no-mapping ghost case). Parameters ---------- From 9f1df4a239ea0059d1360c667972c092f61cdd29 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 03:39:55 +0800 Subject: [PATCH 20/63] fix(spin): force legacy dense routing in SpinModel.call_common The SpinModel backbone (dpa1 attn_layer=0) was being routed to the carry-all graph path by the default-flip (decision #17) introduced in the graph-pair-exclude branch. Virtual/placeholder types injected by get_spin_model double the atom density, making the system sel-binding; the carry-all graph keeps neighbors the capped dense nlist discards, so SpinModel.call_common diverged from call_common_lower (the dense lower used by pt_expt eager inference) by ~2e-6. Fix: pass neighbor_graph_method='legacy' when SpinModel.call_common invokes the backbone, forcing the dense-nlist path until spin-graph support is explicitly implemented. The graph .pt2 export path already has a fail-fast guard for spin (serialization.py:963). Adds a regression test pinning that: - SpinModel.call_common total energy equals backbone(legacy) on the same spin-doubled inputs (exact bit-identity). - The backbone in graph mode gives a DIFFERENT energy at this density, confirming the fixture exercises the diverging regime. --- deepmd/dpmodel/model/spin_model.py | 5 + .../dpmodel/test_spin_model_legacy_routing.py | 154 ++++++++++++++++++ 2 files changed, 159 insertions(+) create mode 100644 source/tests/common/dpmodel/test_spin_model_legacy_routing.py diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 373294bece..6a30434a58 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -628,6 +628,11 @@ def call_common( charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, coord_corr_for_virial=coord_corr_for_virial, + # Spin graph support is not yet implemented; the carry-all graph + # route diverges on sel-binding spin systems (virtual atoms double + # the density). Force the legacy dense-nlist path until spin-graph + # support lands. + neighbor_graph_method="legacy", ) model_output_type = self.backbone_model.model_output_type() if "mask" in model_output_type: diff --git a/source/tests/common/dpmodel/test_spin_model_legacy_routing.py b/source/tests/common/dpmodel/test_spin_model_legacy_routing.py new file mode 100644 index 0000000000..c63efc6f64 --- /dev/null +++ b/source/tests/common/dpmodel/test_spin_model_legacy_routing.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression test: SpinModel backbone must stay on the legacy dense-nlist path. + +A spin system uses virtual/placeholder types whose pair exclusions double the +effective atom density. Before the fix (commit 6c2b007c9), the dpmodel +``call_common`` auto-flip (decision #17) moved graph-eligible mixed_types +backbones to the carry-all graph route. When a spin model's backbone +(dpa1 attn_layer=0) crossed to the graph route, its ``call_common`` diverged +from the dense lower interface (``call_common_lower``) used by pt_expt eager +inference -- the carry-all graph keeps neighbors the capped dense nlist +discards. + +Fix: ``SpinModel.call_common`` forces ``neighbor_graph_method="legacy"`` on +the backbone call so the spin backbone always runs dense regardless of the +default flip. + +This test pins that contract: + +1. ``SpinModel.call_common`` energy == backbone energy computed with explicit + ``neighbor_graph_method="legacy"`` on the spin-doubled coordinate inputs. +2. The backbone in graph mode (``neighbor_graph_method="ase"``) on the same + doubled inputs gives a DIFFERENT energy, confirming the fixture is + sel-binding (i.e., the virtual-atom density really does trigger divergence). +""" + +import numpy as np +import pytest + +from deepmd.dpmodel.model.model import ( + get_spin_model, +) + + +def _spin_dpa1_config() -> dict: + """Minimal dpa1-backed spin model config (sel-binding on doubled inputs).""" + return { + "type": "standard", + "type_map": ["Fe", "H"], + "spin": { + "use_spin": [True, False], + "virtual_scale": [0.3], + }, + "descriptor": { + "type": "dpa1", + "rcut": 4.0, + "rcut_smth": 0.5, + # Small sel: with 8 real + 8 virtual = 16 atoms in a 6 Å box, + # sel=8 is well below the average neighbor count → sel-binding. + "sel": 8, + "ntypes": 4, # expanded to 4 by get_spin_model (2 real + 2 virtual) + "attn_layer": 0, + "axis_neuron": 2, + "neuron": [4, 8], + "seed": 0, + }, + "fitting_net": { + "type": "ener", + "neuron": [4, 4], + "seed": 0, + }, + } + + +def _make_test_frame(rng: np.random.Generator): + """Return (coord, atype, box) for a small PBC cell with 8 atoms.""" + natoms = 8 # 4 Fe + 4 H + coord = rng.random((1, natoms, 3)) * 4.0 # 1 frame + atype = np.array([[0, 0, 0, 0, 1, 1, 1, 1]]) + box = np.eye(3).reshape(1, 9) * 6.0 + return coord, atype, box + + +def test_spin_model_backbone_routes_legacy() -> None: + """SpinModel.call_common energy must equal backbone with explicit legacy. + + Procedure: + - Build the SpinModel and obtain the doubled coords/atypes via + ``process_spin_input`` (the same transform used internally). + - Call the backbone directly with ``neighbor_graph_method="legacy"`` on + those doubled inputs to get the expected dense-path energy. + - Call ``SpinModel.call_common`` and extract its energy. + - Assert the two energies are exactly equal. + - Additionally assert the backbone in ``neighbor_graph_method="ase"`` mode + on the same doubled inputs gives a DIFFERENT energy (sel-binding guard). + """ + pytest.importorskip("ase") # ase builder needed for divergence check + + rng = np.random.default_rng(42) + coord, atype, box = _make_test_frame(rng) + spin = np.zeros_like(coord) + spin[:, :4, 2] = 1.0 # spin-z on Fe atoms only + + model = get_spin_model(_spin_dpa1_config()) + backbone = model.backbone_model + + # --- Get doubled inputs via the model's own transform --- + coord_doubled, atype_doubled, _corr = model.process_spin_input(coord, atype, spin) + + # --- Backbone with explicit legacy routing --- + legacy_ret = backbone.call_common( + coord_doubled, atype_doubled, box, neighbor_graph_method="legacy" + ) + legacy_energy = np.array(legacy_ret["energy"]) + + # --- SpinModel.call_common (must route legacy internally) --- + spin_ret = model.call_common(coord, atype, spin, box) + spin_energy = np.array(spin_ret["energy"]) + + # ``energy`` here is per-atom energy; the backbone returns all 2*nloc + # atoms while the spin model truncates to the first nloc real atoms. + # Compare the total energy (sum over all atoms) which should be equal + # because the virtual-atom energies are zeroed by the exclusion mask. + np.testing.assert_allclose( + float(spin_energy.sum()), + float(legacy_energy.sum()), + rtol=0, + atol=0, + err_msg=( + "SpinModel.call_common total energy does not match backbone(legacy) " + "on doubled inputs; the spin model may be routing through the " + "carry-all graph instead of the legacy dense path." + ), + ) + + # --- Sel-binding guard: backbone graph must DIFFER from backbone legacy --- + graph_ret = backbone.call_common( + coord_doubled, atype_doubled, box, neighbor_graph_method="ase" + ) + graph_energy = np.array(graph_ret["energy"]) + + assert not np.allclose(legacy_energy.sum(), graph_energy.sum(), rtol=1e-6), ( + "Backbone legacy and graph give the same energy on the doubled spin " + "system — sel is not binding with these inputs; the regression fixture " + "is too weak. Reduce sel or increase atom density." + ) + + +def test_spin_model_call_common_deterministic() -> None: + """SpinModel.call_common is deterministic (no stochastic routing).""" + rng = np.random.default_rng(7) + coord, atype, box = _make_test_frame(rng) + spin = np.zeros_like(coord) + spin[:, :4, 2] = 1.0 + + model = get_spin_model(_spin_dpa1_config()) + + ret1 = model.call_common(coord, atype, spin, box) + ret2 = model.call_common(coord, atype, spin, box) + + np.testing.assert_array_equal( + float(np.array(ret1["energy"]).sum()), + float(np.array(ret2["energy"]).sum()), + err_msg="SpinModel.call_common is non-deterministic", + ) From 9a06854e7b6d70746bd574da6c5a3a0e6de05601 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 03:54:29 +0800 Subject: [PATCH 21/63] fix(spin): explicit graph-lower opt-out on backbone descriptor The pair-exclude branch removed exclude_types from DescrptDPA1.uses_graph_lower(), which previously kept spin backbones (they inject exclude_types) on the dense path. As a side effect the descriptor-level dispatch routed spin .pt2/.pte export through the graph kernel, whose scatter/atomic_add tripped a torch-inductor CPU codegen assertion (23 errors in test_deep_eval_spin.py). Add an explicit disable_graph_lower() knob on DescrptDPA1 and set it structurally in SpinModel.__init__ (covers get_spin_model and both dpmodel/pt_expt deserialize paths, so it survives serialize round trips). The flag is not serialized; re-derived at construction. The neighbor_graph_method=legacy kwarg in call_common is kept as belt-and-braces. --- deepmd/dpmodel/descriptor/dpa1.py | 19 ++++++++ deepmd/dpmodel/model/spin_model.py | 15 ++++++- .../dpmodel/test_spin_model_legacy_routing.py | 43 +++++++++++++++++-- 3 files changed, 72 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 9b44eff1d8..64d3589e40 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -356,6 +356,9 @@ def __init__( self.tebd_compress = False self.geo_compress = False self.compress = False + # When set, force the legacy dense lower even if the config would + # otherwise be graph-lower eligible (see ``disable_graph_lower``). + self._graph_lower_disabled = False def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -443,8 +446,24 @@ def uses_graph_lower(self) -> bool: differs from the dense lower by up to ~1e-4 (see the Notes of :meth:`call_graph`). """ + if self._graph_lower_disabled: + return False return self.se_atten.tebd_input_mode == "concat" + def disable_graph_lower(self) -> None: + """Force the legacy dense lower for this descriptor. + + This is an explicit opt-out knob used by contexts where the + graph-native lower is unsupported or undesirable (e.g. spin models, + whose carry-all routing diverges on sel-binding spin systems and + whose ``.pt2``/``.pte`` export trips a torch-inductor scatter/ + atomic_add CPU codegen assertion). After calling this, + :meth:`uses_graph_lower` returns ``False`` regardless of the + descriptor configuration. The flag is not serialized; it is + re-derived structurally at spin-model construction/deserialization. + """ + self._graph_lower_disabled = True + def share_params( self, base_class: "DescrptDPA1", shared_level: int, resume: bool = False ) -> NoReturn: diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 6a30434a58..fa3863320c 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -63,6 +63,16 @@ def __init__( super().__init__() self.backbone_model = backbone_model self.spin = spin + # Spin graph-lower unsupported: carry-all routing diverges on + # sel-binding spin systems and spin export trips inductor scatter + # codegen. Re-derived structurally here so it survives both + # construction and serialize/deserialize round trips (the flag is + # not part of the serialized schema). + dp_atomic_model = self.backbone_model.get_dp_atomic_model() + if dp_atomic_model is not None: + descriptor = getattr(dp_atomic_model, "descriptor", None) + if descriptor is not None and hasattr(descriptor, "disable_graph_lower"): + descriptor.disable_graph_lower() self.ntypes_real = self.spin.ntypes_real self.virtual_scale_mask = self.spin.get_virtual_scale_mask() self.spin_mask = self.spin.get_spin_mask() @@ -630,8 +640,9 @@ def call_common( coord_corr_for_virial=coord_corr_for_virial, # Spin graph support is not yet implemented; the carry-all graph # route diverges on sel-binding spin systems (virtual atoms double - # the density). Force the legacy dense-nlist path until spin-graph - # support lands. + # the density). Belt-and-braces: the backbone descriptor already + # has graph-lower disabled in ``__init__``, but force the legacy + # dense-nlist path here too until spin-graph support lands. neighbor_graph_method="legacy", ) model_output_type = self.backbone_model.model_output_type() diff --git a/source/tests/common/dpmodel/test_spin_model_legacy_routing.py b/source/tests/common/dpmodel/test_spin_model_legacy_routing.py index c63efc6f64..1a309a1ca4 100644 --- a/source/tests/common/dpmodel/test_spin_model_legacy_routing.py +++ b/source/tests/common/dpmodel/test_spin_model_legacy_routing.py @@ -96,6 +96,12 @@ def test_spin_model_backbone_routes_legacy() -> None: # --- Get doubled inputs via the model's own transform --- coord_doubled, atype_doubled, _corr = model.process_spin_input(coord, atype, spin) + # The backbone descriptor has graph-lower disabled by the spin-model + # opt-out knob. Temporarily re-enable it only for the sel-binding + # divergence probe below; the legacy/spin comparison uses the dense path. + backbone_descriptor = backbone.get_dp_atomic_model().descriptor + assert backbone_descriptor._graph_lower_disabled is True + # --- Backbone with explicit legacy routing --- legacy_ret = backbone.call_common( coord_doubled, atype_doubled, box, neighbor_graph_method="legacy" @@ -123,9 +129,14 @@ def test_spin_model_backbone_routes_legacy() -> None: ) # --- Sel-binding guard: backbone graph must DIFFER from backbone legacy --- - graph_ret = backbone.call_common( - coord_doubled, atype_doubled, box, neighbor_graph_method="ase" - ) + # Re-enable graph lower on the backbone descriptor only for this probe. + backbone_descriptor._graph_lower_disabled = False + try: + graph_ret = backbone.call_common( + coord_doubled, atype_doubled, box, neighbor_graph_method="ase" + ) + finally: + backbone_descriptor._graph_lower_disabled = True graph_energy = np.array(graph_ret["energy"]) assert not np.allclose(legacy_energy.sum(), graph_energy.sum(), rtol=1e-6), ( @@ -135,6 +146,32 @@ def test_spin_model_backbone_routes_legacy() -> None: ) +def test_spin_backbone_descriptor_graph_lower_disabled() -> None: + """The spin backbone descriptor must have graph-lower disabled. + + The explicit ``disable_graph_lower`` opt-out knob is set structurally at + spin-model construction and must survive a serialize -> deserialize round + trip (the flag is NOT part of the serialized schema, so it must be + re-derived on deserialization). + """ + model = get_spin_model(_spin_dpa1_config()) + descriptor = model.backbone_model.get_dp_atomic_model().descriptor + # dpa1 with attn_layer=0 concat tebd would otherwise be graph-eligible. + assert descriptor._graph_lower_disabled is True + assert descriptor.uses_graph_lower() is False + + # --- survive serialize -> deserialize --- + data = model.serialize() + from deepmd.dpmodel.model.spin_model import ( + SpinModel, + ) + + model2 = SpinModel.deserialize(data) + descriptor2 = model2.backbone_model.get_dp_atomic_model().descriptor + assert descriptor2._graph_lower_disabled is True + assert descriptor2.uses_graph_lower() is False + + def test_spin_model_call_common_deterministic() -> None: """SpinModel.call_common is deterministic (no stochastic routing).""" rng = np.random.default_rng(7) From 331c9cfec520d4ce0a05ead3aa52f7670c002f26 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 00:58:34 +0000 Subject: [PATCH 22/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/utils/default_neighbor_list.py | 4 +++- deepmd/dpmodel/utils/neighbor_list.py | 4 +++- deepmd/dpmodel/utils/nlist.py | 4 +++- deepmd/pt/utils/nv_nlist.py | 4 +++- source/api_cc/include/DeepPotPTExpt.h | 4 ++-- source/api_cc/include/commonPT.h | 14 ++++++++------ source/api_cc/src/DeepPotPTExpt.cc | 17 ++++++++--------- .../tests/test_deeppot_dpa1_pairexcl_ptexpt.cc | 3 ++- .../common/dpmodel/test_apply_pair_exclusion.py | 4 +++- .../dpmodel/test_apply_pair_exclusion_nlist.py | 17 ++++++++++++----- .../common/dpmodel/test_graph_atomic_parity.py | 6 +++--- .../pt_expt/utils/test_vesin_graph_builder.py | 4 +++- 12 files changed, 53 insertions(+), 32 deletions(-) diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index c759e2f21d..d730b43669 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -26,7 +26,9 @@ ) if TYPE_CHECKING: - from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) class DefaultNeighborList(NeighborList): diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py index 947e8d2d33..c9b6f5923b 100644 --- a/deepmd/dpmodel/utils/neighbor_list.py +++ b/deepmd/dpmodel/utils/neighbor_list.py @@ -23,7 +23,9 @@ ) if TYPE_CHECKING: - from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) @dataclass diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index d4330d21a9..3ee852e40f 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -19,7 +19,9 @@ ) if TYPE_CHECKING: - from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) def _is_ndtensorflow_namespace(xp: Any) -> bool: diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index c6cb7c4ef8..df5f3dc119 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -51,7 +51,9 @@ Iterator, ) - from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) @contextlib.contextmanager diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 1c3400cc4b..6114dca340 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -331,8 +331,8 @@ class DeepPotPTExpt : public DeepPotBackend { // continue to work; GNN archives must be regenerated to opt into // the fail-fast guard against the silent-corruption bug. bool has_message_passing_ = false; - // Flat (ntypes+1)^2 model-level pair-type keep table, rebuilt in ``init`` from - // the ``pair_exclude_types`` metadata field (see + // Flat (ntypes+1)^2 model-level pair-type keep table, rebuilt in ``init`` + // from the ``pair_exclude_types`` metadata field (see // ``deepmd::buildPairExcludeTable``). Empty => no model-level exclusion. // Applied at the C++ ingestion seam (``applyPairExclusion`` graph / // ``applyPairExclusionNlist`` dense) as an idempotent backstop; the compiled diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index c379e9acc9..af0f37df20 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -469,9 +469,9 @@ inline GraphTensorPack buildGraphTensors( * * Inference-path mirror of the Python ``PairExcludeMask`` constructor * (``deepmd/dpmodel/utils/exclude_mask.py``). The table is row-major over - * ``[tj][ti]`` (flat index ``tj * (ntypes+1) + ti``); an entry is ``0`` when the - * ordered pair ``(ti, tj)`` is excluded and ``1`` otherwise. Both ``(ti, tj)`` - * and ``(tj, ti)`` are inserted into the exclude set, so the table is + * ``[tj][ti]`` (flat index ``tj * (ntypes+1) + ti``); an entry is ``0`` when + * the ordered pair ``(ti, tj)`` is excluded and ``1`` otherwise. Both ``(ti, + * tj)`` and ``(tj, ti)`` are inserted into the exclude set, so the table is * symmetric. Type ``ntypes`` is the reserved virtual-atom row/column. * * Returns an empty vector when ``exclude_types`` is empty, so callers can treat @@ -556,9 +556,11 @@ inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, * @brief Dense-nlist pair-type exclusion: erase excluded neighbours to ``-1``. * * Inference-path twin of Python ``apply_pair_exclusion_nlist`` in - * ``deepmd/dpmodel/utils/nlist.py`` + ``PairExcludeMask.build_type_exclude_mask``. - * Same argument order (nlist, atype_ext, ...) and same variable names - * (``type_ij``, ``keep``). Idempotent: erasing ``-1`` a second time is a no-op. + * ``deepmd/dpmodel/utils/nlist.py`` + + * ``PairExcludeMask.build_type_exclude_mask``. Same argument order (nlist, + * atype_ext, ...) and same variable names + * (``type_ij``, ``keep``). Idempotent: erasing ``-1`` a second time is a + * no-op. * * @param nlist (nf, nloc, nnei) int64 neighbour list; ``-1`` == empty slot. * @param atype_ext (nf, nall) int64 extended atom types. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index ad09ae9365..1185dfba80 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -212,8 +212,8 @@ void DeepPotPTExpt::init(const std::string& model, pair_exclude_types.emplace_back(v[0].as_int(), v[1].as_int()); } } - pair_exclude_table_ = deepmd::buildPairExcludeTable(ntypes, - pair_exclude_types); + pair_exclude_table_ = + deepmd::buildPairExcludeTable(ntypes, pair_exclude_types); } if (has_comm_artifact_) { try { @@ -892,10 +892,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); // Model-level pair exclusion at the ingestion seam (idempotent backstop; // the compiled graph already applies the same transform internally). - const at::Tensor graph_edge_mask = - deepmd::applyPairExclusion(edge_tensors.edge_index, - edge_tensors.edge_mask, node_atype, - pair_exclude_table_, ntypes); + const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( + edge_tensors.edge_index, edge_tensors.edge_mask, node_atype, + pair_exclude_table_, ntypes); flat_outputs = run_model_graph(node_atype, n_node_tensor, edge_tensors.edge_index, edge_tensors.edge_vec, graph_edge_mask, fparam_tensor, @@ -905,9 +904,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, // backstop; the compiled dense forward already applies the same erase). const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); - flat_outputs = run_model(coord_Tensor, atype_Tensor, excl_nlist, - mapping_tensor, fparam_tensor, aparam_tensor, - charge_spin_tensor); + flat_outputs = + run_model(coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, + fparam_tensor, aparam_tensor, charge_spin_tensor); } } diff --git a/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc index 18a3016b03..a8f7e7daee 100644 --- a/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc @@ -160,7 +160,8 @@ TYPED_TEST(TestInferDpa1PairExclPtExpt, graph_equals_nlist_route) { // The no-exclusion baseline exercises the EMPTY-table (identity) branch of the // C++ helpers; it must run cleanly and produce an energy that DIFFERS from the -// excluded models (proving pair_exclude_types is genuinely active, not dropped). +// excluded models (proving pair_exclude_types is genuinely active, not +// dropped). TYPED_TEST(TestInferDpa1PairExclPtExpt, exclusion_is_active_vs_baseline) { using VALUETYPE = TypeParam; double e_none, e_g, e_n; diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion.py b/source/tests/common/dpmodel/test_apply_pair_exclusion.py index 2dad917120..d9ceb4a8f5 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -2,7 +2,9 @@ import numpy as np import pytest -from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, apply_pair_exclusion, diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py index fde86e7530..0d4f2c2304 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py @@ -7,13 +7,14 @@ import numpy as np import pytest -from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.nlist import ( apply_pair_exclusion_nlist, build_neighbor_list, ) - # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -212,7 +213,9 @@ def _local_system(): def test_default_neighbor_list_pair_excl_equals_seam() -> None: """DefaultNeighborList(pair_excl=excl) nlist equals build-then-apply.""" - from deepmd.dpmodel.utils.default_neighbor_list import DefaultNeighborList + from deepmd.dpmodel.utils.default_neighbor_list import ( + DefaultNeighborList, + ) coord, atype = _local_system() rcut = 1.5 @@ -282,8 +285,12 @@ def test_nv_nlist_edges_pair_excl_raises(): """ import torch - from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask - from deepmd.pt.utils.nv_nlist import NvNeighborList + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + from deepmd.pt.utils.nv_nlist import ( + NvNeighborList, + ) coord = torch.zeros((1, 4, 3), dtype=torch.float64) atype = torch.zeros((1, 4), dtype=torch.int64) diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index 2608c412c2..1332a36471 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -14,13 +14,13 @@ from deepmd.dpmodel.model.ener_model import ( EnergyModel, ) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( apply_pair_exclusion, from_dense_quartet, ) -from deepmd.dpmodel.utils.exclude_mask import ( - PairExcludeMask, -) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, ) diff --git a/source/tests/pt_expt/utils/test_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py index 41cf0245d3..c216676994 100644 --- a/source/tests/pt_expt/utils/test_vesin_graph_builder.py +++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py @@ -162,7 +162,9 @@ def test_vesin_pair_excl_oracle_set_equality(periodic): def test_vesin_nlist_edges_pair_excl_raises(): """VesinNeighborList.build with return_mode='edges' and pair_excl raises NotImplementedError.""" - from deepmd.pt_expt.utils.vesin_neighbor_list import VesinNeighborList + from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + ) coord = torch.zeros((1, 4, 3), dtype=torch.float64) atype = torch.zeros((1, 4), dtype=torch.int64) From 79d5a8922c61dc29828c7a5b602f16ccae6449f5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 15:38:04 +0800 Subject: [PATCH 23/63] fix(docs+review): numpydoc See Also->Notes for C++ cross-refs; address CodeQL/CodeRabbit - RTD build failed: numpydoc's strict See Also parser rejected the C++-twin cross-reference prose entries ('Error parsing See Also entry ...'). Move the cross-refs to Notes sections (free-form reST) in apply_pair_exclusion and apply_pair_exclusion_nlist. - main.py: drop stale 'no type exclusion' from the graph-eligibility comment (exclude_types is now graph-native; matches the docstring + ValueError). - test_graph_atomic_parity: remove dead ds/ft/am chain (CodeQL unused 'am'). - test_neighbor_graph_builder: drop redundant local 'import unittest' (CodeQL). --- deepmd/dpmodel/utils/neighbor_graph/graph.py | 17 +++++++++-------- deepmd/dpmodel/utils/nlist.py | 17 +++++++++-------- deepmd/pt_expt/entrypoints/main.py | 2 +- .../common/dpmodel/test_graph_atomic_parity.py | 3 --- .../dpmodel/test_neighbor_graph_builder.py | 2 -- 5 files changed, 19 insertions(+), 22 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 9586461312..338eb01182 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -211,14 +211,15 @@ def apply_pair_exclusion( A ``dataclasses.replace`` copy (or the original ``graph`` on early exit) with the exclusion applied. - See Also - -------- - C++ twin ``applyPairExclusion`` in ``source/api_cc/include/commonPT.h`` - The inference-path mirror. Same argument order (edge_index, edge_mask, - atype, ...), same variable names (``type_ij``, ``keep``): it computes - ``type_ij = atype[dst]*(ntypes+1) + atype[src]`` and ANDs the flat - ``(ntypes+1)^2`` table lookup into ``edge_mask`` (mask-only mode; no - compact variant on the compiled path). + Notes + ----- + The C++ inference-path mirror is ``applyPairExclusion`` in + ``source/api_cc/include/commonPT.h``. It uses the same argument order + (edge_index, edge_mask, atype, ...) and the same variable names + (``type_ij``, ``keep``): it computes + ``type_ij = atype[dst]*(ntypes+1) + atype[src]`` and ANDs the flat + ``(ntypes+1)^2`` table lookup into ``edge_mask`` (mask-only mode; no + compact variant on the compiled path). """ import dataclasses diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 3ee852e40f..8516409c1b 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -92,14 +92,6 @@ def apply_pair_exclusion_nlist( This is the nlist-representation counterpart of :func:`deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`. - See Also - -------- - C++ twin ``applyPairExclusionNlist`` in ``source/api_cc/include/commonPT.h`` - The inference-path mirror. Same argument order (nlist, atype_ext, ...), - same variable names (``type_ij``, ``keep``): it computes ``type_ij`` - from the center/neighbor types via the flat ``(ntypes+1)^2`` table and - replaces excluded entries with ``-1``. - Parameters ---------- nlist : Array @@ -115,6 +107,15 @@ def apply_pair_exclusion_nlist( Array Neighbor list of the same shape with excluded entries set to ``-1``. Erasing ``-1`` entries a second time is a no-op (idempotent). + + Notes + ----- + The C++ inference-path mirror is ``applyPairExclusionNlist`` in + ``source/api_cc/include/commonPT.h``. It uses the same argument order + (nlist, atype_ext, ...) and the same variable names (``type_ij``, + ``keep``): it computes ``type_ij`` from the center/neighbor types via + the flat ``(ntypes+1)^2`` table and replaces excluded entries with + ``-1``. """ if pair_excl is None or len(pair_excl.exclude_types) == 0: return nlist diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index f74c911ff5..2e4f747ccb 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -569,7 +569,7 @@ def freeze( m.eval() # The graph lower is opt-in and only valid for graph-eligible models - # (dpa1 with concat tebd and no type exclusion; attention layers included + # (dpa1 with concat tebd, incl. attention layers and exclude_types # -- the carry-all pair enumeration exports via unbacked SymInts). Fail # fast with a clear message rather than emitting a broken .pt2. Enable the # per-atom virial for the graph form -- it is near-free there (one extra diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index 1332a36471..64bfe82f78 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -326,9 +326,6 @@ def test_apply_pair_exclusion_idempotent(pair_exclude_types): rng = np.random.default_rng(42) coord = rng.normal(size=(1, 5, 3)) * 1.5 atype = np.array([[0, 1, 0, 1, 0]], dtype=np.int64) - ds = DescrptDPA1(rcut=4.0, rcut_smth=0.5, sel=[200], ntypes=2, attn_layer=0) - ft = InvarFitting("energy", 2, ds.get_dim_out(), 1, mixed_types=True) - am = DPAtomicModel(ds, ft, type_map=["a", "b"]) ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( coord, atype, 4.0, [200], mixed_types=True, box=None ) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 30e1516a75..8f91f5af8c 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -421,8 +421,6 @@ def setUpClass(cls) -> None: try: import ase # noqa: F401 except ImportError as e: - import unittest - raise unittest.SkipTest("ase not installed") from e def setUp(self) -> None: From cd466790cdea25d92c0c9afd8c000b2f3242aa66 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 00:24:25 +0800 Subject: [PATCH 24/63] feat(dpmodel): compact-mode apply_pair_exclusion supports angle fields Replace the NotImplementedError guard with a real angle remap: after edge compaction, remap angle_index onto the compacted edge axis via an exclusive prefix-sum over surviving edges and drop any angle whose constituent edges were excluded. angle_mask set without angle_index is rejected (nothing to remap). Covers the dpa3/se_t angle channel for the eager/dynamic-nedge path; the compiled/C++ path stays mask-only. --- deepmd/dpmodel/utils/neighbor_graph/graph.py | 58 ++++++++++++---- .../dpmodel/test_apply_pair_exclusion.py | 68 ++++++++++++++++--- 2 files changed, 103 insertions(+), 23 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 338eb01182..b841d0797c 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -203,7 +203,10 @@ def apply_pair_exclusion( (shape-static; the ONLY mode allowed in compiled / AOTI paths). If ``True``, additionally drop masked edges so the returned graph has no padding on the edge axis (data-dependent shape; eager / - dynamic-nedge only). + dynamic-nedge only). When the graph carries angle fields, angles are + remapped onto the compacted edge axis and any angle whose constituent + edges were excluded is dropped (``angle_mask`` present without + ``angle_index`` is rejected: there are no indices to remap). Returns ------- @@ -232,21 +235,46 @@ def apply_pair_exclusion( edge_mask=xp.logical_and(graph.edge_mask, xp.astype(keep, xp.bool)), ) if compact: - if graph.angle_index is not None or graph.angle_mask is not None: - raise NotImplementedError( - "apply_pair_exclusion(compact=True) is not supported when the " - "NeighborGraph carries angle fields (angle_index / angle_mask). " - "Angle indices reference pre-compaction edge positions and would " - "become silently wrong after edge compaction. Either use " - "compact=False (mask-only mode) or strip the angle fields first." - ) (keep_idx,) = xp.nonzero(out.edge_mask) - out = dataclasses.replace( - out, - edge_index=out.edge_index[:, keep_idx], - edge_vec=xp.take(out.edge_vec, keep_idx, axis=0), - edge_mask=xp.take(out.edge_mask, keep_idx, axis=0), - ) + fields = { + "edge_index": out.edge_index[:, keep_idx], + "edge_vec": xp.take(out.edge_vec, keep_idx, axis=0), + "edge_mask": xp.take(out.edge_mask, keep_idx, axis=0), + } + if out.angle_index is not None: + # Angles reference PRE-compaction edge positions; remap them to the + # compacted axis and drop any angle whose constituent edges were + # excluded. ``new_pos`` maps old edge position -> new position via an + # exclusive prefix sum over the survivors (-1 for dropped edges). + surv = xp.astype(out.edge_mask, out.edge_index.dtype) # (E,) 0/1 + rank = xp.cumulative_sum(surv, axis=0) - surv # survivors before me + new_pos = xp.where(out.edge_mask, rank, xp.full_like(rank, -1)) + a_new = xp.take(new_pos, out.angle_index[0, :], axis=0) + b_new = xp.take(new_pos, out.angle_index[1, :], axis=0) + both_survive = xp.logical_and(a_new >= 0, b_new >= 0) + if out.angle_mask is not None: + angle_keep = xp.logical_and( + xp.astype(out.angle_mask, xp.bool), both_survive + ) + else: + angle_keep = both_survive + (angle_keep_idx,) = xp.nonzero(angle_keep) + fields["angle_index"] = xp.stack( + [ + xp.take(a_new, angle_keep_idx, axis=0), + xp.take(b_new, angle_keep_idx, axis=0), + ], + axis=0, + ) + if out.angle_mask is not None: + fields["angle_mask"] = xp.take(out.angle_mask, angle_keep_idx, axis=0) + elif out.angle_mask is not None: + raise ValueError( + "apply_pair_exclusion(compact=True): angle_mask is set without " + "angle_index; cannot remap angles onto the compacted edge axis. " + "Provide angle_index too, or strip angle_mask first." + ) + out = dataclasses.replace(out, **fields) return out diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion.py b/source/tests/common/dpmodel/test_apply_pair_exclusion.py index d9ceb4a8f5..9d55e81a51 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -171,7 +171,7 @@ def test_compact_invariance_vs_mask_only() -> None: # --------------------------------------------------------------------------- -# compact=True with angle fields — must raise NotImplementedError +# compact=True with angle fields — remap onto the compacted edge axis # --------------------------------------------------------------------------- @@ -180,34 +180,86 @@ def _toy_graph_with_angles(): g = _toy_graph() import dataclasses - # Two toy angles (pairs of edges sharing a center) + # Two toy angles (pairs of edges sharing a center), into edge positions [0,5) + # angle0 = (edge0, edge1); angle1 = (edge1, edge2) angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) angle_mask = np.array([1, 1], dtype=np.int32) return dataclasses.replace(g, angle_index=angle_index, angle_mask=angle_mask) -def test_compact_raises_when_angle_index_present() -> None: - """compact=True must raise NotImplementedError when angle_index is set.""" +def test_compact_drops_angles_touching_excluded_edges() -> None: + """compact=True remaps angle_index and drops angles whose edges were excluded.""" g = _toy_graph_with_angles() atype = np.array([0, 1, 0, 1], dtype=np.int64) - with pytest.raises(NotImplementedError, match="angle"): - apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + # exclusion (0,1) masks edges 0 and 3 (see the mask-only test): survivors are + # old edges [1, 2] -> new positions [0, 1]; padding edge 4 is dropped too. + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + # edges compacted to the two survivors + assert out.edge_index.shape[1] == 2 + # angle0 = (edge0, edge1): edge0 excluded -> angle0 DROPPED. + # angle1 = (edge1, edge2): both survive -> kept, remapped (edge1->0, edge2->1). + np.testing.assert_array_equal(out.angle_index, [[0], [1]]) + np.testing.assert_array_equal(out.angle_mask, [1]) + + +def test_compact_remaps_angles_when_no_angle_dropped() -> None: + """When exclusion drops no referenced edge, angles are remapped, none lost.""" + g = _toy_graph_with_angles() + atype = np.array([0, 0, 0, 0], dtype=np.int64) # (0,1) matches nothing + # all 4 real edges kept (new positions == old for [0,1,2,3]); padding dropped. + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + assert out.edge_index.shape[1] == 4 + # both angles survive; indices unchanged (survivor ranks equal old positions) + np.testing.assert_array_equal(out.angle_index, [[0, 1], [1, 2]]) + np.testing.assert_array_equal(out.angle_mask, [1, 1]) + + +def test_compact_angle_index_without_mask_defaults_all_real() -> None: + """angle_index present, angle_mask None: treated as all-real, still remapped.""" + import dataclasses + + g = _toy_graph() + angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) + g2 = dataclasses.replace(g, angle_index=angle_index) # angle_mask stays None + atype = np.array([0, 1, 0, 1], dtype=np.int64) + out = apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + # angle0 touches excluded edge0 -> dropped; angle1 kept & remapped + np.testing.assert_array_equal(out.angle_index, [[0], [1]]) + assert out.angle_mask is None # stays None (all remaining are real) def test_compact_raises_when_only_angle_mask_present() -> None: - """compact=True must raise even when only angle_mask (not angle_index) is set.""" + """compact=True rejects angle_mask set without angle_index (nothing to remap).""" import dataclasses g = _toy_graph() angle_mask = np.array([1], dtype=np.int32) g_with_mask = dataclasses.replace(g, angle_mask=angle_mask) atype = np.array([0, 1, 0, 1], dtype=np.int64) - with pytest.raises(NotImplementedError, match="angle"): + with pytest.raises(ValueError, match="angle_mask is set without"): apply_pair_exclusion( g_with_mask, atype, PairExcludeMask(2, [(0, 1)]), compact=True ) +def test_compact_angle_torch_smoke() -> None: + """compact-mode angle remap runs under the torch namespace.""" + torch = pytest.importorskip("torch") + g = _toy_graph_with_angles() + gt = NeighborGraph( + n_node=torch.from_numpy(g.n_node), + edge_index=torch.from_numpy(g.edge_index), + edge_vec=torch.from_numpy(g.edge_vec), + edge_mask=torch.from_numpy(g.edge_mask), + angle_index=torch.from_numpy(g.angle_index), + angle_mask=torch.from_numpy(g.angle_mask), + ) + atype = torch.tensor([0, 1, 0, 1], dtype=torch.int64) + out = apply_pair_exclusion(gt, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + np.testing.assert_array_equal(out.angle_index.numpy(), [[0], [1]]) + np.testing.assert_array_equal(out.angle_mask.numpy(), [1]) + + def test_compact_works_when_angle_fields_are_none() -> None: """compact=True must NOT raise when angle_index and angle_mask are both None.""" g = _toy_graph() # angle_index=None, angle_mask=None by default From a051090301b52ef80e57ee9df64434b25e885087 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 00:24:45 +0800 Subject: [PATCH 25/63] refactor(graph): apply model-level pair_exclude at graph BUILD, not in the atomic model Model-level pair_exclude_types is a canonical NeighborGraph BUILD transform (decision #18): fold it into edge_mask where the graph is constructed, so the graph lower / exported .pt2 consumes a pre-excluded graph and never re-applies it. Previously it was applied inside forward_common_atomic_graph, which the exported lower routes through -> the .pt2 double-applied it (C++ already masks at build via applyPairExclusion). Idempotent, so numerically identical; this removes the redundant table baked into the .pt2 and unifies all paths on the build-time seam. - dpmodel _call_common_graph: pass pair_excl to the builder (was relying on the atomic-model application; now aligned with pt_expt/C++) - pt_expt DeepEval _build_eval_graph: build pair_excl from the dpmodel/metadata and pass to all 4 backends (dense/ase/vesin/nv) -- previously .pt2 DeepEval inference relied on the in-model application, which is now gone - pt_expt compiled-training graph forward: pass pair_excl at build so eager==compiled holds for pair-excluded models - base_atomic_model.forward_common_atomic_graph: drop the apply_pair_exclusion call (+ unused import); the lower now consumes a pre-excluded graph - direct-lower callers (test_dpa1_graph_lower) apply exclusion at build via apply_pair_exclusion, mirroring the C++ transform after from_dense_quartet - new contract test: the lower does NOT re-apply model-level exclusion --- .../dpmodel/atomic_model/base_atomic_model.py | 17 +++--- deepmd/dpmodel/model/make_model.py | 14 ++++- deepmd/pt_expt/infer/deep_eval.py | 43 +++++++++++-- deepmd/pt_expt/train/training.py | 8 ++- .../dpmodel/test_graph_atomic_parity.py | 60 +++++++++++++++++++ .../pt_expt/model/test_dpa1_graph_lower.py | 5 ++ 6 files changed, 130 insertions(+), 17 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 240d4d6720..19842b91d2 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -9,10 +9,6 @@ Any, ) -from deepmd.dpmodel.utils.neighbor_graph import ( - apply_pair_exclusion, -) - if TYPE_CHECKING: from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, @@ -329,10 +325,10 @@ def forward_common_atomic_graph( The node axis is flat ``(N,)`` (``N = sum(graph.n_node)``); masking and out-stat operate per node. Reuses :meth:`_finalize_atomic_ret`, so virtual-atom masking, ``atom_excl`` and ``apply_out_stat`` match the dense - path. Model-level ``pair_exclude_types`` is graph-native: when - ``self.pair_excl is not None``, an edge-keep mask is ANDed into - ``graph.edge_mask`` before the descriptor forward, so excluded type-pairs - contribute zero to the segment_sum. Descriptor-level ``exclude_types`` is + path. Model-level ``pair_exclude_types`` is applied at graph BUILD time + (decision #18) — folded into ``graph.edge_mask`` by the NeighborGraph + builder (Python) or ``applyPairExclusion`` (C++), NOT here; this method + consumes a pre-excluded graph. Descriptor-level ``exclude_types`` is handled inside the descriptor's ``call_graph`` (graph-native). Parameters @@ -359,7 +355,10 @@ def forward_common_atomic_graph( atype = xp.asarray(atype, device=array_api_compat.device(graph.edge_vec)) atom_mask = self.make_atom_mask(atype) # (N,) bool atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype)) - graph = apply_pair_exclusion(graph, atype_clamped, self.pair_excl) + # NOTE: model-level ``pair_exclude_types`` is NOT applied here. It is a + # graph-BUILD transform (decision #18) already folded into + # ``graph.edge_mask`` by the NeighborGraph builder (Python) or + # ``applyPairExclusion`` (C++); this method consumes a pre-excluded graph. ret_dict = self.forward_atomic_graph( graph, atype_clamped, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index cd70e318b6..f346fada9f 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -460,10 +460,20 @@ def _call_common_graph( "neighbor_graph_method requires a mixed_types descriptor with a " "graph lower (e.g. dpa1 attn_layer=0)" ) + # Model-level ``pair_exclude_types`` is a graph-BUILD transform + # (decision #18): apply it here, at the seam where the NeighborGraph + # is constructed, so the graph lower / exported ``.pt2`` consumes an + # already-excluded ``edge_mask`` and never re-applies it. Mirrors the + # pt_expt eager path and the C++ ``applyPairExclusion`` at build. + pair_excl = getattr(self.atomic_model, "pair_excl", None) if method == "dense": - ng = build_neighbor_graph(cc, atype, bb, self.get_rcut()) + ng = build_neighbor_graph( + cc, atype, bb, self.get_rcut(), pair_excl=pair_excl + ) elif method == "ase": - ng = build_neighbor_graph_ase(cc, atype, bb, self.get_rcut()) + ng = build_neighbor_graph_ase( + cc, atype, bb, self.get_rcut(), pair_excl=pair_excl + ) else: raise ValueError( f"unknown neighbor_graph_method {method!r}; the dpmodel/jax backend " diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 89a724fc2a..989fa83688 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -65,6 +65,9 @@ if TYPE_CHECKING: import ase.neighborlist + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) @@ -1772,19 +1775,26 @@ def _build_eval_graph( selection is a pure performance choice and results are unchanged. """ method = self._neighbor_graph_method + # Model-level ``pair_exclude_types`` is a graph-BUILD transform + # (decision #18): apply it here so the exported ``.pt2`` lower consumes a + # pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++ + # ``applyPairExclusion`` and the eager dpmodel/pt_expt build path). + pair_excl = self._graph_pair_excl() if method == "dense": from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, ) - return build_neighbor_graph(coord_input, atom_types, box_input, self._rcut) + return build_neighbor_graph( + coord_input, atom_types, box_input, self._rcut, pair_excl=pair_excl + ) if method == "ase": from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph_ase, ) return build_neighbor_graph_ase( - coord_input, atom_types, box_input, self._rcut + coord_input, atom_types, box_input, self._rcut, pair_excl=pair_excl ) if method in ("vesin", "nv"): cc = torch.as_tensor(coord_input, dtype=torch.float64, device=device) @@ -1801,17 +1811,42 @@ def _build_eval_graph( build_neighbor_graph_vesin, ) - return build_neighbor_graph_vesin(cc, aa, bb, self._rcut) + return build_neighbor_graph_vesin( + cc, aa, bb, self._rcut, pair_excl=pair_excl + ) from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, ) - return build_neighbor_graph_nv(cc, aa, bb, self._rcut) + return build_neighbor_graph_nv(cc, aa, bb, self._rcut, pair_excl=pair_excl) raise ValueError( f"unknown neighbor_graph_method {method!r}; " "use 'dense', 'ase', 'vesin', or 'nv'" ) + def _graph_pair_excl(self) -> "PairExcludeMask | None": + """Model-level ``pair_exclude_types`` as a ``PairExcludeMask`` (or None). + + Applied at graph BUILD time (decision #18), NOT inside the exported + ``.pt2`` lower. Prefers the loaded dpmodel's mask; otherwise rebuilds it + from the ``pair_exclude_types`` field in ``metadata.json``. + + Returns + ------- + PairExcludeMask | None + The exclusion mask, or ``None`` when the model excludes no pairs. + """ + if self._dpmodel is not None: + return getattr(self._dpmodel.atomic_model, "pair_excl", None) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + pet = self.metadata.get("pair_exclude_types", []) + if not pet: + return None + return PairExcludeMask(len(self._type_map), [tuple(p) for p in pet]) + def _get_output_shape( self, odef: OutputVariableDef, nframes: int, natoms: int ) -> list[int]: diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 818ff9c46a..9edd3e4d4b 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1154,8 +1154,12 @@ def _forward_graph( ) # Carry-all graph (dynamic E, no edge_capacity) — identical to the eager - # uncompiled ``_call_common_graph`` builder so the two paths match. - ng = build_neighbor_graph(coord_3d, atype, box_flat, rcut) + # uncompiled ``_call_common_graph`` builder so the two paths match. Model- + # level pair_exclude is a graph-BUILD transform (decision #18): fold it + # into edge_mask here so the compiled lower consumes a pre-excluded graph + # (the lower no longer re-applies it), matching the eager path exactly. + pair_excl = getattr(_model.atomic_model, "pair_excl", None) + ng = build_neighbor_graph(coord_3d, atype, box_flat, rcut, pair_excl=pair_excl) atype_flat = atype.reshape(nframes * nloc) # Lazy compile of the GRAPH lower (cached per structure key). diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index 64bfe82f78..545f639f80 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -19,6 +19,7 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( apply_pair_exclusion, + build_neighbor_graph, from_dense_quartet, ) from deepmd.dpmodel.utils.nlist import ( @@ -166,6 +167,65 @@ def test_model_pair_exclude_types_graph_matches_dense(): ), "pair exclusion must change the graph energy (same weights)" +def test_model_pair_exclude_applied_at_build_not_in_lower(): + """Seam contract (decision #18): model-level pair_exclude is a graph-BUILD + transform. The graph lower must NOT re-apply it — it consumes whatever + ``edge_mask`` the builder produced. Feeding a NON-excluded graph to the lower + on a model that HAS ``pair_exclude_types`` yields the SAME result as a model + with no exclusion; the exclusion only takes effect when applied at build. + """ + rng = np.random.default_rng(4) + nloc = 6 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0, 1]], dtype=np.int64) + box = np.eye(3).reshape(1, 9) * 20.0 + ds = DescrptDPA1(rcut=4.0, rcut_smth=0.5, sel=[200], ntypes=2, attn_layer=0) + ft = InvarFitting("energy", 2, ds.get_dim_out(), 1, mixed_types=True) + model = EnergyModel(ds, ft, type_map=["a", "b"], pair_exclude_types=[(0, 1)]) + assert model.atomic_model.pair_excl is not None + + # RAW graph: built WITHOUT pair_excl (no exclusion baked into edge_mask). + ng_raw = build_neighbor_graph(coord, atype, box, model.get_rcut()) + kw = { + "atype": atype.reshape(-1), + "n_node": ng_raw.n_node, + "edge_index": ng_raw.edge_index, + "edge_vec": ng_raw.edge_vec, + "edge_mask": ng_raw.edge_mask, + } + out_with_excl_model = model.call_lower_graph(**kw) + # clear the model-level exclusion; the lower output must be UNCHANGED, proving + # the lower never consulted ``pair_excl`` (exclusion is not applied here). + model.atomic_model.reinit_pair_exclude([]) + assert model.atomic_model.pair_excl is None + out_no_excl_model = model.call_lower_graph(**kw) + np.testing.assert_allclose( + np.asarray(out_with_excl_model["energy_redu"]), + np.asarray(out_no_excl_model["energy_redu"]), + rtol=1e-12, + atol=1e-12, + ) + + # Positive control: applying exclusion at BUILD (excluded edge_mask) DOES + # change the same lower's output. + ng_excl = build_neighbor_graph( + coord, atype, box, model.get_rcut(), pair_excl=PairExcludeMask(2, [(0, 1)]) + ) + out_built_excl = model.call_lower_graph( + atype=atype.reshape(-1), + n_node=ng_excl.n_node, + edge_index=ng_excl.edge_index, + edge_vec=ng_excl.edge_vec, + edge_mask=ng_excl.edge_mask, + ) + assert not np.allclose( + np.asarray(out_built_excl["energy_redu"]), + np.asarray(out_no_excl_model["energy_redu"]), + rtol=1e-9, + atol=1e-9, + ), "build-time pair exclusion must change the graph energy" + + def test_graph_matches_dense_with_fparam(): """Frame parameter is gathered to nodes by frame_id in forward_atomic_graph and fed to the fitting's call_graph; the graph path must match dense at 1e-12 diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index 8189384c62..4160993885 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -22,6 +22,7 @@ import torch from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, from_dense_quartet, ) from deepmd.dpmodel.utils.nlist import ( @@ -223,6 +224,10 @@ def test_force_virial_parity_vs_legacy( # returned edge_vec is already a torch tensor on env.DEVICE. ng = from_dense_quartet(ext_coord, nlist, mapping) atype_local = ext_atype[:, :nloc].reshape(nf * nloc) + # Model-level pair_exclude is a BUILD-time transform (decision #18): the + # converter does not bake it in and the lower no longer re-applies it, so + # apply it to the graph here (mirrors the C++ ``applyPairExclusion`` step). + ng = apply_pair_exclusion(ng, atype_local, model.atomic_model.pair_excl) graph = model.forward_common_lower_graph( atype_local, ng.n_node, From b5435318ca06a02c644e936f83dbe5e08e26860a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 00:27:39 +0800 Subject: [PATCH 26/63] fix(dpmodel): fail fast on inconsistent angle fields in compact apply_pair_exclusion angle_index and angle_mask are a coupled pair. compact=True now validates them up front and raises ValueError on any inconsistent state (only one set; A dims disagree; angle_index not (2, A)) instead of silently defaulting a missing mask to all-real. Prevents a partial/mismatched pair from remapping silently wrong. --- deepmd/dpmodel/utils/neighbor_graph/graph.py | 50 ++++++++++++------- .../dpmodel/test_apply_pair_exclusion.py | 40 ++++++++++++--- 2 files changed, 65 insertions(+), 25 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index b841d0797c..d8fcfb780e 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -205,8 +205,9 @@ def apply_pair_exclusion( has no padding on the edge axis (data-dependent shape; eager / dynamic-nedge only). When the graph carries angle fields, angles are remapped onto the compacted edge axis and any angle whose constituent - edges were excluded is dropped (``angle_mask`` present without - ``angle_index`` is rejected: there are no indices to remap). + edges were excluded is dropped. ``angle_index`` and ``angle_mask`` must + be a consistent pair (both set or both ``None``; matching ``A``); + anything else raises ``ValueError``. Returns ------- @@ -235,13 +236,38 @@ def apply_pair_exclusion( edge_mask=xp.logical_and(graph.edge_mask, xp.astype(keep, xp.bool)), ) if compact: + # Angle fields are a coupled pair (produced together by the angle + # builder): both present or both None. Fail fast on any inconsistent + # state — a partial or shape-mismatched pair is a caller bug that would + # otherwise remap silently wrong. + has_ai = out.angle_index is not None + has_am = out.angle_mask is not None + if has_ai != has_am: + raise ValueError( + "apply_pair_exclusion(compact=True): angle_index and angle_mask " + "must both be set or both be None; got " + f"angle_index={'set' if has_ai else 'None'}, " + f"angle_mask={'set' if has_am else 'None'}." + ) + if has_ai: + if out.angle_index.ndim != 2 or out.angle_index.shape[0] != 2: + raise ValueError( + "apply_pair_exclusion(compact=True): angle_index must have " + f"shape (2, A); got {tuple(out.angle_index.shape)}." + ) + if out.angle_index.shape[1] != out.angle_mask.shape[0]: + raise ValueError( + "apply_pair_exclusion(compact=True): angle_index (2, A) and " + f"angle_mask (A,) disagree on A: {out.angle_index.shape[1]} " + f"vs {out.angle_mask.shape[0]}." + ) (keep_idx,) = xp.nonzero(out.edge_mask) fields = { "edge_index": out.edge_index[:, keep_idx], "edge_vec": xp.take(out.edge_vec, keep_idx, axis=0), "edge_mask": xp.take(out.edge_mask, keep_idx, axis=0), } - if out.angle_index is not None: + if has_ai: # Angles reference PRE-compaction edge positions; remap them to the # compacted axis and drop any angle whose constituent edges were # excluded. ``new_pos`` maps old edge position -> new position via an @@ -252,12 +278,9 @@ def apply_pair_exclusion( a_new = xp.take(new_pos, out.angle_index[0, :], axis=0) b_new = xp.take(new_pos, out.angle_index[1, :], axis=0) both_survive = xp.logical_and(a_new >= 0, b_new >= 0) - if out.angle_mask is not None: - angle_keep = xp.logical_and( - xp.astype(out.angle_mask, xp.bool), both_survive - ) - else: - angle_keep = both_survive + angle_keep = xp.logical_and( + xp.astype(out.angle_mask, xp.bool), both_survive + ) (angle_keep_idx,) = xp.nonzero(angle_keep) fields["angle_index"] = xp.stack( [ @@ -266,14 +289,7 @@ def apply_pair_exclusion( ], axis=0, ) - if out.angle_mask is not None: - fields["angle_mask"] = xp.take(out.angle_mask, angle_keep_idx, axis=0) - elif out.angle_mask is not None: - raise ValueError( - "apply_pair_exclusion(compact=True): angle_mask is set without " - "angle_index; cannot remap angles onto the compacted edge axis. " - "Provide angle_index too, or strip angle_mask first." - ) + fields["angle_mask"] = xp.take(out.angle_mask, angle_keep_idx, axis=0) out = dataclasses.replace(out, **fields) return out diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion.py b/source/tests/common/dpmodel/test_apply_pair_exclusion.py index 9d55e81a51..b73f4c8d38 100644 --- a/source/tests/common/dpmodel/test_apply_pair_exclusion.py +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -214,34 +214,58 @@ def test_compact_remaps_angles_when_no_angle_dropped() -> None: np.testing.assert_array_equal(out.angle_mask, [1, 1]) -def test_compact_angle_index_without_mask_defaults_all_real() -> None: - """angle_index present, angle_mask None: treated as all-real, still remapped.""" +def test_compact_raises_when_only_angle_index_present() -> None: + """compact=True fails fast when angle_index is set but angle_mask is None.""" import dataclasses g = _toy_graph() angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) g2 = dataclasses.replace(g, angle_index=angle_index) # angle_mask stays None atype = np.array([0, 1, 0, 1], dtype=np.int64) - out = apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) - # angle0 touches excluded edge0 -> dropped; angle1 kept & remapped - np.testing.assert_array_equal(out.angle_index, [[0], [1]]) - assert out.angle_mask is None # stays None (all remaining are real) + with pytest.raises(ValueError, match="both be set or both be None"): + apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) def test_compact_raises_when_only_angle_mask_present() -> None: - """compact=True rejects angle_mask set without angle_index (nothing to remap).""" + """compact=True fails fast when angle_mask is set but angle_index is None.""" import dataclasses g = _toy_graph() angle_mask = np.array([1], dtype=np.int32) g_with_mask = dataclasses.replace(g, angle_mask=angle_mask) atype = np.array([0, 1, 0, 1], dtype=np.int64) - with pytest.raises(ValueError, match="angle_mask is set without"): + with pytest.raises(ValueError, match="both be set or both be None"): apply_pair_exclusion( g_with_mask, atype, PairExcludeMask(2, [(0, 1)]), compact=True ) +def test_compact_raises_on_angle_dim_mismatch() -> None: + """compact=True fails fast when angle_index (2,A) and angle_mask (A,) disagree.""" + import dataclasses + + g = _toy_graph() + angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) # A == 2 + angle_mask = np.array([1], dtype=np.int32) # A == 1 (mismatch) + g2 = dataclasses.replace(g, angle_index=angle_index, angle_mask=angle_mask) + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(ValueError, match="disagree on A"): + apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + + +def test_compact_raises_on_bad_angle_index_shape() -> None: + """compact=True fails fast when angle_index is not (2, A).""" + import dataclasses + + g = _toy_graph() + angle_index = np.array([[0, 1, 2]], dtype=np.int64) # (1, 3), not (2, A) + angle_mask = np.array([1, 1, 1], dtype=np.int32) + g2 = dataclasses.replace(g, angle_index=angle_index, angle_mask=angle_mask) + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(ValueError, match=r"shape \(2, A\)"): + apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + + def test_compact_angle_torch_smoke() -> None: """compact-mode angle remap runs under the torch namespace.""" torch = pytest.importorskip("torch") From 6ff11d201ab8dd3f790a613fbdce40c8f34afb46 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 00:41:45 +0800 Subject: [PATCH 27/63] fix(pt_expt): DeepEval graph pair_excl uses a numpy-backed mask (device-safe) The DeepEval graph builder applies model-level pair_exclude at build. Reusing the loaded dpmodel's pt_expt-wrapped pair_excl carries a torch (CUDA) type_mask buffer that fails to convert onto a numpy atype (dense/ase build path): 'can't convert cuda:0 tensor to numpy'. Build a fresh numpy PairExcludeMask from the exclude types instead -- it converts cleanly to numpy (dense/ase) or torch (vesin/nv) atype. Refresh the stale gen_dpa1_pairexcl docstring (exclusion is a build-time transform, not baked into the .pt2). --- deepmd/pt_expt/infer/deep_eval.py | 19 ++++++++++++++----- source/tests/infer/gen_dpa1_pairexcl.py | 11 ++++++----- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 989fa83688..5b8df237cc 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1828,21 +1828,30 @@ def _graph_pair_excl(self) -> "PairExcludeMask | None": """Model-level ``pair_exclude_types`` as a ``PairExcludeMask`` (or None). Applied at graph BUILD time (decision #18), NOT inside the exported - ``.pt2`` lower. Prefers the loaded dpmodel's mask; otherwise rebuilds it - from the ``pair_exclude_types`` field in ``metadata.json``. + ``.pt2`` lower. Reads the excluded pairs from the loaded dpmodel (if any) + or the ``pair_exclude_types`` field in ``metadata.json``, and returns a + FRESH numpy-backed mask. + + A numpy ``type_mask`` converts cleanly onto whichever namespace/device the + builder's ``atype`` uses (dense/ase pass numpy; vesin/nv pass torch). The + dpmodel's own ``pair_excl`` is NOT reused: as a pt_expt module attribute + its ``type_mask`` is a torch (possibly CUDA) buffer, which cannot convert + to a numpy ``atype`` on the dense/ase build path. Returns ------- PairExcludeMask | None The exclusion mask, or ``None`` when the model excludes no pairs. """ - if self._dpmodel is not None: - return getattr(self._dpmodel.atomic_model, "pair_excl", None) from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, ) - pet = self.metadata.get("pair_exclude_types", []) + if self._dpmodel is not None: + pe = getattr(self._dpmodel.atomic_model, "pair_excl", None) + pet = pe.get_exclude_types() if pe is not None else [] + else: + pet = self.metadata.get("pair_exclude_types", []) if not pet: return None return PairExcludeMask(len(self._type_map), [tuple(p) for p in pet]) diff --git a/source/tests/infer/gen_dpa1_pairexcl.py b/source/tests/infer/gen_dpa1_pairexcl.py index 609ad91ddd..c37b3d23f1 100644 --- a/source/tests/infer/gen_dpa1_pairexcl.py +++ b/source/tests/infer/gen_dpa1_pairexcl.py @@ -12,11 +12,12 @@ The pair models exercise the C++ pair-exclusion ingestion seam: ``applyPairExclusion`` (graph route) / ``applyPairExclusionNlist`` (dense route) plus the ``pair_exclude_types`` metadata round-trip in ``DeepPotPTExpt::init``. -The compiled ``.pt2`` forward ALREADY applies the exclusion internally (the seam -transform is traced into the exported artifact), so the C++ seam is an -idempotent backstop; the gtest validates C++ energy/force vs the Python -``DeepEval`` reference at 1e-10 and, by comparing against the ``_none`` baseline, -confirms the exclusion is actually active. +Model-level pair exclusion is a graph-BUILD transform (decision #18): it is +folded into ``edge_mask`` at build time (``applyPairExclusion`` in C++, the +NeighborGraph builder in Python ``DeepEval``), NOT inside the exported ``.pt2`` +lower. The gtest validates C++ energy/force vs the Python ``DeepEval`` reference +at 1e-10 and, by comparing against the ``_none`` baseline, confirms the exclusion +is actually active. Reference sidecar files (.expected) consumed by the C++ gtest are written from the Python ``DeepEval`` evaluation of each pair model (PBC + NoPBC sections). From 4f5fdba3c73c8ed37c476ad0771fb84aa8ae4a31 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 00:58:56 +0800 Subject: [PATCH 28/63] DEBUG: instrument applyPairExclusion (REVERT) --- source/api_cc/include/commonPT.h | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index af0f37df20..c14d19db44 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -549,7 +549,21 @@ inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, .clone() .to(device); const auto keep = table.index_select(0, type_ij).to(torch::kBool); - return torch::logical_and(edge_mask, keep); + const auto result = torch::logical_and(edge_mask, keep); + std::cerr << "[DBG applyPairExclusion] table_sz=" << type_mask_table.size() + << " ntypes=" << ntypes << " E=" << edge_index.size(1) + << " atypeN=" << atype.size(0) + << " edge_mask.sum=" << edge_mask.to(torch::kInt64).sum().item() + << " keep.sum=" << keep.to(torch::kInt64).sum().item() + << " result.sum=" << result.to(torch::kInt64).sum().item() + << " atype[:8]="; + { + auto a = atype.to(torch::kCPU).to(torch::kInt64); + int64_t n = std::min(8, a.size(0)); + for (int64_t i = 0; i < n; ++i) std::cerr << a[i].item() << ","; + } + std::cerr << std::endl; + return result; } /** From 1dbe3132f3387f99ee06fe9ec963798bd9a0ea39 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 16:59:58 +0000 Subject: [PATCH 29/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_cc/include/commonPT.h | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index c14d19db44..2d566b0557 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -552,15 +552,17 @@ inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, const auto result = torch::logical_and(edge_mask, keep); std::cerr << "[DBG applyPairExclusion] table_sz=" << type_mask_table.size() << " ntypes=" << ntypes << " E=" << edge_index.size(1) - << " atypeN=" << atype.size(0) - << " edge_mask.sum=" << edge_mask.to(torch::kInt64).sum().item() + << " atypeN=" << atype.size(0) << " edge_mask.sum=" + << edge_mask.to(torch::kInt64).sum().item() << " keep.sum=" << keep.to(torch::kInt64).sum().item() << " result.sum=" << result.to(torch::kInt64).sum().item() << " atype[:8]="; { auto a = atype.to(torch::kCPU).to(torch::kInt64); int64_t n = std::min(8, a.size(0)); - for (int64_t i = 0; i < n; ++i) std::cerr << a[i].item() << ","; + for (int64_t i = 0; i < n; ++i) { + std::cerr << a[i].item() << ","; + } } std::cerr << std::endl; return result; From 931f2097450a8fec50bb72479cf3c85c45a2f8ca Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 01:11:07 +0800 Subject: [PATCH 30/63] Revert "DEBUG: instrument applyPairExclusion (REVERT)" The debug probe confirmed the C++ graph-route pair exclusion IS active (edge_mask 30 -> 14 on the gtest system); the earlier gtest failure was a stale installed libdeepmd_cc.so, not a code bug. All 8 Dpa1PairExcl gtests pass on the Tesla T4 after a full rebuild + reinstall. --- source/api_cc/include/commonPT.h | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 2d566b0557..af0f37df20 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -549,23 +549,7 @@ inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, .clone() .to(device); const auto keep = table.index_select(0, type_ij).to(torch::kBool); - const auto result = torch::logical_and(edge_mask, keep); - std::cerr << "[DBG applyPairExclusion] table_sz=" << type_mask_table.size() - << " ntypes=" << ntypes << " E=" << edge_index.size(1) - << " atypeN=" << atype.size(0) << " edge_mask.sum=" - << edge_mask.to(torch::kInt64).sum().item() - << " keep.sum=" << keep.to(torch::kInt64).sum().item() - << " result.sum=" << result.to(torch::kInt64).sum().item() - << " atype[:8]="; - { - auto a = atype.to(torch::kCPU).to(torch::kInt64); - int64_t n = std::min(8, a.size(0)); - for (int64_t i = 0; i < n; ++i) { - std::cerr << a[i].item() << ","; - } - } - std::cerr << std::endl; - return result; + return torch::logical_and(edge_mask, keep); } /** From e172298ac38bb9eace6c0664fb03e5dd01dddff4 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 03:46:00 +0800 Subject: [PATCH 31/63] refactor(dense): apply model-level pair_exclude at nlist BUILD, completing A4 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decision #18/A4 (user, 2026-07-04): exclusion is applied ONCE when the neighbor list is built, not re-masked per forward — SAME design as the graph route. The A4 execution had wired pair_excl into the builders but left the atomic-model seam as an 'idempotent backstop' and never connected the callers, so in-tree dense ran on the backstop alone. This completes the port and removes the backstop: - base_atomic_model.forward_common_atomic: DROP the internal apply_pair_exclusion_nlist; the dense lower now consumes a pre-excluded nlist (contract documented; negative-contract test added) - dpmodel model_call_from_call_lower: pass pair_excl to builder.build (was never passed); extend_input_and_build_neighbor_list gains pair_excl - out-stat model_forward helper: build with pair_excl - SpinModel.process_spin_input_lower: the virtual-atom nlist extension IS the build site of the spin-extended nlist — fold the backbone's pair_excl in there (universal spin forward==forward_lower holds, 527 tests) - pt_expt: DeepEval dense builders (native strategy + inline + ASE) and the compiled-training dense branch build with pair_excl; _graph_pair_excl renamed _model_pair_excl (serves both routes) - jax: jax2tf TF wrapper gains a TF twin of the erasure (flat keep-table gather); jax2tf serialization + HLO wrapper (pair_excl rebuilt from model_def_script) + trainer prepare_input pass pair_excl at build - C++: applyPairExclusionNlist RESTORED as the single application site on the C++ dense route (its earlier removal was misaligned with A4); all 'idempotent backstop' comments rewritten as build-time ownership statements - tests: dense negative-contract test (lower must NOT re-apply); consistency TestEnerLower harness pre-excludes at build (legacy pt/pd re-apply internally = idempotent no-op, so cross-backend equality holds); test_dp_atomic_model excl-consistency feeds md0 a pre-excluded nlist; pt_expt graph-vs-dense lower parity pre-excludes the dense side --- .../dpmodel/atomic_model/base_atomic_model.py | 18 ++++-- deepmd/dpmodel/model/make_model.py | 13 +++- deepmd/dpmodel/model/spin_model.py | 20 +++++++ deepmd/dpmodel/utils/nlist.py | 16 +++-- deepmd/jax/jax2tf/make_model.py | 42 ++++++++++++- deepmd/jax/jax2tf/serialization.py | 3 + deepmd/jax/model/hlo.py | 22 +++++++ deepmd/jax/train/trainer.py | 11 ++++ deepmd/pt_expt/infer/deep_eval.py | 26 ++++++-- deepmd/pt_expt/train/training.py | 3 + deepmd/pt_expt/utils/serialization.py | 20 +++---- source/api_cc/include/DeepPotPTExpt.h | 7 ++- source/api_cc/include/commonPT.h | 16 +++-- source/api_cc/src/DeepPotPTExpt.cc | 34 +++++++---- .../test_deeppot_dpa1_pairexcl_ptexpt.cc | 16 ++--- .../common/dpmodel/test_dp_atomic_model.py | 17 ++++-- .../dpmodel/test_graph_atomic_parity.py | 59 +++++++++++++++++++ source/tests/consistent/model/test_ener.py | 11 ++++ source/tests/infer/gen_dpa1_pairexcl.py | 22 ++++--- .../pt_expt/model/test_dpa1_graph_lower.py | 8 +++ 20 files changed, 314 insertions(+), 70 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 19842b91d2..2a65906fcc 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -32,7 +32,6 @@ from deepmd.dpmodel.utils import ( AtomExcludeMask, PairExcludeMask, - apply_pair_exclusion_nlist, ) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, @@ -270,7 +269,11 @@ def forward_common_atomic( extended atom typs, shape: nf x nall for a type < 0 indicating the atomic is virtual. nlist - neighbor list, shape: nf x nloc x nsel + neighbor list, shape: nf x nloc x nsel. CONTRACT: model-level + ``pair_exclude_types`` is already folded in (excluded entries are + ``-1``) — exclusion is a nlist-BUILD transform (decision #18/A4) + applied by the NeighborList builders (Python) or + ``applyPairExclusionNlist`` (C++ ingestion), never re-applied here. mapping extended to local index mapping, shape: nf x nall fparam @@ -294,9 +297,11 @@ def forward_common_atomic( xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist) _, nloc, _ = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) - # idempotent backstop: externally-supplied nlists (C++/LAMMPS, call_lower - # users) bypass the in-tree builders and land here still unfiltered. - nlist = apply_pair_exclusion_nlist(nlist, extended_atype, self.pair_excl) + # NOTE: model-level ``pair_exclude_types`` is NOT applied here. It is a + # nlist-BUILD transform (decision #18/A4, same as the graph route): + # already folded into the nlist by the NeighborList builders (Python) + # or ``applyPairExclusionNlist`` (C++ ingestion); this method consumes + # a pre-excluded nlist. ext_atom_mask = self.make_atom_mask(extended_atype) ret_dict = self.forward_atomic( @@ -668,6 +673,9 @@ def model_forward( self.get_sel(), mixed_types=self.mixed_types(), box=box, + # exclusion is a nlist-BUILD transform (decision #18/A4); + # forward_common_atomic consumes a pre-excluded nlist. + pair_excl=self.pair_excl, ) atomic_ret = self.forward_common_atomic( extended_coord, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index f346fada9f..f44d87940f 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -11,6 +11,9 @@ from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) import array_api_compat import numpy as np @@ -85,6 +88,7 @@ def model_call_from_call_lower( coord_corr_for_virial: Array | None = None, charge_spin: Array | None = None, neighbor_list: NeighborList | None = None, + pair_excl: "PairExcludeMask | None" = None, ) -> dict[str, Array]: """Return model prediction from lower interface. @@ -109,6 +113,11 @@ def model_call_from_call_lower( historical behavior. An alternative strategy (e.g. an O(N) cell list) may be injected to speed up neighbor-list construction; it returns the same extended representation, so model outputs are unchanged. + pair_excl + Model-level pair-type exclusion mask. Exclusion is a nlist-BUILD + transform (decision #18/A4): it is folded into the nlist here, at the + build seam, and ``call_lower`` consumes a pre-excluded nlist without + re-applying it. Returns ------- @@ -122,7 +131,7 @@ def model_call_from_call_lower( del coord, box, fparam, aparam builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() extended_coord, extended_atype, nlist, mapping = builder.build( - cc, atype, bb, rcut, sel + cc, atype, bb, rcut, sel, pair_excl=pair_excl ) extended_coord = extended_coord.reshape(nframes, -1, 3) if coord_corr_for_virial is not None: @@ -377,6 +386,8 @@ def call_common( coord_corr_for_virial=coord_corr_for_virial, charge_spin=cs, neighbor_list=neighbor_list, + # exclusion is a nlist-BUILD transform (decision #18/A4) + pair_excl=getattr(self.atomic_model, "pair_excl", None), ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index fa3863320c..0d9da07355 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -180,6 +180,26 @@ def process_spin_input_lower( mapping_updated = None # extend the nlist nlist_updated = self.extend_nlist(extended_atype, nlist) + # This extension CREATES the virtual-atom nlist entries, so it is the + # BUILD site of the spin-extended nlist: fold the backbone's model-level + # pair exclusion in here (decision #18/A4 — the lower consumes a + # pre-excluded nlist and never re-applies it). No-op when the backbone + # has no pair_exclude_types. + pair_excl = getattr( + self.backbone_model.atomic_model + if hasattr(self.backbone_model, "atomic_model") + else self.backbone_model, + "pair_excl", + None, + ) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist_updated = apply_pair_exclusion_nlist( + nlist_updated, extended_atype_updated, pair_excl + ) return ( extended_coord_updated, extended_atype_updated, diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 8516409c1b..2697e57388 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -52,6 +52,7 @@ def extend_input_and_build_neighbor_list( sel: list[int], mixed_types: bool = False, box: Array | None = None, + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(coord, atype) nframes, nloc = atype.shape[:2] @@ -72,6 +73,7 @@ def extend_input_and_build_neighbor_list( rcut, sel, distinguish_types=(not mixed_types), + pair_excl=pair_excl, ) extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) return extended_coord, extended_atype, mapping, nlist @@ -110,12 +112,14 @@ def apply_pair_exclusion_nlist( Notes ----- - The C++ inference-path mirror is ``applyPairExclusionNlist`` in - ``source/api_cc/include/commonPT.h``. It uses the same argument order - (nlist, atype_ext, ...) and the same variable names (``type_ij``, - ``keep``): it computes ``type_ij`` from the center/neighbor types via - the flat ``(ntypes+1)^2`` table and replaces excluded entries with - ``-1``. + Exclusion is a nlist-BUILD transform (decision #18/A4), same as the graph + route: it is applied ONCE where the nlist is built — the Python builders + (``build_neighbor_list(pair_excl=...)`` / the NeighborList strategies) or + the C++ ingestion seam (``applyPairExclusionNlist`` in + ``source/api_cc/include/commonPT.h``, same table lookup and variable + names). The lower interfaces (``call_lower`` / ``forward_common_atomic`` + and the exported artifacts) consume a pre-excluded nlist and never + re-apply it. """ if pair_excl is None or len(pair_excl.exclude_types) == 0: return nlist diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index dba6d6946b..5b3716f541 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -11,9 +11,17 @@ from collections.abc import ( Callable, ) +from typing import ( + TYPE_CHECKING, +) import tensorflow as tf +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + from deepmd.dpmodel.output_def import ( ModelOutputDef, ) @@ -52,8 +60,15 @@ def model_call_from_call_lower( fparam: tf.Tensor, aparam: tf.Tensor, do_atomic_virial: bool = False, + pair_excl: "PairExcludeMask | None" = None, ) -> dict[str, tf.Tensor]: - """Return model prediction from lower interface.""" + """Return model prediction from lower interface. + + ``pair_excl`` is the model-level pair-type exclusion mask. Exclusion is a + nlist-BUILD transform (decision #18/A4): it is folded into the nlist here, + in the traced TF wrapper, because the lower JAX model consumes a + pre-excluded nlist and never re-applies it. + """ atype_shape = tf.shape(atype) nframes, nloc = atype_shape[0], atype_shape[1] cc, bb, fp, ap = coord, box, fparam, aparam @@ -78,6 +93,31 @@ def model_call_from_call_lower( # need to be distinguished here distinguish_types=False, ) + if pair_excl is not None and len(pair_excl.get_exclude_types()) > 0: + # TF twin of ``apply_pair_exclusion_nlist`` (nlist-BUILD transform, + # decision #18/A4): erase excluded type-pairs to -1 via the flat + # (ntypes+1)^2 keep table, exactly like the numpy / C++ versions. + n1 = pair_excl.ntypes + 1 + table = tf.constant(pair_excl.type_mask, dtype=tf.int32) # ((n1*n1),) + nall = tf.shape(extended_atype)[1] + # map -1 (empty slot) to the virtual atom appended at index nall + ae = tf.concat( + [ + extended_atype, + tf.fill([nframes, 1], tf.constant(pair_excl.ntypes, atype.dtype)), + ], + axis=1, + ) + nlist_for_type = tf.where( + nlist == -1, tf.fill(tf.shape(nlist), tf.cast(nall, nlist.dtype)), nlist + ) + type_j = tf.gather(ae, nlist_for_type, batch_dims=1) # (nf, nloc, nnei) + type_i = ae[:, :nloc] * n1 # (nf, nloc) + type_ij = type_i[:, :, None] + type_j + keep = tf.gather(table, tf.cast(type_ij, tf.int32)) + nlist = tf.where( + keep == 1, nlist, tf.fill(tf.shape(nlist), tf.cast(-1, nlist.dtype)) + ) extended_coord = tf.reshape(extended_coord, [nframes, -1, 3]) model_predict_lower = call_lower( extended_coord, diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index dd496dedf0..6b06ec81e0 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -176,6 +176,9 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + # exclusion is a nlist-BUILD transform (decision #18/A4); + # the traced lower consumes a pre-excluded nlist. + pair_excl=getattr(model.atomic_model, "pair_excl", None), ) return call diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 8c1e85c59c..953624c466 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -84,6 +84,25 @@ def __init__( self.model_def_script = model_def_script self._has_default_fparam = has_default_fparam self.default_fparam = default_fparam + # Model-level pair_exclude_types, rebuilt from the training config so + # the outer wrapper can fold it into the nlist at BUILD time (decision + # #18/A4) — the serialized StableHLO lower consumes a pre-excluded + # nlist and never re-applies it. + import json + + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + pet = [] + if model_def_script: + try: + pet = json.loads(model_def_script).get("pair_exclude_types", []) + except (ValueError, AttributeError): + pet = [] + self._pair_excl = ( + PairExcludeMask(len(type_map), [tuple(p) for p in pet]) if pet else None + ) def __call__( self, @@ -167,6 +186,9 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + # exclusion is a nlist-BUILD transform (decision #18/A4); the + # serialized StableHLO lower consumes a pre-excluded nlist. + pair_excl=self._pair_excl, ) def model_output_def(self) -> ModelOutputDef: diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index c77bc944b5..655bf5c581 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -17,9 +17,15 @@ Path, ) from typing import ( + TYPE_CHECKING, Any, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + import numpy as np import optax import orbax.checkpoint as ocp @@ -773,6 +779,7 @@ def _prepare_batch( box=jax_data["box"] if jax_data["find_box"] else None, fparam=jax_data.get("fparam", None), aparam=jax_data.get("aparam", None), + pair_excl=getattr(model.atomic_model, "pair_excl", None), ) return jax_data, extended_coord, extended_atype, nlist, mapping, fp, ap @@ -1007,6 +1014,7 @@ def prepare_input( box: np.ndarray | None = None, fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[ np.ndarray, np.ndarray, @@ -1038,6 +1046,9 @@ def prepare_input( # types will be distinguished in the lower interface, # so it doesn't need to be distinguished here distinguish_types=False, + # model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4); the lower consumes a pre-excluded nlist. + pair_excl=pair_excl, ) extended_coord = extended_coord.reshape(nframes, -1, 3) return extended_coord, extended_atype, nlist, mapping, fp, ap diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 5b8df237cc..5a84d0e881 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -957,6 +957,10 @@ def _build_nlist_native( sel = self._sel mixed_types = self._mixed_types + # Model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): fold it in here; the exported dense lower consumes a + # pre-excluded nlist and never re-applies it. + pair_excl = self._model_pair_excl() if self._nlist_builder is not None: # O(N) cell-list strategy (e.g. vesin): builds the same extended # representation. Match the native builder's type handling @@ -966,7 +970,7 @@ def _build_nlist_native( # type-distinguished nlist a non-mixed-type descriptor expects. The # main eval path is unaffected (its ``format_nlist`` re-formats). extended_coord, extended_atype, nlist, mapping = self._nlist_builder.build( - coords, atom_types, cells, rcut, sel + coords, atom_types, cells, rcut, sel, pair_excl=pair_excl ) if not mixed_types: nlist = nlist_distinguish_types(nlist, extended_atype, sel) @@ -991,6 +995,7 @@ def _build_nlist_native( rcut, sel, distinguish_types=not mixed_types, + pair_excl=pair_excl, ) extended_coord = extended_coord.reshape(nframes, -1, 3) return extended_coord, extended_atype, nlist, mapping @@ -1050,10 +1055,21 @@ def _build_nlist_ase( ext_atypes.append(ea) nlists.append(nl) mappings.append(mp) + extended_atype = np.stack(ext_atypes, axis=0) + nlist = np.stack(nlists, axis=0) + # Model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): fold it in here, like the native builder path. + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist( + nlist, extended_atype, self._model_pair_excl() + ) return ( np.stack(ext_coords, axis=0), - np.stack(ext_atypes, axis=0), - np.stack(nlists, axis=0), + extended_atype, + nlist, np.stack(mappings, axis=0), ) @@ -1779,7 +1795,7 @@ def _build_eval_graph( # (decision #18): apply it here so the exported ``.pt2`` lower consumes a # pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++ # ``applyPairExclusion`` and the eager dpmodel/pt_expt build path). - pair_excl = self._graph_pair_excl() + pair_excl = self._model_pair_excl() if method == "dense": from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, @@ -1824,7 +1840,7 @@ def _build_eval_graph( "use 'dense', 'ase', 'vesin', or 'nv'" ) - def _graph_pair_excl(self) -> "PairExcludeMask | None": + def _model_pair_excl(self) -> "PairExcludeMask | None": """Model-level ``pair_exclude_types`` as a ``PairExcludeMask`` (or None). Applied at graph BUILD time (decision #18), NOT inside the exported diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9edd3e4d4b..5c07d438f7 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -941,6 +941,9 @@ def forward( rcut, sel, distinguish_types=False, + # model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4); the compiled dense lower consumes a pre-excluded nlist. + pair_excl=getattr(self.original_model.atomic_model, "pair_excl", None), ) ext_coord = ext_coord.reshape(nframes, -1, 3) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index ef7a1d69b6..c6d26dcde8 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -740,17 +740,15 @@ def _probe_has_message_passing(obj: object) -> bool | None: meta["lower_input_kind"] = "graph" if lower_kind == "graph" else "nlist" # Model-level pair-type exclusion (``pair_exclude_types``): a list of - # ``[ti, tj]`` type pairs whose interaction is dropped. The compiled AOTI - # forward ALREADY applies this exclusion internally (the graph seam - # ``apply_pair_exclusion`` / dense seam ``apply_pair_exclusion_nlist`` is - # traced into the exported artifact), so this field is redundant for the - # compiled math. It is serialized so that the C++ loaders - # (``DeepPotPTExpt::init``) can rebuild the flat ``(ntypes+1)^2`` keep table - # and re-apply the SAME transform (``applyPairExclusion`` / - # ``applyPairExclusionNlist`` in ``commonPT.h``) at the ingestion seam as an - # idempotent backstop, keeping the C++ ingestion path side-by-side - # reviewable with the Python transforms. Descriptor-level ``exclude_types`` - # needs NO metadata: it is fully inside the compiled graph. + # ``[ti, tj]`` type pairs whose interaction is dropped. Exclusion is a + # BUILD-time transform on BOTH routes (decision #18/A4): the exported + # lower (graph edge_mask / dense nlist) consumes a pre-excluded input and + # never re-applies it, so every feeder — Python builders, DeepEval, C++ + # ``applyPairExclusion`` / ``applyPairExclusionNlist`` — MUST fold the + # exclusion in at build. This metadata field is what lets external + # feeders (C++ ``DeepPotPTExpt::init``, metadata-only DeepEval) rebuild + # the mask. Descriptor-level ``exclude_types`` needs NO metadata: it is + # fully inside the compiled artifact. pair_exclude_types: list[list[int]] = [] for obj in ( getattr(model, "atomic_model", None), diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 6114dca340..4f14e0178f 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -334,9 +334,10 @@ class DeepPotPTExpt : public DeepPotBackend { // Flat (ntypes+1)^2 model-level pair-type keep table, rebuilt in ``init`` // from the ``pair_exclude_types`` metadata field (see // ``deepmd::buildPairExcludeTable``). Empty => no model-level exclusion. - // Applied at the C++ ingestion seam (``applyPairExclusion`` graph / - // ``applyPairExclusionNlist`` dense) as an idempotent backstop; the compiled - // .pt2 graph already applies the same transform internally. + // Exclusion is a BUILD-time transform (decision #18/A4): the C++ ingestion + // seam is the single application site (``applyPairExclusion`` graph / + // ``applyPairExclusionNlist`` dense); the exported .pt2 lowers consume + // pre-excluded inputs and never re-apply it. std::vector pair_exclude_table_; std::unique_ptr with_comm_tempfile_; std::unique_ptr with_comm_loader; diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index af0f37df20..48cf984cf1 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -515,9 +515,10 @@ inline std::vector buildPairExcludeTable( * same argument order (edge_index, edge_mask, atype, ...) and same variable * names (``type_ij``, ``keep``). * - * The compiled ``.pt2`` graph already applies this exclusion internally (the - * seam transform is traced into the exported forward), so this call is an - * idempotent backstop at the C++ ingestion seam. + * OWNERSHIP: exclusion is a BUILD-time transform (decision #18/A4). The + * exported ``.pt2`` graph lower consumes a pre-excluded ``edge_mask`` and + * never re-applies it, so this call is the SINGLE application site on the + * C++ graph route. * * @param edge_index (2, E) int64 ``[src, dst]``; src = neighbor, dst = center. * @param edge_mask (E,) bool real-vs-padding mask to be ANDed in place. @@ -558,9 +559,12 @@ inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, * Inference-path twin of Python ``apply_pair_exclusion_nlist`` in * ``deepmd/dpmodel/utils/nlist.py`` + * ``PairExcludeMask.build_type_exclude_mask``. Same argument order (nlist, - * atype_ext, ...) and same variable names - * (``type_ij``, ``keep``). Idempotent: erasing ``-1`` a second time is a - * no-op. + * atype_ext, ...) and same variable names (``type_ij``, ``keep``). + * + * OWNERSHIP: exclusion is a BUILD-time transform (decision #18/A4). The + * exported ``.pt2`` dense lower consumes a pre-excluded nlist and never + * re-applies it, so this call is the SINGLE application site on the C++ + * dense route. * * @param nlist (nf, nloc, nnei) int64 neighbour list; ``-1`` == empty slot. * @param atype_ext (nf, nall) int64 extended atom types. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 1185dfba80..9dd544bbad 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -201,10 +201,12 @@ void DeepPotPTExpt::init(const std::string& model, has_message_passing_ = metadata.obj_val.count("has_message_passing") && metadata["has_message_passing"].as_bool(); - // Model-level pair-type exclusion table. ``pair_exclude_types`` is a list of - // [ti, tj] pairs; rebuild the flat (ntypes+1)^2 keep table exactly like the - // Python ``PairExcludeMask`` ctor so the ingestion seam can re-apply the same - // exclusion (idempotent backstop; the compiled graph already applies it). + // Model-level pair-type exclusion table. ``pair_exclude_types`` is a list + // of [ti, tj] pairs; rebuild the flat (ntypes+1)^2 keep table exactly like + // the Python ``PairExcludeMask`` ctor. Exclusion is a BUILD-time transform + // (decision #18/A4): the C++ ingestion seam is the single application site + // (applyPairExclusion graph / applyPairExclusionNlist dense); the exported + // lowers consume pre-excluded inputs and never re-apply it. { std::vector> pair_exclude_types; if (metadata.obj_val.count("pair_exclude_types")) { @@ -890,8 +892,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, torch::full({1}, n_node_count, int_option).to(device); at::Tensor node_atype = atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); - // Model-level pair exclusion at the ingestion seam (idempotent backstop; - // the compiled graph already applies the same transform internally). + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported graph lower consumes a pre-excluded edge_mask + // and never re-applies it; this is the single application site on the + // C++ graph route. const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( edge_tensors.edge_index, edge_tensors.edge_mask, node_atype, pair_exclude_table_, ntypes); @@ -900,8 +904,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_vec, graph_edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); } else { - // Model-level pair exclusion at the dense ingestion seam (idempotent - // backstop; the compiled dense forward already applies the same erase). + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported dense lower consumes a pre-excluded nlist and + // never re-applies it; this is the single application site on the C++ + // dense route. const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = @@ -1262,8 +1268,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_index_ext, edge_tensors.edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); } else if (lower_input_is_graph_) { - // Model-level pair exclusion at the ingestion seam (idempotent backstop; - // the compiled graph already applies the same transform internally). + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported graph lower consumes a pre-excluded edge_mask and + // never re-applies it; this is the single application site on the C++ + // graph route. const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( graph_tensors.edge_index, graph_tensors.edge_mask, graph_tensors.atype, pair_exclude_table_, ntypes); @@ -1272,8 +1280,10 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, graph_tensors.edge_vec, graph_edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); } else { - // Model-level pair exclusion at the dense ingestion seam (idempotent - // backstop; the compiled dense forward already applies the same erase). + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported dense lower consumes a pre-excluded nlist and + // never re-applies it; this is the single application site on the C++ + // dense route. const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( nlist_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = diff --git a/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc index a8f7e7daee..0b2934ecf9 100644 --- a/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc @@ -8,13 +8,15 @@ // A no-exclusion baseline (deeppot_dpa1_pairexcl_none.pt2, empty exclude table) // exercises the identity/pre-change branch of both helpers. // -// The compiled .pt2 forward ALREADY applies the exclusion internally (the seam -// transform is traced into the exported artifact), so the C++ helpers are an -// idempotent backstop; the reference values (.expected sidecars) come from the -// Python DeepEval of the SAME .pt2, so a 1e-10 match validates the whole chain -// (pair_exclude_types metadata round-trip + init table build + seam apply + -// compiled math). A separate assertion (excluded energy != baseline energy) -// proves the exclusion is genuinely active and not silently dropped. +// Exclusion is a BUILD-time transform (decision #18/A4): the exported .pt2 +// lowers consume pre-excluded inputs (graph edge_mask / dense nlist) and never +// re-apply it, so the C++ helpers here are the SINGLE application site — these +// tests are what proves they are load-bearing. The reference values +// (.expected sidecars) come from the Python DeepEval of the SAME .pt2, so a +// 1e-10 match validates the whole chain (pair_exclude_types metadata +// round-trip + init table build + build-time apply + compiled math). A +// separate assertion (excluded energy != baseline energy) proves the exclusion +// is genuinely active and not silently dropped. #include #include diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index ae5a9c610d..45b5fa7efa 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -10,6 +10,9 @@ from deepmd.dpmodel.descriptor import ( DescrptSeA, ) +from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, +) from deepmd.dpmodel.fitting import ( InvarFitting, ) @@ -111,10 +114,16 @@ def test_excl_consistency(self) -> None: md1.descriptor.reinit_exclude(pair_excl) md1.fitting_net.reinit_exclude(atom_excl) - # check energy consistency - args = [self.coord_ext, self.atype_ext, self.nlist] - ret0 = md0.forward_common_atomic(*args) - ret1 = md1.forward_common_atomic(*args) + # check energy consistency. Model-level pair exclusion is a + # nlist-BUILD transform (decision #18/A4): forward_common_atomic + # consumes a pre-excluded nlist, so fold md0's exclusion into its + # nlist here (the descriptor-level md1 keeps the raw nlist; its + # emask applies inside the descriptor). + nlist0 = apply_pair_exclusion_nlist( + self.nlist, self.atype_ext, md0.pair_excl + ) + ret0 = md0.forward_common_atomic(self.coord_ext, self.atype_ext, nlist0) + ret1 = md1.forward_common_atomic(self.coord_ext, self.atype_ext, self.nlist) np.testing.assert_allclose( ret0["energy"], ret1["energy"], diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index 545f639f80..fa5e09e163 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -226,6 +226,65 @@ def test_model_pair_exclude_applied_at_build_not_in_lower(): ), "build-time pair exclusion must change the graph energy" +def test_model_pair_exclude_applied_at_build_not_in_dense_lower(): + """Dense-route seam contract (decision #18/A4, mirror of the graph test): + model-level pair_exclude is a nlist-BUILD transform. The dense lower + (``call_lower``) must NOT re-apply it — it consumes whatever nlist the + builder produced. Feeding a RAW nlist to the lower on a model that HAS + ``pair_exclude_types`` yields the SAME result as a model with no exclusion; + the exclusion only takes effect when folded in at build. + """ + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + extend_input_and_build_neighbor_list, + ) + + rng = np.random.default_rng(4) + nloc = 6 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0, 1]], dtype=np.int64) + box = np.eye(3).reshape(1, 9) * 20.0 + ds = DescrptDPA1(rcut=4.0, rcut_smth=0.5, sel=[200], ntypes=2, attn_layer=0) + ft = InvarFitting("energy", 2, ds.get_dim_out(), 1, mixed_types=True) + model = EnergyModel(ds, ft, type_map=["a", "b"], pair_exclude_types=[(0, 1)]) + assert model.atomic_model.pair_excl is not None + + # RAW quartet: built WITHOUT pair_excl (no exclusion folded into the nlist). + coord_ext, atype_ext, mapping, nlist = extend_input_and_build_neighbor_list( + coord.reshape(1, -1), + atype, + model.get_rcut(), + model.get_sel(), + mixed_types=True, + box=box, + ) + # EnergyModel.call_lower translates keys: reduced energy -> "energy" + out_with_excl_model = model.call_lower(coord_ext, atype_ext, nlist, mapping) + # clear the model-level exclusion; the lower output must be UNCHANGED, + # proving the lower never consulted ``pair_excl``. + model.atomic_model.reinit_pair_exclude([]) + assert model.atomic_model.pair_excl is None + out_no_excl_model = model.call_lower(coord_ext, atype_ext, nlist, mapping) + np.testing.assert_allclose( + np.asarray(out_with_excl_model["energy"]), + np.asarray(out_no_excl_model["energy"]), + rtol=1e-12, + atol=1e-12, + ) + + # Positive control: folding the exclusion in at BUILD changes the output. + nlist_excl = apply_pair_exclusion_nlist( + nlist, atype_ext, PairExcludeMask(2, [(0, 1)]) + ) + out_built_excl = model.call_lower(coord_ext, atype_ext, nlist_excl, mapping) + assert not np.allclose( + np.asarray(out_built_excl["energy"]), + np.asarray(out_no_excl_model["energy"]), + rtol=1e-9, + atol=1e-9, + ), "build-time pair exclusion must change the dense energy" + + def test_graph_matches_dense_with_fparam(): """Frame parameter is gathered to nodes by frame_id in forward_atomic_graph and fed to the fitting's call_graph; the graph path must match dense at 1e-12 diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index 800eaa072d..85466c1919 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -11,6 +11,9 @@ ) from deepmd.dpmodel.model.ener_model import EnergyModel as EnergyModelDP from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -484,6 +487,14 @@ def setUp(self) -> None: 6.0, [20, 20], distinguish_types=True, + # model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): the dpmodel-family lowers consume a pre-excluded nlist; + # legacy pt/pd re-apply internally, which is an idempotent no-op. + pair_excl=PairExcludeMask( + self.ntypes, [tuple(p) for p in self.data["pair_exclude_types"]] + ) + if self.data["pair_exclude_types"] + else None, ) extended_coord = extended_coord.reshape(nframes, -1, 3) self.nlist = nlist diff --git a/source/tests/infer/gen_dpa1_pairexcl.py b/source/tests/infer/gen_dpa1_pairexcl.py index c37b3d23f1..8c2291db4d 100644 --- a/source/tests/infer/gen_dpa1_pairexcl.py +++ b/source/tests/infer/gen_dpa1_pairexcl.py @@ -9,15 +9,19 @@ - deeppot_dpa1_pairexcl_nlist.pt2 (lower_kind="nlist", pair_exclude=[[0,1]]) - deeppot_dpa1_pairexcl_none.pt2 (lower_kind="graph", NO exclusion) -The pair models exercise the C++ pair-exclusion ingestion seam: -``applyPairExclusion`` (graph route) / ``applyPairExclusionNlist`` (dense route) -plus the ``pair_exclude_types`` metadata round-trip in ``DeepPotPTExpt::init``. -Model-level pair exclusion is a graph-BUILD transform (decision #18): it is -folded into ``edge_mask`` at build time (``applyPairExclusion`` in C++, the -NeighborGraph builder in Python ``DeepEval``), NOT inside the exported ``.pt2`` -lower. The gtest validates C++ energy/force vs the Python ``DeepEval`` reference -at 1e-10 and, by comparing against the ``_none`` baseline, confirms the exclusion -is actually active. +The pair models exercise the C++ pair-exclusion ownership (decision #18/A4: +exclusion is a BUILD-time transform on BOTH routes) plus the +``pair_exclude_types`` metadata round-trip in ``DeepPotPTExpt::init``: + +- graph route: folded into ``edge_mask`` at build by ``applyPairExclusion`` + (C++) / the NeighborGraph builder (Python DeepEval). +- dense route: folded into the nlist at build by ``applyPairExclusionNlist`` + (C++) / ``build_neighbor_list(pair_excl=...)`` (Python DeepEval). + +The exported lowers consume pre-excluded inputs and never re-apply the +exclusion. The gtest validates C++ energy/force vs the Python ``DeepEval`` +reference at 1e-10 and, by comparing against the ``_none`` baseline, confirms +the exclusion is actually active. Reference sidecar files (.expected) consumed by the C++ gtest are written from the Python ``DeepEval`` evaluation of each pair model (PBC + NoPBC sections). diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index 4160993885..d85e6dae27 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -26,6 +26,7 @@ from_dense_quartet, ) from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, build_neighbor_list, extend_coord_with_ghosts, ) @@ -211,6 +212,13 @@ def test_force_virial_parity_vs_legacy( nf = ext_coord.shape[0] nloc = self.natoms + # Model-level pair_exclude is a nlist-BUILD transform (decision + # #18/A4): BOTH lowers consume pre-excluded inputs, so fold the + # exclusion into the dense nlist here (mirrors the C++ + # ``applyPairExclusionNlist`` build step). + nlist = apply_pair_exclusion_nlist( + nlist, ext_atype, model.atomic_model.pair_excl + ) legacy = model.forward_common_lower( ext_coord.clone().requires_grad_(True), ext_atype, From 83dd109fdf7f87fe5c47db1cb5c870900fda821d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 13:32:26 +0800 Subject: [PATCH 32/63] refactor(jax2tf): reuse dpmodel apply_pair_exclusion_nlist via ndtensorflow The jax2tf SavedModel wrapper previously carried a hand-written TensorFlow twin of the model-level pair-exclusion nlist transform (decision #18/A4), pinned to the canonical numpy/C++ implementation only by a value test. Replace it with a call to the canonical dpmodel apply_pair_exclusion_nlist through the vendored ndtensorflow array-API namespace -- the same mechanism the TF2 backend uses to run dpmodel code on TensorFlow. Unlike the neighbor-list *build* (which has data-dependent Python control flow and is deliberately kept as a TF twin, see jax2tf/nlist.py), the exclusion's only branch is on the static exclude_types config, so it traces cleanly under tf.saved_model.save. This removes the second implementation: the exclusion transform now has a single owner (dpmodel) reused across dpmodel / pt_expt / native-jax / TF2 / jax2tf, with the C++ ingestion seam the only remaining twin (unavoidable). Add source/tests/consistent/io/test_pair_exclude_savedmodel.py: exports a pair-excluded se_e2_a model to .savedmodel and checks energy vs the dpmodel reference (fp64), force/virial vs pytorch (the numpy dpmodel DeepEval does not compute usable forces), the identity/no-exclusion branch, and that the exclusion is genuinely active. --- deepmd/jax/jax2tf/make_model.py | 38 ++-- .../io/test_pair_exclude_savedmodel.py | 172 ++++++++++++++++++ 2 files changed, 188 insertions(+), 22 deletions(-) create mode 100644 source/tests/consistent/io/test_pair_exclude_savedmodel.py diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index 5b3716f541..86923f7e99 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -94,30 +94,24 @@ def model_call_from_call_lower( distinguish_types=False, ) if pair_excl is not None and len(pair_excl.get_exclude_types()) > 0: - # TF twin of ``apply_pair_exclusion_nlist`` (nlist-BUILD transform, - # decision #18/A4): erase excluded type-pairs to -1 via the flat - # (ntypes+1)^2 keep table, exactly like the numpy / C++ versions. - n1 = pair_excl.ntypes + 1 - table = tf.constant(pair_excl.type_mask, dtype=tf.int32) # ((n1*n1),) - nall = tf.shape(extended_atype)[1] - # map -1 (empty slot) to the virtual atom appended at index nall - ae = tf.concat( - [ - extended_atype, - tf.fill([nframes, 1], tf.constant(pair_excl.ntypes, atype.dtype)), - ], - axis=1, + # Reuse the canonical dpmodel nlist-BUILD transform (decision #18/A4) + # via the vendored ``ndtensorflow`` array-API namespace -- the same way + # the TF2 backend (``deepmd/tf2``) runs dpmodel array-API code on + # TensorFlow. Unlike the neighbor-list *build* (see the docstring of + # ``jax2tf/nlist.py``), the exclusion has no data-dependent Python + # control flow: its only branch is on the static ``exclude_types`` + # config, so it traces cleanly under SavedModel export and does not + # need a hand-written TF twin. + from deepmd._vendors import ( + ndtensorflow as ndtf, ) - nlist_for_type = tf.where( - nlist == -1, tf.fill(tf.shape(nlist), tf.cast(nall, nlist.dtype)), nlist - ) - type_j = tf.gather(ae, nlist_for_type, batch_dims=1) # (nf, nloc, nnei) - type_i = ae[:, :nloc] * n1 # (nf, nloc) - type_ij = type_i[:, :, None] + type_j - keep = tf.gather(table, tf.cast(type_ij, tf.int32)) - nlist = tf.where( - keep == 1, nlist, tf.fill(tf.shape(nlist), tf.cast(-1, nlist.dtype)) + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, ) + + nlist = apply_pair_exclusion_nlist( + ndtf.asarray(nlist), ndtf.asarray(extended_atype), pair_excl + ).unwrap() extended_coord = tf.reshape(extended_coord, [nframes, -1, 3]) model_predict_lower = call_lower( extended_coord, diff --git a/source/tests/consistent/io/test_pair_exclude_savedmodel.py b/source/tests/consistent/io/test_pair_exclude_savedmodel.py new file mode 100644 index 0000000000..64186fc327 --- /dev/null +++ b/source/tests/consistent/io/test_pair_exclude_savedmodel.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Validate that the jax2tf ``.savedmodel`` export reuses the canonical dpmodel +model-level pair-exclusion transform. + +Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision +#18/A4). The jax2tf SavedModel wrapper folds it into the neighbor list by +calling :func:`deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist` through +the vendored ``ndtensorflow`` array-API namespace, rather than a hand-written +TensorFlow twin. These tests prove that reuse: + +* the exported SavedModel traces (the reused array-API code is convertible + under ``tf.saved_model.save``), and +* it matches the dpmodel reference to fp64 tolerance (same math, applied at the + same nlist-BUILD seam), +* while a no-exclusion baseline exercises the identity branch and proves the + exclusion is genuinely active (excluded energy differs from baseline). +""" + +import copy +import shutil +import unittest +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.backend.backend import ( + Backend, +) +from deepmd.dpmodel.model.model import ( + get_model, +) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer.deep_eval import ( + DeepEval, +) + +_JAX_BACKEND = Backend.get_backend("jax")() + + +def _model_def_script(pair_exclude_types: list[list[int]]) -> dict: + md = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "ener", + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "atom_ener": [], + "seed": 1, + }, + } + if pair_exclude_types: + md["pair_exclude_types"] = pair_exclude_types + return md + + +@unittest.skipIf( + not _JAX_BACKEND.is_available(), + "jax2tf SavedModel export requires the jax backend (jax + tensorflow).", +) +class TestPairExcludeSavedModel(unittest.TestCase): + """jax2tf ``.savedmodel`` export of model-level pair exclusion.""" + + def setUp(self) -> None: + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + self._artifacts: list[Path] = [] + + def tearDown(self) -> None: + for path in self._artifacts: + if path.is_file(): + path.unlink() + elif path.is_dir(): + shutil.rmtree(path) + + def _eval(self, backend_name: str, md: dict, atomic: bool = False) -> tuple: + """Serialize *md*, write it through *backend_name*, DeepEval it.""" + backend = Backend.get_backend(backend_name)() + suffix = ".savedmodel" if backend_name == "jax" else backend.suffixes[0] + prefix = f"test_pairexcl_savedmodel_{backend_name}_{len(self._artifacts)}" + model_file = prefix + suffix + model = get_model(copy.deepcopy(md)) + data = { + "model": model.serialize(), + "backend": "test", + "model_def_script": md, + } + backend.deserialize_hook(model_file, data) + self._artifacts.append(Path(model_file)) + deep_eval = DeepEval(model_file) + return deep_eval.eval(self.coords, self.box, self.atype, atomic=atomic) + + def test_savedmodel_matches_reference_excluded(self) -> None: + """SavedModel export matches the references when [0, 1] is excluded. + + Energy is compared against the dpmodel reference: both apply the + exclusion at the nlist-BUILD seam, so the neighbor list fed to the + lower model is identical and the energy must match to fp64. Force and + virial are compared against the pytorch backend instead, because the + numpy dpmodel DeepEval does not compute usable forces (returns NaN) + regardless of exclusion; pytorch applies the same exclusion + (idempotently) and produces finite gradients. + """ + md = _model_def_script([[0, 1]]) + e_ref = self._eval("dpmodel", md)[0] + e_sm, f_sm, v_sm = self._eval("jax", md)[:3] + np.testing.assert_allclose(e_sm, e_ref, atol=1e-10) + if Backend.get_backend("pytorch")().is_available(): + _, f_pt, v_pt = self._eval("pytorch", md)[:3] + np.testing.assert_allclose(f_sm, f_pt, atol=1e-10) + np.testing.assert_allclose(v_sm, v_pt, atol=1e-10) + + def test_savedmodel_matches_dpmodel_no_exclusion(self) -> None: + """Identity branch: no exclusion still round-trips and matches dpmodel.""" + md = _model_def_script([]) + e_ref = self._eval("dpmodel", md)[0] + e_sm = self._eval("jax", md)[0] + np.testing.assert_allclose(e_sm, e_ref, atol=1e-10) + + def test_exclusion_is_active(self) -> None: + """Excluding [0, 1] changes the SavedModel energy vs the baseline.""" + e_excl = self._eval("jax", _model_def_script([[0, 1]]))[0] + e_none = self._eval("jax", _model_def_script([]))[0] + # O-H pairs dominate this water-like cluster; excluding them moves the + # energy well above the fp64 tolerance. + self.assertGreater(float(np.abs(e_excl - e_none).max()), 1e-6) + + +if __name__ == "__main__": + unittest.main() From 01254151d7786668b5c61a5a6d64cc054de90e3c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 17:42:37 +0800 Subject: [PATCH 33/63] fix(tf2): apply model-level pair_exclude at nlist BUILD; align io test Replace the bespoke test_pair_exclude_savedmodel.py with an aligned IOTest subclass (TestDeepPotPairExclude) in test_io.py: it supplies only the model dict and inherits test_data_equal / test_deep_eval, which export through every backend and cross-compare at rtol/atol 1e-12 (with the built-in all-NaN skip covering the numpy dpmodel force path). Gated on DP_TEST_TF2_ONLY, the mode in which test_deep_eval exercises the jax2tf '.savedmodel' path -- the TF v1 backend raises NotImplementedError on pair_exclude_types, so it must not run there. That aligned cross-backend eval exposed a real gap: the tf2 backend ('.savedmodeltf') silently ignored model-level pair_exclude_types (its SavedModel returned the non-excluded energy 2.847 vs the correct 2.843). tf2 was outside the original PR scope (dpmodel/pt_expt/jax) and its outer TF wrapper (deepmd/tf2/make_model.py) never folded the exclusion into the nlist. Fix it the same way as every other backend: thread pair_excl into model_call_from_call_lower and apply the canonical dpmodel apply_pair_exclusion_nlist at the nlist-BUILD seam (the tensors are already ndtensorflow arrays, so no wrap/unwrap). All five backends now agree. --- deepmd/tf2/make_model.py | 15 ++ deepmd/tf2/utils/serialization.py | 3 + source/tests/consistent/io/test_io.py | 56 ++++++ .../io/test_pair_exclude_savedmodel.py | 172 ------------------ 4 files changed, 74 insertions(+), 172 deletions(-) delete mode 100644 source/tests/consistent/io/test_pair_exclude_savedmodel.py diff --git a/deepmd/tf2/make_model.py b/deepmd/tf2/make_model.py index 458e0b5a06..e08890247e 100644 --- a/deepmd/tf2/make_model.py +++ b/deepmd/tf2/make_model.py @@ -3,6 +3,7 @@ Callable, ) from typing import ( + TYPE_CHECKING, Any, ) @@ -18,6 +19,7 @@ NeighborList, ) from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, nlist_distinguish_types, ) from deepmd.tf2.common import ( @@ -38,6 +40,11 @@ normalize_coord, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + def _unwrap_tuple(values: tuple[Array, ...]) -> tuple[tf.Tensor, ...]: return tuple(to_tf_tensor(value) for value in values) @@ -68,6 +75,7 @@ def model_call_from_call_lower( charge_spin: Array | None = None, neighbor_list: NeighborList | None = None, pass_lower_kwargs: bool = False, + pair_excl: "PairExcludeMask | None" = None, ) -> dict[str, Array]: """Return model prediction from lower interface. @@ -124,6 +132,13 @@ def model_call_from_call_lower( charge_spin=charge_spin, neighbor_list=neighbor_list, ) + if pair_excl is not None: + # Model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): fold it into the freshly built nlist here (already an + # ndtensorflow array, so no wrap/unwrap) via the canonical dpmodel + # helper, before the lower consumes it. Identity when nothing is + # excluded. + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) lower_kwargs: dict[str, Any] = {"fparam": fp, "aparam": ap} if pass_lower_kwargs: if nlist_is_formatted: diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index 5de81a2bdf..71d9dd699a 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -390,6 +390,9 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + # exclusion is a nlist-BUILD transform (decision #18/A4); + # the traced lower consumes a pre-excluded nlist. + pair_excl=getattr(model.atomic_model, "pair_excl", None), ) ) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 9601ac7607..be1808c0a9 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -311,3 +311,59 @@ def setUp(self) -> None: def tearDown(self) -> None: IOTest.tearDown(self) + + +@unittest.skipIf( + not DP_TEST_TF2_ONLY, + "pair_exclude_types is not supported by the TF v1 backend; it is validated " + "in the TF2-only job, where test_deep_eval also exercises the jax2tf " + "'.savedmodel' export path (see the backend table in test_deep_eval).", +) +class TestDeepPotPairExclude(unittest.TestCase, IOTest): + """Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision + #18/A4). Every backend folds it in where the neighbor list is built (the + jax2tf ``.savedmodel`` export reuses the dpmodel + ``apply_pair_exclusion_nlist`` via the ``ndtensorflow`` namespace), so the + exported models must still eval-agree across backends. + """ + + def setUp(self) -> None: + model_def_script = { + "type_map": ["O", "H"], + "pair_exclude_types": [[0, 1]], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 3, + 6, + ], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "ener", + "neuron": [ + 5, + 5, + ], + "resnet_dt": True, + "precision": "float64", + "atom_ener": [], + "seed": 1, + }, + } + model = get_model(copy.deepcopy(model_def_script)) + self.data = { + "model": model.serialize(), + "backend": "test", + "model_def_script": model_def_script, + } + + def tearDown(self) -> None: + IOTest.tearDown(self) diff --git a/source/tests/consistent/io/test_pair_exclude_savedmodel.py b/source/tests/consistent/io/test_pair_exclude_savedmodel.py deleted file mode 100644 index 64186fc327..0000000000 --- a/source/tests/consistent/io/test_pair_exclude_savedmodel.py +++ /dev/null @@ -1,172 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Validate that the jax2tf ``.savedmodel`` export reuses the canonical dpmodel -model-level pair-exclusion transform. - -Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision -#18/A4). The jax2tf SavedModel wrapper folds it into the neighbor list by -calling :func:`deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist` through -the vendored ``ndtensorflow`` array-API namespace, rather than a hand-written -TensorFlow twin. These tests prove that reuse: - -* the exported SavedModel traces (the reused array-API code is convertible - under ``tf.saved_model.save``), and -* it matches the dpmodel reference to fp64 tolerance (same math, applied at the - same nlist-BUILD seam), -* while a no-exclusion baseline exercises the identity branch and proves the - exclusion is genuinely active (excluded energy differs from baseline). -""" - -import copy -import shutil -import unittest -from pathlib import ( - Path, -) - -import numpy as np - -from deepmd.backend.backend import ( - Backend, -) -from deepmd.dpmodel.model.model import ( - get_model, -) -from deepmd.env import ( - GLOBAL_NP_FLOAT_PRECISION, -) -from deepmd.infer.deep_eval import ( - DeepEval, -) - -_JAX_BACKEND = Backend.get_backend("jax")() - - -def _model_def_script(pair_exclude_types: list[list[int]]) -> dict: - md = { - "type_map": ["O", "H"], - "descriptor": { - "type": "se_e2_a", - "sel": [20, 20], - "rcut_smth": 0.50, - "rcut": 6.00, - "neuron": [3, 6], - "resnet_dt": False, - "axis_neuron": 2, - "precision": "float64", - "type_one_side": True, - "seed": 1, - }, - "fitting_net": { - "type": "ener", - "neuron": [5, 5], - "resnet_dt": True, - "precision": "float64", - "atom_ener": [], - "seed": 1, - }, - } - if pair_exclude_types: - md["pair_exclude_types"] = pair_exclude_types - return md - - -@unittest.skipIf( - not _JAX_BACKEND.is_available(), - "jax2tf SavedModel export requires the jax backend (jax + tensorflow).", -) -class TestPairExcludeSavedModel(unittest.TestCase): - """jax2tf ``.savedmodel`` export of model-level pair exclusion.""" - - def setUp(self) -> None: - self.coords = np.array( - [ - 12.83, - 2.56, - 2.18, - 12.09, - 2.87, - 2.74, - 00.25, - 3.32, - 1.68, - 3.36, - 3.00, - 1.81, - 3.51, - 2.51, - 2.60, - 4.27, - 3.22, - 1.56, - ], - dtype=GLOBAL_NP_FLOAT_PRECISION, - ).reshape(1, -1, 3) - self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) - self.box = np.array( - [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], - dtype=GLOBAL_NP_FLOAT_PRECISION, - ).reshape(1, 9) - self._artifacts: list[Path] = [] - - def tearDown(self) -> None: - for path in self._artifacts: - if path.is_file(): - path.unlink() - elif path.is_dir(): - shutil.rmtree(path) - - def _eval(self, backend_name: str, md: dict, atomic: bool = False) -> tuple: - """Serialize *md*, write it through *backend_name*, DeepEval it.""" - backend = Backend.get_backend(backend_name)() - suffix = ".savedmodel" if backend_name == "jax" else backend.suffixes[0] - prefix = f"test_pairexcl_savedmodel_{backend_name}_{len(self._artifacts)}" - model_file = prefix + suffix - model = get_model(copy.deepcopy(md)) - data = { - "model": model.serialize(), - "backend": "test", - "model_def_script": md, - } - backend.deserialize_hook(model_file, data) - self._artifacts.append(Path(model_file)) - deep_eval = DeepEval(model_file) - return deep_eval.eval(self.coords, self.box, self.atype, atomic=atomic) - - def test_savedmodel_matches_reference_excluded(self) -> None: - """SavedModel export matches the references when [0, 1] is excluded. - - Energy is compared against the dpmodel reference: both apply the - exclusion at the nlist-BUILD seam, so the neighbor list fed to the - lower model is identical and the energy must match to fp64. Force and - virial are compared against the pytorch backend instead, because the - numpy dpmodel DeepEval does not compute usable forces (returns NaN) - regardless of exclusion; pytorch applies the same exclusion - (idempotently) and produces finite gradients. - """ - md = _model_def_script([[0, 1]]) - e_ref = self._eval("dpmodel", md)[0] - e_sm, f_sm, v_sm = self._eval("jax", md)[:3] - np.testing.assert_allclose(e_sm, e_ref, atol=1e-10) - if Backend.get_backend("pytorch")().is_available(): - _, f_pt, v_pt = self._eval("pytorch", md)[:3] - np.testing.assert_allclose(f_sm, f_pt, atol=1e-10) - np.testing.assert_allclose(v_sm, v_pt, atol=1e-10) - - def test_savedmodel_matches_dpmodel_no_exclusion(self) -> None: - """Identity branch: no exclusion still round-trips and matches dpmodel.""" - md = _model_def_script([]) - e_ref = self._eval("dpmodel", md)[0] - e_sm = self._eval("jax", md)[0] - np.testing.assert_allclose(e_sm, e_ref, atol=1e-10) - - def test_exclusion_is_active(self) -> None: - """Excluding [0, 1] changes the SavedModel energy vs the baseline.""" - e_excl = self._eval("jax", _model_def_script([[0, 1]]))[0] - e_none = self._eval("jax", _model_def_script([]))[0] - # O-H pairs dominate this water-like cluster; excluding them moves the - # energy well above the fp64 tolerance. - self.assertGreater(float(np.abs(e_excl - e_none).max()), 1e-6) - - -if __name__ == "__main__": - unittest.main() From 85e1001e9e73ba9cca2d518170fcd527a1047f0d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 17:50:47 +0800 Subject: [PATCH 34/63] style: isort import order in test_dp_atomic_model (reapply dropped pre-commit.ci autofix) --- source/tests/common/dpmodel/test_dp_atomic_model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index 45b5fa7efa..e4592c2087 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -10,12 +10,12 @@ from deepmd.dpmodel.descriptor import ( DescrptSeA, ) -from deepmd.dpmodel.utils.nlist import ( - apply_pair_exclusion_nlist, -) from deepmd.dpmodel.fitting import ( InvarFitting, ) +from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, +) from .case_single_frame_with_nlist import ( TestCaseSingleFrameWithNlist, From fd6b4752d51c0435604499602a95fcb973df37b2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 09:54:25 +0000 Subject: [PATCH 35/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/jax/jax2tf/make_model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index 86923f7e99..dbd521dad2 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -102,9 +102,7 @@ def model_call_from_call_lower( # control flow: its only branch is on the static ``exclude_types`` # config, so it traces cleanly under SavedModel export and does not # need a hand-written TF twin. - from deepmd._vendors import ( - ndtensorflow as ndtf, - ) + from deepmd._vendors import ndtensorflow as ndtf from deepmd.dpmodel.utils.nlist import ( apply_pair_exclusion_nlist, ) From d7c83ed2facffd8e9a111b9297e058a54adca477 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 19:42:52 +0800 Subject: [PATCH 36/63] fix(jax2tf,tf2): guard model.atomic_model when threading pair_excl test_savedmodel_export_contains_xla_call_module passes a DummyModel with no atomic_model attribute; getattr(model.atomic_model, 'pair_excl', None) still evaluated model.atomic_model first and raised AttributeError. Guard the atomic_model access with a nested getattr in both the jax2tf and tf2 serialization callers. --- deepmd/jax/jax2tf/serialization.py | 7 +++++-- deepmd/tf2/utils/serialization.py | 7 +++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index 6b06ec81e0..928798baeb 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -177,8 +177,11 @@ def call( aparam=aparam, do_atomic_virial=do_atomic_virial, # exclusion is a nlist-BUILD transform (decision #18/A4); - # the traced lower consumes a pre-excluded nlist. - pair_excl=getattr(model.atomic_model, "pair_excl", None), + # the traced lower consumes a pre-excluded nlist. Guard + # atomic_model too: test doubles (DummyModel) lack it. + pair_excl=getattr( + getattr(model, "atomic_model", None), "pair_excl", None + ), ) return call diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index 71d9dd699a..2bc7e0d42a 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -391,8 +391,11 @@ def call( aparam=aparam, do_atomic_virial=do_atomic_virial, # exclusion is a nlist-BUILD transform (decision #18/A4); - # the traced lower consumes a pre-excluded nlist. - pair_excl=getattr(model.atomic_model, "pair_excl", None), + # the traced lower consumes a pre-excluded nlist. Guard + # atomic_model too: test doubles (DummyModel) lack it. + pair_excl=getattr( + getattr(model, "atomic_model", None), "pair_excl", None + ), ) ) From 8f6b79692d6dd6c21021e51a9de14047449cf350 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 20:13:33 +0800 Subject: [PATCH 37/63] docs(se-atten): exclude_types no longer disqualifies from the graph path CodeRabbit flagged the pt_expt graph-eligibility sentence as stale: it listed descriptor-level exclude_types as a disqualifier, but this PR removed that gate. uses_graph_lower() now only requires tebd_input_mode == 'concat' and folds exclude_types into the neighbor graph. Drop the exclude_types clause and state its support explicitly. --- doc/model/train-se-atten.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 177c652bed..56bde5dbdc 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -157,7 +157,8 @@ In other backends, type embedding is within this descriptor with the {ref}`tebd_ TensorFlow and other backends have different implementations for {ref}`smooth_type_embedding `. The results are inconsistent when `smooth_type_embedding` is `true`. -In the pt_expt backend, graph-eligible descriptors (mixed types, `tebd_input_mode` `"concat"`, no descriptor-level `exclude_types` or compression) are evaluated by default through the carry-all neighbor-graph path instead of the legacy dense neighbor list. +In the pt_expt backend, graph-eligible descriptors (mixed types, `tebd_input_mode` `"concat"`, not compressed) are evaluated by default through the carry-all neighbor-graph path instead of the legacy dense neighbor list. +Descriptor-level `exclude_types` is supported on the graph path (it is folded into the neighbor graph), so it no longer disqualifies a descriptor from graph evaluation. The graph path considers all neighbors within the cutoff, so its result does not depend on {ref}`sel `. When `smooth_type_embedding` is `true` and {ref}`attn_layer ` is larger than 0 (the defaults), the dense path keeps `sel`-padding phantom terms in the attention softmax denominator while the graph path drops them, so checkpoints trained under the dense semantics shift by up to about 1e-4 in energy when evaluated on the graph path. Passing `neighbor_graph_method="legacy"` to the model forward (or the corresponding evaluation option) restores the dense-path numbers exactly. From f2d750b28962a349b11df5356a910bf860edc4fc Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 20:29:32 +0800 Subject: [PATCH 38/63] fix(stat): model-level pair_exclude enters input stat at nlist BUILD (Piece A) The descriptor input stat (EnvMatStatSe.iter) applied model-level pair_exclude_types as an accumulation-DESELECT (excluded pairs dropped from the count), while the model forward feeds the descriptor a pre-excluded nlist (excluded pairs -> -1, treated like empty slots: env_mat 0, still counted). Those are different stat semantics. Align to the forward (decision #18/A4): fold model-level exclusion into the neighbor list EnvMatStatSe builds (pair_excl on extend_input_and_build_ neighbor_list) and drop the deselect block. Excluded pairs now zero-and-count, identical to descriptor-level exclude_types and to empty slots. This SHIFTS stored davg/dstd for models with pair_exclude_types (intended: it was misaligned with the forward). New invariant test asserts model-level == descriptor-level exclusion bit-identically, plus an exclusion-active control. --- deepmd/dpmodel/utils/env_mat_stat.py | 32 +++++----- .../dpmodel/test_env_mat_stat_exclude.py | 60 +++++++++++++++++++ 2 files changed, 77 insertions(+), 15 deletions(-) create mode 100644 source/tests/common/dpmodel/test_env_mat_stat_exclude.py diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 8d53602b18..999aea25ae 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -203,6 +203,11 @@ def iter( system["box"], ) nframes, nloc = atype.shape[:2] + pair_excl = None + if "pair_exclude_types" in system: + pair_excl = PairExcludeMask( + self.descriptor.get_ntypes(), system["pair_exclude_types"] + ) ( extended_coord, extended_atype, @@ -215,6 +220,13 @@ def iter( self.descriptor.get_sel(), mixed_types=self.descriptor.mixed_types(), box=box, + # Model-level pair exclusion is a nlist-BUILD transform + # (decision #18/A4): fold it in here so the input stat matches + # the model forward, which feeds the descriptor a pre-excluded + # nlist. Excluded pairs then behave exactly like empty slots + # (env_mat 0, still counted) -- identical to descriptor-level + # exclude_types, replacing the previous accumulation-deselect. + pair_excl=pair_excl, ) env_mat_caller = EnvMat( self.descriptor.get_rcut(), @@ -258,21 +270,11 @@ def iter( (-1, 1), ), ) - if "pair_exclude_types" in system: - pair_exclude_mask = PairExcludeMask( - self.descriptor.get_ntypes(), system["pair_exclude_types"] - ) - pair_exclude_mask.type_mask = xp.asarray( - pair_exclude_mask.type_mask, - device=array_api_compat.device(atype), - ) - # shape: (1, nloc, nnei) - exclude_mask = xp.reshape( - pair_exclude_mask.build_type_exclude_mask(nlist, extended_atype), - (1, nframes * nloc, -1), - ) - # shape: (ntypes, nloc, nnei) - type_idx = xp.logical_and(type_idx[..., None], exclude_mask) + # NOTE: model-level ``pair_exclude_types`` is NOT re-applied here. + # It is folded into the neighbor list at BUILD time above + # (decision #18/A4), so excluded pairs already have env_mat == 0 + # and are counted like empty slots -- the same treatment the model + # forward gives them. for type_i in range(self.descriptor.get_ntypes()): dd = env_mat[type_idx[type_i, ...]] dd = xp.reshape( diff --git a/source/tests/common/dpmodel/test_env_mat_stat_exclude.py b/source/tests/common/dpmodel/test_env_mat_stat_exclude.py new file mode 100644 index 0000000000..2f76602c4f --- /dev/null +++ b/source/tests/common/dpmodel/test_env_mat_stat_exclude.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Piece A: model-level pair exclusion enters the descriptor INPUT stat at +nlist BUILD, matching the model forward. + +Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision +#18/A4). In the forward the descriptor is fed a pre-excluded nlist, so an +excluded pair is treated exactly like an empty slot (env_mat 0, still counted). +The input stat now folds the same exclusion into the neighbor list it builds +(``EnvMatStatSe.iter``) instead of deselecting excluded pairs during +accumulation. + +The load-bearing invariant: excluding a type-pair via model-level +``pair_exclude_types`` must give the SAME stat (davg/dstd) as excluding it via +descriptor-level ``exclude_types`` -- both zero-and-count the excluded pairs. +Before this change they diverged (descriptor-level zero-counted, model-level +deselected). A no-exclusion control proves the exclusion is genuinely active. +""" + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) + + +def _sample() -> dict: + rng = np.random.default_rng(0) + nf, nloc = 2, 6 + coord = rng.normal(size=(nf, nloc, 3)) * 2.0 + # both types present so the (0, 1) exclusion actually removes pairs + atype = np.array([[0, 1, 0, 1, 0, 1], [1, 0, 1, 0, 1, 0]], dtype=np.int64) + box = np.tile((np.eye(3) * 12.0).reshape(1, 9), (nf, 1)) + return {"coord": coord, "atype": atype, "box": box} + + +def _stats(exclude_types, pair_exclude_types): + """Compute (davg, dstd) for a se_e2_a descriptor under the given exclusions.""" + ds = DescrptSeA(6.0, 0.5, [10, 10], exclude_types=exclude_types) + sample = _sample() + if pair_exclude_types is not None: + sample["pair_exclude_types"] = pair_exclude_types + ds.compute_input_stats([sample]) + return np.asarray(ds.davg), np.asarray(ds.dstd) + + +def test_model_level_equals_descriptor_level_exclusion() -> None: + """Model-level pair_exclude of (0,1) == descriptor-level exclude_types of (0,1).""" + davg_model, dstd_model = _stats(exclude_types=[], pair_exclude_types=[[0, 1]]) + davg_desc, dstd_desc = _stats(exclude_types=[[0, 1]], pair_exclude_types=None) + # Both zero-and-count the same excluded pairs, so the env-matrix + # distribution -- hence the stat -- is bit-identical. + np.testing.assert_allclose(davg_model, davg_desc, rtol=1e-14, atol=1e-14) + np.testing.assert_allclose(dstd_model, dstd_desc, rtol=1e-14, atol=1e-14) + + +def test_exclusion_is_active() -> None: + """Excluding (0,1) shifts the stat vs no exclusion (proves it is applied).""" + davg_none, dstd_none = _stats(exclude_types=[], pair_exclude_types=None) + davg_excl, dstd_excl = _stats(exclude_types=[], pair_exclude_types=[[0, 1]]) + assert not np.allclose(dstd_none, dstd_excl) From dd5de26202ac8a8b7b8a84f4a7cdd5dc9c04e00d Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 22:21:30 +0800 Subject: [PATCH 39/63] feat(stat): compute dpa1 input stat env matrix via NeighborGraph (Piece B) The dpa1 forward computes its env matrix through the NeighborGraph (from_dense_quartet -> edge_env_mat); the input stat used the dense EnvMat. Route the dpa1 block's input stat through the SAME graph path so stat and forward share one env-matrix implementation, with both exclusions folded in exactly as the forward does (model-level via the pre-excluded nlist from Piece A; descriptor-level via the emask mask). BIT-IDENTICAL to the dense path, so stored davg/dstd are unchanged: from_dense_quartet(compact=False) reuses the same neighbor set + padding (row-major (frame,center,slot) edges), edge_env_mat mirrors EnvMat.call, and the (E,4) output reshapes 1:1 to the dense (nf,nloc,nsel,4) tensor. Opt-in via EnvMatStatSe(use_graph=True); se_e2_a/se_r are untouched. pt_expt inherits it via autowrap; legacy pt stays dense (bit-identical, so cross-backend parity holds). Tests: graph==dense stat bit-identical (1e-15) for se_e2_a and the dpa1 block, under no exclusion, model-level pair_exclude, and descriptor-level exclude. --- deepmd/dpmodel/descriptor/dpa1.py | 6 +- deepmd/dpmodel/utils/env_mat_stat.py | 101 +++++++++++++++--- .../common/dpmodel/test_env_mat_stat_graph.py | 84 +++++++++++++++ 3 files changed, 176 insertions(+), 15 deletions(-) create mode 100644 source/tests/common/dpmodel/test_env_mat_stat_graph.py diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 64d3589e40..b03da78118 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -1401,7 +1401,11 @@ def compute_input_stats( The path to the stat file. """ - env_mat_stat = EnvMatStatSe(self) + # dpa1's forward computes its env matrix through the NeighborGraph + # (from_dense_quartet -> edge_env_mat); run the input stat through the + # SAME path so stat and forward share one env-matrix implementation. + # Bit-identical to the dense EnvMat (see test_env_mat_stat_graph.py). + env_mat_stat = EnvMatStatSe(self, use_graph=True) if path is not None: path = path / env_mat_stat.get_hash() if path is None or not path.is_dir(): diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 999aea25ae..36a42fc441 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -15,6 +15,7 @@ ) from deepmd.dpmodel.array_api import ( Array, + xp_take_first_n, ) from deepmd.dpmodel.common import ( get_xp_precision, @@ -22,6 +23,10 @@ from deepmd.dpmodel.utils.env_mat import ( EnvMat, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + edge_env_mat, + from_dense_quartet, +) from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, ) @@ -143,12 +148,73 @@ class EnvMatStatSe(EnvMatStat): The descriptor of the model. """ - def __init__(self, descriptor: Union["Descriptor", "DescriptorBlock"]) -> None: + def __init__( + self, + descriptor: Union["Descriptor", "DescriptorBlock"], + use_graph: bool = False, + ) -> None: super().__init__() self.descriptor = descriptor self.last_dim = ( self.descriptor.ndescrpt // self.descriptor.nnei ) # se_r=1, se_a=4 + # ``use_graph`` computes the env matrix through the NeighborGraph path + # (``from_dense_quartet`` -> ``edge_env_mat``) instead of the dense + # ``EnvMat``, so the input stat runs the SAME machinery the dpa1 graph + # forward uses. It is BIT-IDENTICAL to the dense path (same neighbor + # set + padding, ``edge_env_mat`` mirrors ``EnvMat.call``, row-major + # ``(frame, center, slot)`` edges reshape 1:1 to ``(nf, nloc, nsel)``); + # only se_a-type (``last_dim == 4``) descriptors may opt in. + self.use_graph = use_graph + + def _graph_env_mat( + self, + extended_coord: Array, + extended_atype: Array, + mapping: Array, + nlist: Array, + ) -> Array: + """Env matrix via the NeighborGraph, shaped ``(nf, nloc, nsel, last_dim)``. + + Bit-identical to the dense ``EnvMat.call`` with zero mean / unit std: + ``from_dense_quartet(compact=False)`` reuses the same neighbor set and + padding (row-major ``(frame, center, slot)`` edges), ``edge_env_mat`` + mirrors ``EnvMat.call``, and padding / model-excluded edges (already + ``-1`` in the pre-excluded ``nlist``) carry ``edge_mask=False`` and are + zeroed -- so the ``(E, 4)`` output reshapes 1:1 back to the dense + ``(nf, nloc, nsel, 4)`` env-matrix tensor. + """ + xp = array_api_compat.array_namespace(extended_coord, nlist) + dev = array_api_compat.device(extended_coord) + nframes, nloc, nsel = nlist.shape + nall = extended_atype.shape[1] + ntypes = self.descriptor.get_ntypes() + coord_ext_3 = xp.reshape(extended_coord, (nframes, nall, 3)) + mapping_g = xp.reshape(mapping, (nframes, nall)) + graph = from_dense_quartet(coord_ext_3, nlist, mapping_g, compact=False) + # local center type per edge (dst is the local center index) + atype_local = xp.reshape( + xp_take_first_n(extended_atype, 1, nloc), (nframes * nloc,) + ) + center_type = xp.take(atype_local, graph.edge_index[1, :], axis=0) + zero2 = xp.zeros((ntypes, 4), dtype=graph.edge_vec.dtype, device=dev) + one2 = xp.ones((ntypes, 4), dtype=graph.edge_vec.dtype, device=dev) + em = edge_env_mat( + graph.edge_vec, + center_type, + zero2, + one2, + self.descriptor.get_rcut(), + self.descriptor.get_rcut_smth(), + protection=self.descriptor.get_env_protection(), + edge_mask=graph.edge_mask, + return_sw=False, + ) # (E, 4) + # zero padding / model-excluded edges (edge_mask=False) so they count + # as 0 -- exactly like empty slots in the dense path. + em = em * xp.astype(graph.edge_mask[:, None], em.dtype) + # row-major (frame, center, slot) -> dense (nf, nloc, nsel, last_dim) + return xp.reshape(em, (nframes, nloc, nsel, self.last_dim)) def iter( self, data: list[dict[str, np.ndarray | list[tuple[int, int]]]] @@ -228,19 +294,26 @@ def iter( # exclude_types, replacing the previous accumulation-deselect. pair_excl=pair_excl, ) - env_mat_caller = EnvMat( - self.descriptor.get_rcut(), - self.descriptor.get_rcut_smth(), - protection=self.descriptor.get_env_protection(), - ) - env_mat, _, _ = env_mat_caller.call( - extended_coord, - extended_atype, - nlist, - zero_mean, - one_stddev, - radial_only, - ) + if self.use_graph: + # NeighborGraph env matrix (bit-identical to the dense EnvMat + # below): the SAME machinery the dpa1 graph forward uses. + env_mat = self._graph_env_mat( + extended_coord, extended_atype, mapping, nlist + ) + else: + env_mat_caller = EnvMat( + self.descriptor.get_rcut(), + self.descriptor.get_rcut_smth(), + protection=self.descriptor.get_env_protection(), + ) + env_mat, _, _ = env_mat_caller.call( + extended_coord, + extended_atype, + nlist, + zero_mean, + one_stddev, + radial_only, + ) # apply excluded_types exclude_mask = self.descriptor.emask.build_type_exclude_mask( nlist, extended_atype diff --git a/source/tests/common/dpmodel/test_env_mat_stat_graph.py b/source/tests/common/dpmodel/test_env_mat_stat_graph.py new file mode 100644 index 0000000000..0f0243aa11 --- /dev/null +++ b/source/tests/common/dpmodel/test_env_mat_stat_graph.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Piece B: the input stat's env matrix computed via the NeighborGraph +(``from_dense_quartet`` -> ``edge_env_mat``) is BIT-IDENTICAL to the dense +``EnvMat`` path. + +This is the same machinery the dpa1 graph forward uses. ``from_dense_quartet`` +reuses the same neighbor set and padding, ``edge_env_mat`` mirrors +``EnvMat.call``, and the row-major ``(frame, center, slot)`` edge order maps 1:1 +back to the dense ``(nf, nloc, nsel, 4)`` env-matrix tensor -- so the stored +``davg``/``dstd`` are unchanged. The parity must hold with no exclusion, with +model-level ``pair_exclude_types`` (folded into the nlist at BUILD, Piece A), +and with descriptor-level ``exclude_types``. +""" + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptDPA1, + DescrptSeA, +) +from deepmd.dpmodel.utils.env_mat_stat import ( + EnvMatStatSe, +) + + +def _sample() -> dict: + rng = np.random.default_rng(0) + nf, nloc = 2, 6 + coord = rng.normal(size=(nf, nloc, 3)) * 2.0 + atype = np.array([[0, 1, 0, 1, 0, 1], [1, 0, 1, 0, 1, 0]], dtype=np.int64) + box = np.tile((np.eye(3) * 12.0).reshape(1, 9), (nf, 1)) + return {"coord": coord, "atype": atype, "box": box} + + +def _stats(exclude_types, pair_exclude_types, use_graph): + ds = DescrptSeA(6.0, 0.5, [10, 10], exclude_types=exclude_types) + sample = _sample() + if pair_exclude_types is not None: + sample["pair_exclude_types"] = pair_exclude_types + st = EnvMatStatSe(ds, use_graph=use_graph) + st.load_or_compute_stats([sample], None) + mean, stddev = st() + return np.asarray(mean), np.asarray(stddev) + + +def _assert_graph_equals_dense(exclude_types, pair_exclude_types) -> None: + m_dense, s_dense = _stats(exclude_types, pair_exclude_types, use_graph=False) + m_graph, s_graph = _stats(exclude_types, pair_exclude_types, use_graph=True) + np.testing.assert_allclose(m_graph, m_dense, rtol=1e-15, atol=1e-15) + np.testing.assert_allclose(s_graph, s_dense, rtol=1e-15, atol=1e-15) + + +def test_graph_equals_dense_no_exclusion() -> None: + _assert_graph_equals_dense(exclude_types=[], pair_exclude_types=None) + + +def test_graph_equals_dense_model_level_exclusion() -> None: + _assert_graph_equals_dense(exclude_types=[], pair_exclude_types=[[0, 1]]) + + +def test_graph_equals_dense_descriptor_level_exclusion() -> None: + _assert_graph_equals_dense(exclude_types=[[0, 1]], pair_exclude_types=None) + + +def _dpa1_block_stats(pair_exclude_types, use_graph): + """Stats for the actual dpa1 (mixed_types) block, the wired descriptor.""" + ds = DescrptDPA1(6.0, 0.5, 20, ntypes=2, attn_layer=0) + block = ds.se_atten + sample = _sample() + if pair_exclude_types is not None: + sample["pair_exclude_types"] = pair_exclude_types + st = EnvMatStatSe(block, use_graph=use_graph) + st.load_or_compute_stats([sample], None) + mean, stddev = st() + return np.asarray(mean), np.asarray(stddev) + + +def test_dpa1_block_graph_equals_dense() -> None: + """The wired dpa1 block (use_graph=True) is bit-identical to the dense path.""" + for pair_exclude_types in (None, [[0, 1]]): + m_dense, s_dense = _dpa1_block_stats(pair_exclude_types, use_graph=False) + m_graph, s_graph = _dpa1_block_stats(pair_exclude_types, use_graph=True) + np.testing.assert_allclose(m_graph, m_dense, rtol=1e-15, atol=1e-15) + np.testing.assert_allclose(s_graph, s_dense, rtol=1e-15, atol=1e-15) From 06e6ede4d9d5964dd41565141a5e0eefcfa654c0 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 22:29:18 +0800 Subject: [PATCH 40/63] refactor(neighbor_graph): extract graph_from_dense_quartet helper (DRY) _graph_env_mat (input stat) duplicated the dense-quartet -> (graph, atype_local) setup from DescrptDPA1._call_graph_adapter almost verbatim (coord reshape, mapping-None identity, from_dense_quartet compact=False, xp_take_first_n local atype). Extract it as neighbor_graph.graph_from_dense_ quartet(coord_ext, atype_ext, nlist, mapping) -> (graph, atype_local) and route both call sites through it. Pure extraction, bit-identical (dpa1 adapter + stat parity tests unchanged). --- deepmd/dpmodel/descriptor/dpa1.py | 20 ++----- deepmd/dpmodel/utils/env_mat_stat.py | 13 ++--- .../dpmodel/utils/neighbor_graph/__init__.py | 2 + .../dpmodel/utils/neighbor_graph/builder.py | 56 +++++++++++++++++++ 4 files changed, 67 insertions(+), 24 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index b03da78118..3227bbc39e 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -643,26 +643,16 @@ def _call_graph_adapter( The smooth switch function. shape: nf x nloc x nnei x 1 """ from deepmd.dpmodel.utils.neighbor_graph import ( - from_dense_quartet, + graph_from_dense_quartet, ) xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) - dev = array_api_compat.device(coord_ext) nf, nloc, nnei = nlist.shape - nall = xp.reshape(coord_ext, (nf, -1)).shape[1] // 3 - coord_ext_3 = xp.reshape(coord_ext, (nf, nall, 3)) - if mapping is None: - # default identity mapping (ext == loc, e.g. no-PBC nall == nloc) - mapping_g = xp.broadcast_to( - xp.arange(nall, dtype=xp.int64, device=dev)[None, :], (nf, nall) - ) - else: - mapping_g = xp.reshape(mapping, (nf, nall)) - graph = from_dense_quartet( - coord_ext_3, nlist, mapping_g, layout=None, compact=False + # shape-static graph + flat local center types from the dense quartet + # (shared with the input-stat graph path, see graph_from_dense_quartet). + graph, atype_local = graph_from_dense_quartet( + coord_ext, atype_ext, nlist, mapping ) - # local atom types, flat (nf * nloc,) - atype_local = xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (nf * nloc,)) grrg_flat, rot_mat_flat = self.call_graph( graph, atype_local, diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 36a42fc441..d03af7dd88 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -15,7 +15,6 @@ ) from deepmd.dpmodel.array_api import ( Array, - xp_take_first_n, ) from deepmd.dpmodel.common import ( get_xp_precision, @@ -25,7 +24,7 @@ ) from deepmd.dpmodel.utils.neighbor_graph import ( edge_env_mat, - from_dense_quartet, + graph_from_dense_quartet, ) from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, @@ -187,15 +186,11 @@ def _graph_env_mat( xp = array_api_compat.array_namespace(extended_coord, nlist) dev = array_api_compat.device(extended_coord) nframes, nloc, nsel = nlist.shape - nall = extended_atype.shape[1] ntypes = self.descriptor.get_ntypes() - coord_ext_3 = xp.reshape(extended_coord, (nframes, nall, 3)) - mapping_g = xp.reshape(mapping, (nframes, nall)) - graph = from_dense_quartet(coord_ext_3, nlist, mapping_g, compact=False) - # local center type per edge (dst is the local center index) - atype_local = xp.reshape( - xp_take_first_n(extended_atype, 1, nloc), (nframes * nloc,) + graph, atype_local = graph_from_dense_quartet( + extended_coord, extended_atype, nlist, mapping ) + # local center type per edge (dst is the local center index) center_type = xp.take(atype_local, graph.edge_index[1, :], axis=0) zero2 = xp.zeros((ntypes, 4), dtype=graph.edge_vec.dtype, device=dev) one2 = xp.ones((ntypes, 4), dtype=graph.edge_vec.dtype, device=dev) diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 5cb84d4e9d..f29738c05c 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -15,6 +15,7 @@ from .builder import ( build_neighbor_graph, from_dense_quartet, + graph_from_dense_quartet, ) from .derivatives import ( edge_force_virial, @@ -54,6 +55,7 @@ "edge_force_virial", "frame_id_from_n_node", "from_dense_quartet", + "graph_from_dense_quartet", "neighbor_graph_from_ijs", "node_validity_mask", "pad_and_guard_edges", diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index 0b39c1ddcb..54f70e0bc8 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -39,6 +39,10 @@ import array_api_compat +from deepmd.dpmodel.array_api import ( + xp_take_first_n, +) + from .graph import ( GraphLayout, NeighborGraph, @@ -202,6 +206,58 @@ def from_dense_quartet( ) +def graph_from_dense_quartet( + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None, +) -> tuple[NeighborGraph, Array]: + """Shape-static NeighborGraph + flat local center atype from a dense quartet. + + Convenience wrapper over :func:`from_dense_quartet` (``compact=False``) that + also fills the ``mapping is None`` identity case (no-PBC, ``nall == nloc``) + and derives the ghost-free flat local atom types ``(nf * nloc,)``. Shared by + the dpa1 dense->graph forward adapter (``DescrptDPA1._call_graph_adapter``) + and the input-stat graph path (``EnvMatStatSe._graph_env_mat``) so both build + the graph identically. + + Parameters + ---------- + coord_ext + Extended coordinates, ``(nf, nall x 3)`` or ``(nf, nall, 3)``. + atype_ext + Extended atom types, ``(nf, nall)``. + nlist + Dense neighbor list into the extended atoms, ``(nf, nloc, nsel)``; ``-1`` + is padding. + mapping + Extended -> local-owner index, ``(nf, nall)``; ``None`` means the + identity mapping (``nall == nloc``). + + Returns + ------- + graph + Shape-static :class:`NeighborGraph` over the LOCAL atoms. + atype_local + Flat local atom types, ``(nf * nloc,)``. + """ + xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) + dev = array_api_compat.device(coord_ext) + nf, nloc, _ = nlist.shape + nall = atype_ext.shape[1] + coord_ext_3 = xp.reshape(coord_ext, (nf, nall, 3)) + if mapping is None: + # default identity mapping (ext == loc, e.g. no-PBC nall == nloc) + mapping_g = xp.broadcast_to( + xp.arange(nall, dtype=xp.int64, device=dev)[None, :], (nf, nall) + ) + else: + mapping_g = xp.reshape(mapping, (nf, nall)) + graph = from_dense_quartet(coord_ext_3, nlist, mapping_g, compact=False) + atype_local = xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (nf * nloc,)) + return graph, atype_local + + def build_neighbor_graph( coord: Array, atype: Array, From c564e2d066660aa289557365119be311140d32b9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 22:57:29 +0000 Subject: [PATCH 41/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/utils/env_mat_stat.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index d03af7dd88..2fde517421 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -22,13 +22,13 @@ from deepmd.dpmodel.utils.env_mat import ( EnvMat, ) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( edge_env_mat, graph_from_dense_quartet, ) -from deepmd.dpmodel.utils.exclude_mask import ( - PairExcludeMask, -) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, ) From ff283707999a2d7528c2a18e504ef0d86439c908 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 8 Jul 2026 16:53:36 +0800 Subject: [PATCH 42/63] fix(pt-expt): apply pair exclusion on the dense multi-rank (with-comm) path Address the pair_exclude_types review on #5733 (njzjz-bot #1, OutisLi). Decision #18/A4 moved model-level pair_exclude_types to a build-time C++ ingestion seam, but the seam was wired into only the single-rank dense (applyPairExclusionNlist) and graph (applyPairExclusion) routes. The multi-rank dense path (run_model_with_comm, DeepPotPTExpt.cc) passed the raw nlist, so a message-passing .pt2 with pair_exclude_types silently included excluded pairs multi-rank (multi-rank != single-rank). Apply the SAME seam before run_model_with_comm, mirroring the single-rank dense site exactly (the cross-rank ghost exchange happens inside run_model_with_comm and does not change the nlist's meaning, so per-rank pre-exclusion is correct). Scope: the edge-schema lowers (run_model_edges*) are intentionally NOT touched -- the edge lower is a pt-backend export (SeZM/DPA4 via pt/entrypoints/freeze_pt2.py) and the pt backend still applies exclusion at consume time, so the edge .pt2 bakes exclusion into the graph. The MP message-passing GRAPH lower stays fail-fast (PR-G). Perf (OutisLi): buildPairExcludeTable's flat table was rebuilt and copied H2D on every compute() (MD) step. Upload it once in init as a device torch::Tensor member (pair_exclude_table_); the seam helpers now take the prebuilt tensor and index_select it directly (undefined tensor => identity, mirroring the old empty-vector early-exit). No per-step allocation or transfer. Docs: reconcile the stale "idempotent backstop" wording (vesin_neighbor_list, pt_expt/model/make_model) with the single-owner design -- both dpmodel seams (forward_common_atomic{,_graph}) now say "NOT applied here"; the build site is the sole owner. Test (cell 5): add a dpa1 MP-graph + pair_exclude_types parity test to test_lammps_dpa1_graph_pt2.py -- MP (-n 2) == SP (-n 1) on the excluded model, plus a cross-check that the excluded run differs from the same-weights no-exclusion baseline (so a silently-dropped exclusion on both ranks cannot pass trivially). Reuses the existing MPI harness with a pb= override. Verified locally (CPU): 8/8 single-rank pairexcl gtests pass after the perf refactor; all 6 tests in test_lammps_dpa1_graph_pt2.py pass (incl. the new multi-rank exclusion test). --- deepmd/pt_expt/model/make_model.py | 6 +- deepmd/pt_expt/utils/vesin_neighbor_list.py | 8 +-- source/api_cc/include/DeepPotPTExpt.h | 18 ++--- source/api_cc/include/commonPT.h | 52 +++++++------- source/api_cc/src/DeepPotPTExpt.cc | 36 ++++++++-- .../lmp/tests/test_lammps_dpa1_graph_pt2.py | 67 ++++++++++++++++++- 6 files changed, 138 insertions(+), 49 deletions(-) diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index ee89b7bd78..a080600709 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -468,8 +468,10 @@ def _call_common_graph( "graph lower (e.g. dpa1 attn_layer=0)" ) rcut = self.get_rcut() - # Model-level pair_exclude_types — apply at build time so the seam - # backstop in forward_atomic_graph acts as an idempotent identity. + # Model-level pair_exclude_types is a graph-BUILD transform + # (decision #18): apply it here, at the single owning site, so the + # exported lower (forward_common_atomic_graph, which no longer + # re-applies it) consumes a pre-excluded edge_mask. pair_excl = getattr(self.atomic_model, "pair_excl", None) if method == "dense": ng = build_neighbor_graph(cc, atype, bb, rcut, pair_excl=pair_excl) diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 4965dbf77a..0c7e5a22ec 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -79,10 +79,10 @@ def build( When provided, excluded type pairs are erased from the returned neighbor list (entries set to ``-1``) by :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. - The atomic-model seam applies the same filter as an idempotent - backstop, so passing ``pair_excl`` here is a build-time - optimization that avoids re-scanning per forward call. - ``return_mode='edges'`` does not support ``pair_excl``; a + This is the OWNING application site (decision #18/A4): the exported + lower (``forward_common_atomic``) no longer re-applies model-level + ``pair_exclude_types`` -- it is applied once here at nlist BUILD + time. ``return_mode='edges'`` does not support ``pair_excl``; a :class:`NotImplementedError` is raised in that combination. """ if return_mode == "edges" and pair_excl is not None: diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 4f14e0178f..e2b92c2924 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -331,14 +331,16 @@ class DeepPotPTExpt : public DeepPotBackend { // continue to work; GNN archives must be regenerated to opt into // the fail-fast guard against the silent-corruption bug. bool has_message_passing_ = false; - // Flat (ntypes+1)^2 model-level pair-type keep table, rebuilt in ``init`` - // from the ``pair_exclude_types`` metadata field (see - // ``deepmd::buildPairExcludeTable``). Empty => no model-level exclusion. - // Exclusion is a BUILD-time transform (decision #18/A4): the C++ ingestion - // seam is the single application site (``applyPairExclusion`` graph / - // ``applyPairExclusionNlist`` dense); the exported .pt2 lowers consume - // pre-excluded inputs and never re-apply it. - std::vector pair_exclude_table_; + // Device-resident (ntypes+1)^2 model-level pair-type keep table, uploaded + // ONCE in ``init`` from the ``pair_exclude_types`` metadata field (see + // ``deepmd::buildPairExcludeTable``). An UNDEFINED tensor => no model-level + // exclusion (identity). The device is fixed at ``init`` (``gpu_id`` / + // ``gpu_enabled``), so the seam helpers ``index_select`` it directly with no + // per-step CPU clone / H2D upload. Exclusion is a BUILD-time transform + // (decision #18/A4): the C++ ingestion seam is the single application site + // (``applyPairExclusion`` graph / ``applyPairExclusionNlist`` dense); the + // exported .pt2 lowers consume pre-excluded inputs and never re-apply it. + torch::Tensor pair_exclude_table_; std::unique_ptr with_comm_tempfile_; std::unique_ptr with_comm_loader; diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 48cf984cf1..fd88d933ac 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -517,39 +517,37 @@ inline std::vector buildPairExcludeTable( * * OWNERSHIP: exclusion is a BUILD-time transform (decision #18/A4). The * exported ``.pt2`` graph lower consumes a pre-excluded ``edge_mask`` and - * never re-applies it, so this call is the SINGLE application site on the - * C++ graph route. + * never re-applies it, so the C++ ingestion seam is the sole owner: this helper + * is called once on each dispatched graph run path (single-rank and the + * non-message-passing multi-rank extended-region path), never twice on the same + * tensors. * * @param edge_index (2, E) int64 ``[src, dst]``; src = neighbor, dst = center. * @param edge_mask (E,) bool real-vs-padding mask to be ANDed in place. * @param atype (N,) int64 flat node types (clamped >= 0). - * @param type_mask_table Flat ``(ntypes+1)^2`` table from - * ``buildPairExcludeTable``. Empty => identity (returns ``edge_mask``). + * @param type_mask_table Device-resident flat ``(ntypes+1)^2`` int32 keep table + * uploaded once in ``init`` (``buildPairExcludeTable`` + one H2D copy). An + * UNDEFINED tensor => identity (returns ``edge_mask``), mirroring the old + * empty-vector early-exit. Already on the model device (== ``edge_mask`` + * device), so ``index_select`` needs no per-call transfer. * @param ntypes Number of real atom types. * @return New ``edge_mask`` with excluded edges cleared. */ inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, const torch::Tensor& edge_mask, const torch::Tensor& atype, - const std::vector& type_mask_table, + const torch::Tensor& type_mask_table, const int ntypes) { - if (type_mask_table.empty()) { + if (!type_mask_table.defined()) { return edge_mask; } - const auto device = edge_mask.device(); const auto src = edge_index.index({0}); // (E,) neighbour const auto dst = edge_index.index({1}); // (E,) center const auto src_t = atype.index_select(0, src); const auto dst_t = atype.index_select(0, dst); // type_ij = atype[dst] * (ntypes + 1) + atype[src] (matches Python) const auto type_ij = dst_t * (ntypes + 1) + src_t; - const auto table = - torch::from_blob(const_cast(type_mask_table.data()), - {static_cast(type_mask_table.size())}, - torch::TensorOptions().dtype(torch::kInt32)) - .clone() - .to(device); - const auto keep = table.index_select(0, type_ij).to(torch::kBool); + const auto keep = type_mask_table.index_select(0, type_ij).to(torch::kBool); return torch::logical_and(edge_mask, keep); } @@ -563,25 +561,27 @@ inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, * * OWNERSHIP: exclusion is a BUILD-time transform (decision #18/A4). The * exported ``.pt2`` dense lower consumes a pre-excluded nlist and never - * re-applies it, so this call is the SINGLE application site on the C++ - * dense route. + * re-applies it, so the C++ ingestion seam is the sole owner: this helper is + * called once on each dispatched dense run path (single-rank ``run_model`` and + * multi-rank ``run_model_with_comm``), never twice on the same tensors. * * @param nlist (nf, nloc, nnei) int64 neighbour list; ``-1`` == empty slot. * @param atype_ext (nf, nall) int64 extended atom types. - * @param type_mask_table Flat ``(ntypes+1)^2`` table from - * ``buildPairExcludeTable``. Empty => identity (returns ``nlist``). + * @param type_mask_table Device-resident flat ``(ntypes+1)^2`` int32 keep table + * uploaded once in ``init``. An UNDEFINED tensor => identity (returns + * ``nlist``). Already on the model device (== ``nlist`` device), so + * ``index_select`` needs no per-call transfer. * @param ntypes Number of real atom types. * @return New neighbour list with excluded entries set to ``-1``. */ inline torch::Tensor applyPairExclusionNlist( const torch::Tensor& nlist, const torch::Tensor& atype_ext, - const std::vector& type_mask_table, + const torch::Tensor& type_mask_table, const int ntypes) { - if (type_mask_table.empty()) { + if (!type_mask_table.defined()) { return nlist; } - const auto device = nlist.device(); const std::int64_t nf = nlist.size(0); const std::int64_t nloc = nlist.size(1); const std::int64_t nnei = nlist.size(2); @@ -598,14 +598,8 @@ inline torch::Tensor applyPairExclusionNlist( // type_ij = type_i * (ntypes + 1) + type_j (matches Python: type_i already // scaled above; here just add the neighbour type). const auto type_ij = type_i.unsqueeze(2) + type_j; // (nf, nloc, nnei) - const auto table = - torch::from_blob(const_cast(type_mask_table.data()), - {static_cast(type_mask_table.size())}, - torch::TensorOptions().dtype(torch::kInt32)) - .clone() - .to(device); - const auto keep = - table.index_select(0, type_ij.reshape({-1})).reshape({nf, nloc, nnei}); + const auto keep = type_mask_table.index_select(0, type_ij.reshape({-1})) + .reshape({nf, nloc, nnei}); return torch::where(keep == 1, nlist, torch::full_like(nlist, -1)); } diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 9dd544bbad..f6f73ae95c 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -207,6 +207,11 @@ void DeepPotPTExpt::init(const std::string& model, // (decision #18/A4): the C++ ingestion seam is the single application site // (applyPairExclusion graph / applyPairExclusionNlist dense); the exported // lowers consume pre-excluded inputs and never re-apply it. + // + // Upload the table to the model device ONCE here (device is fixed by + // ``gpu_id`` / ``gpu_enabled``), so the per-step seam helpers only + // ``index_select`` it -- no per-``compute()`` CPU clone + H2D copy. Leave + // ``pair_exclude_table_`` undefined (=> identity) when there is no exclusion. { std::vector> pair_exclude_types; if (metadata.obj_val.count("pair_exclude_types")) { @@ -214,8 +219,20 @@ void DeepPotPTExpt::init(const std::string& model, pair_exclude_types.emplace_back(v[0].as_int(), v[1].as_int()); } } - pair_exclude_table_ = + std::vector tbl = deepmd::buildPairExcludeTable(ntypes, pair_exclude_types); + if (!tbl.empty()) { + torch::Device device(torch::kCUDA, gpu_id); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + } + pair_exclude_table_ = + torch::from_blob(tbl.data(), + {static_cast(tbl.size())}, + torch::TensorOptions().dtype(torch::kInt32)) + .clone() + .to(device); + } } if (has_comm_artifact_) { try { @@ -839,8 +856,19 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor, charge_spin_tensor, comm_tensors); } } else { + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported dense lower consumes a pre-excluded nlist and + // never re-applies it. The multi-rank (with-comm) dense route shares the + // same dense nlist as the single-rank path below, so it applies the SAME + // seam -- otherwise a message-passing .pt2 with pair_exclude_types would + // silently include excluded pairs on the with-comm path (multi-rank != + // single-rank). The cross-rank ghost exchange happens inside + // run_model_with_comm and does not change the nlist's meaning, so + // pre-excluding it is correct per rank. + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = run_model_with_comm( - coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, + coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); } } else { @@ -906,8 +934,8 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, } else { // Model-level pair exclusion is a BUILD-time transform (decision // #18/A4): the exported dense lower consumes a pre-excluded nlist and - // never re-applies it; this is the single application site on the C++ - // dense route. + // never re-applies it. Single-rank dense application site; the + // multi-rank (with-comm) dense sibling above applies the same seam. const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = diff --git a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py index ab898e15f6..b54dbd8cce 100644 --- a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py @@ -44,6 +44,21 @@ pb_file = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa1_graph.pt2" ) +# Graph-lower dpa1 with model-level pair_exclude_types=[[0,1]] and its +# same-weights no-exclusion baseline (source/tests/infer/gen_dpa1_pairexcl.py). +# Used to prove exclusion survives the extended-region multi-rank graph path. +pb_file_pairexcl = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa1_pairexcl_graph.pt2" +) +pb_file_pairexcl_none = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa1_pairexcl_none.pt2" +) ref_file = ( Path(__file__).parent.parent.parent / "tests" @@ -193,6 +208,7 @@ def _run_mpi_subprocess( data_path: Path | None = None, processors: str | None = None, runner_args: list[str] | None = None, + pb: Path | None = None, ) -> dict: """Invoke the (backend-agnostic) DPA3 MPI runner under ``mpirun -n `` against the dpa1 graph .pt2 and return @@ -200,10 +216,13 @@ def _run_mpi_subprocess( ``nprocs == 1`` forces ``--processors 1 1 1`` so the C++ side sees ``nprocs == 1`` and routes to the single-rank graph path — a - same-archive reference for the multi-rank comparison. + same-archive reference for the multi-rank comparison. ``pb`` overrides + the model archive (defaults to the no-exclusion ``deeppot_dpa1_graph.pt2``). """ if data_path is None: data_path = data_file + if pb is None: + pb = pb_file with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: out_path = f.name try: @@ -214,7 +233,7 @@ def _run_mpi_subprocess( sys.executable, str(mpi_runner), str(data_path.resolve()), - str(pb_file.resolve()), + str(pb.resolve()), out_path, ] if processors is not None: @@ -317,3 +336,47 @@ def test_pair_deepmd_mpi_dpa1_graph_empty_subdomain() -> None: out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=0 ) assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +@pytest.mark.skipif( + not pb_file_pairexcl.exists(), + reason="gen_dpa1_pairexcl.py .pt2 fixtures not generated", +) +def test_pair_deepmd_mpi_dpa1_pairexcl_graph_matches_single_rank() -> None: + """Model-level ``pair_exclude_types`` must survive the extended-region + multi-rank graph path (cell 5). + + Exclusion is a BUILD-time transform applied at the C++ ingestion seam + (``applyPairExclusion`` on ``edge_mask``); the SAME seam serves the + single-rank (``n_node == nloc``) and multi-rank (``n_node == nall_real``, + extended region) graph routes, so multi-rank must equal single-rank on the + excluded model. A regression that skipped the seam on the extended-region + path -- or fed it the wrong (local vs extended) atypes -- would diverge here. + + Two checks: + 1. MP (``-n 2``) ≡ SP (``-n 1``) on the excluded archive. + 2. The excluded run differs from the SAME-weights no-exclusion baseline + (``deeppot_dpa1_pairexcl_none.pt2``), so a silently-dropped exclusion + on BOTH ranks cannot pass check 1 trivially. + """ + out_mpi = _run_mpi_subprocess(nprocs=2, pb=pb_file_pairexcl) + out_ref = _run_mpi_subprocess(nprocs=1, pb=pb_file_pairexcl) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=0 + ) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + + # Exclusion must be ACTIVE on the multi-rank path: the excluded energy must + # differ from the same-weights no-exclusion baseline (O-H pairs dropped). + out_none = _run_mpi_subprocess(nprocs=2, pb=pb_file_pairexcl_none) + assert abs(out_mpi["pe"] - out_none["pe"]) > 1e-6, ( + "pair_exclude_types had no effect on the multi-rank graph path " + f"(|E_excl - E_none| = {abs(out_mpi['pe'] - out_none['pe']):.2e})" + ) From 0ebe3e23ce12dd5467daaacad6698aa19b9ad195 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 8 Jul 2026 17:55:43 +0800 Subject: [PATCH 43/63] test(lammps): DPA3 dense multi-rank pair-exclusion cell-3 test Add deeppot_dpa3_pairexcl_mpi.pt2 (gen_dpa3.py, same weights as deeppot_dpa3_mpi.pt2 + pair_exclude_types=[[0,1]]) and an MP==SP + active-vs- baseline test exercising the run_model_with_comm pair-exclusion seam (cell 3). To be validated on a clean env; local build_cpu MPI is inconsistent. --- source/lmp/tests/test_lammps_dpa3_pt2.py | 54 ++++++++++++++++++++++++ source/tests/infer/gen_dpa3.py | 27 ++++++++++++ 2 files changed, 81 insertions(+) diff --git a/source/lmp/tests/test_lammps_dpa3_pt2.py b/source/lmp/tests/test_lammps_dpa3_pt2.py index a55d80561c..b359ba5028 100644 --- a/source/lmp/tests/test_lammps_dpa3_pt2.py +++ b/source/lmp/tests/test_lammps_dpa3_pt2.py @@ -35,6 +35,16 @@ pb_file_mpi = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa3_mpi.pt2" ) +# Same as deeppot_dpa3_mpi.pt2 but with model-level pair_exclude_types=[[0,1]] +# (identical weights). Used to check that model-level exclusion survives the +# dense multi-rank (run_model_with_comm) path; deeppot_dpa3_mpi.pt2 is its +# no-exclusion baseline. Produced by source/tests/infer/gen_dpa3.py. +pb_file_pairexcl_mpi = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa3_pairexcl_mpi.pt2" +) ref_file = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa3.expected" ) @@ -540,6 +550,50 @@ def test_pair_deepmd_mpi_dpa3() -> None: ) +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +@pytest.mark.skipif( + not pb_file_pairexcl_mpi.exists(), + reason="gen_dpa3.py pair-exclude .pt2 fixture not generated", +) +def test_pair_deepmd_mpi_dpa3_pairexcl_matches_single_rank() -> None: + """Model-level ``pair_exclude_types`` must survive the dense multi-rank + (with-comm) path (cell 3). + + DPA3 is message-passing, so ``use_loc_mapping=False`` routes multi-rank + through ``run_model_with_comm`` (the dense with-comm lower). Decision + #18/A4 removed exclusion from the exported lower, so it must be applied at + the C++ ingestion seam (``applyPairExclusionNlist``) on THIS path too -- + otherwise a message-passing model with ``pair_exclude_types`` silently + includes excluded pairs multi-rank (multi-rank != single-rank). + + Two checks: + 1. MP (``-n 2``) ≡ SP (``-n 1``) on the excluded archive. + 2. The excluded run differs from the SAME-weights no-exclusion baseline + (``deeppot_dpa3_mpi.pt2``), so a silently-dropped exclusion on BOTH + ranks cannot pass check 1 trivially. + """ + out_mpi = _run_mpi_subprocess(nprocs=2, pb_path=pb_file_pairexcl_mpi) + out_ref = _run_mpi_subprocess(nprocs=1, pb_path=pb_file_pairexcl_mpi) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=0 + ) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + + # Exclusion must be ACTIVE on the with-comm path: energy must differ from + # the same-weights no-exclusion baseline (O-H pairs dropped). + out_none = _run_mpi_subprocess(nprocs=2, pb_path=pb_file_mpi) + assert abs(out_mpi["pe"] - out_none["pe"]) > 1e-6, ( + "pair_exclude_types had no effect on the dense multi-rank path " + f"(|E_excl - E_none| = {abs(out_mpi['pe'] - out_none['pe']):.2e})" + ) + + @pytest.mark.skipif( shutil.which("mpirun") is None, reason="MPI is not installed on this system" ) diff --git a/source/tests/infer/gen_dpa3.py b/source/tests/infer/gen_dpa3.py index c0a4434e33..3bbf7fbc8b 100644 --- a/source/tests/infer/gen_dpa3.py +++ b/source/tests/infer/gen_dpa3.py @@ -111,6 +111,33 @@ def main(): pt2_mpi_path, copy.deepcopy(data_mpi), do_atomic_virial=True ) + # Multi-rank variant WITH model-level ``pair_exclude_types`` — same weights + # as ``deeppot_dpa3_mpi.pt2`` above (pair_exclude_types is a build-time + # nlist mask, NOT a descriptor weight), so that model is the exact + # no-exclusion baseline. Exercises the C++ dense multi-rank + # (``run_model_with_comm``) pair-exclusion seam: DPA3 is message-passing, so + # use_loc_mapping=False routes multi-rank through ``run_model_with_comm``, + # where model-level exclusion must be applied at the C++ ingestion seam + # (the exported lower no longer bakes it in; decision #18/A4). + # See test_lammps_dpa3_pt2.py::test_pair_deepmd_mpi_dpa3_pairexcl_*. + config_pairexcl_mpi = copy.deepcopy(config_mpi) + config_pairexcl_mpi["pair_exclude_types"] = [[0, 1]] + model_pairexcl_mpi = get_model(copy.deepcopy(config_pairexcl_mpi)) + data_pairexcl_mpi = { + "model": model_pairexcl_mpi.serialize(), + "model_def_script": config_pairexcl_mpi, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + pt2_pairexcl_mpi_path = os.path.join(base_dir, "deeppot_dpa3_pairexcl_mpi.pt2") + print(f"Exporting to {pt2_pairexcl_mpi_path} ...") # noqa: T201 + pt_expt_deserialize_to_file( + pt2_pairexcl_mpi_path, + copy.deepcopy(data_pairexcl_mpi), + do_atomic_virial=True, + ) + # Float32 multi-rank variant — same architecture as the float64 # MPI fixture but with ``precision: float32``. Used by # source/lmp/tests/test_lammps_dpa3_pt2_fp32.py to validate that From b09fa836e74954739a7f502d983968c21b142149 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 8 Jul 2026 09:56:45 +0000 Subject: [PATCH 44/63] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_cc/src/DeepPotPTExpt.cc | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index f6f73ae95c..53276a0360 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -227,8 +227,7 @@ void DeepPotPTExpt::init(const std::string& model, device = torch::Device(torch::kCPU); } pair_exclude_table_ = - torch::from_blob(tbl.data(), - {static_cast(tbl.size())}, + torch::from_blob(tbl.data(), {static_cast(tbl.size())}, torch::TensorOptions().dtype(torch::kInt32)) .clone() .to(device); @@ -868,8 +867,8 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = run_model_with_comm( - coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, - fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); + coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, fparam_tensor, + aparam_tensor, charge_spin_tensor, comm_tensors); } } else { if (lower_input_is_edge_) { From 5ad1f051c5ccc248fe52345ea04b0cfe6faca389 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 8 Jul 2026 19:47:56 +0800 Subject: [PATCH 45/63] feat(dpmodel): eager fail-safe guard at the exclusion seam (fail-closed) Address iProzd's architectural review on #5733: decision #18/A4 moved pair_exclude_types to the build seam and dropped the consume-time backstop in forward_common_atomic{,_graph}, making exclusion fail-OPEN -- any ingestion path that skips the build seam would silently INCLUDE excluded pairs (the dangerous direction for an exclusion feature). Add _assert_nlist_pair_excluded / _assert_graph_pair_excluded: when pair_excl is set, verify no real neighbour/active edge carries an excluded type pair (a leak = the build seam was skipped) and raise a clear AssertionError otherwise. numpy-eager ONLY -- gated on array_api_compat.is_numpy_array, so it is a no-op under torch.export / jax jit (data-dependent, can't be traced) and in compiled production; the exported-.pt2 / C++ ingestion paths are covered by their own ingestion-site regression tests. Verified pt_expt make_fx/export still traces (guard skipped for torch tensors). Rewrite the two seam-contract tests (test_graph_atomic_parity) to assert the guard REJECTS a non-excluded input (the previously-untested contract boundary), keeping the build-time positive control; fix test_dp_atomic_model test_self_consistency to feed a pre-excluded nlist (matches test_excl_consistency). 641 common/dpmodel tests pass. --- .../dpmodel/atomic_model/base_atomic_model.py | 65 ++++++++++++++++- .../common/dpmodel/test_dp_atomic_model.py | 11 ++- .../dpmodel/test_graph_atomic_parity.py | 73 +++++++++---------- 3 files changed, 107 insertions(+), 42 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 2a65906fcc..b16dd31a05 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -301,7 +301,10 @@ def forward_common_atomic( # nlist-BUILD transform (decision #18/A4, same as the graph route): # already folded into the nlist by the NeighborList builders (Python) # or ``applyPairExclusionNlist`` (C++ ingestion); this method consumes - # a pre-excluded nlist. + # a pre-excluded nlist. Fail-safe (eager only): guard against a caller + # that skipped the build seam, which would silently INCLUDE excluded + # pairs (fail-open). + self._assert_nlist_pair_excluded(nlist, extended_atype) ext_atom_mask = self.make_atom_mask(extended_atype) ret_dict = self.forward_atomic( @@ -364,6 +367,9 @@ def forward_common_atomic_graph( # graph-BUILD transform (decision #18) already folded into # ``graph.edge_mask`` by the NeighborGraph builder (Python) or # ``applyPairExclusion`` (C++); this method consumes a pre-excluded graph. + # Fail-safe (eager only): guard against a caller that skipped the build + # seam, which would silently INCLUDE excluded pairs (fail-open). + self._assert_graph_pair_excluded(graph, atype_clamped) ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -373,6 +379,63 @@ def forward_common_atomic_graph( ) return self._finalize_atomic_ret(ret_dict, atom_mask, atype) + def _assert_nlist_pair_excluded(self, nlist: Array, extended_atype: Array) -> None: + """Fail-safe: assert the nlist reaching the dense seam is pre-excluded. + + Decision #18/A4 makes the build seam the SOLE owner of model-level + ``pair_exclude_types``; this method no longer re-applies it. A caller + that skips the build seam would therefore silently INCLUDE excluded + pairs (fail-open) -- the dangerous direction for an exclusion feature. + This guard turns that into a loud error. + + Eager (numpy) only: it is a data-dependent check, so it is skipped under + ``torch.export`` / jax ``jit`` (where it cannot be traced) and in + compiled production. The C++ / exported-``.pt2`` ingestion paths are + covered by their own ingestion-site regression tests instead. + """ + if self.pair_excl is None or not array_api_compat.is_numpy_array(nlist): + return + xp = array_api_compat.array_namespace(nlist) + # keep == 0 marks an excluded type pair; a pre-excluded nlist has already + # set those neighbours to -1, so any *real* (>= 0) neighbour with keep==0 + # is a leak (the build seam was skipped). + keep = self.pair_excl.build_type_exclude_mask(nlist, extended_atype) + leaked = xp.astype(nlist != -1, xp.bool) & (keep == 0) + if bool(xp.any(leaked)): + n_leak = int(xp.sum(xp.astype(leaked, xp.int64))) + raise AssertionError( + f"forward_common_atomic received a nlist that is NOT " + f"pair-excluded: {n_leak} excluded-type neighbour(s) still " + "present. Model-level pair_exclude_types is a nlist-BUILD " + "transform (decision #18/A4) -- apply it at neighbor-list build " + "(Python builders / C++ applyPairExclusionNlist); this seam does " + "not re-apply it." + ) + + def _assert_graph_pair_excluded(self, graph: "NeighborGraph", atype: Array) -> None: + """Fail-safe graph analogue of :meth:`_assert_nlist_pair_excluded`. + + A pre-excluded graph has ``edge_mask == False`` on every excluded edge, + so any *active* edge whose type pair is excluded is a leak. Eager + (numpy) only, for the same reasons. + """ + if self.pair_excl is None or not array_api_compat.is_numpy_array( + graph.edge_mask + ): + return + xp = array_api_compat.array_namespace(graph.edge_mask) + keep = self.pair_excl.build_edge_exclude_mask(graph.edge_index, atype) + leaked = xp.astype(graph.edge_mask, xp.bool) & (keep == 0) + if bool(xp.any(leaked)): + n_leak = int(xp.sum(xp.astype(leaked, xp.int64))) + raise AssertionError( + f"forward_common_atomic_graph received a graph that is NOT " + f"pair-excluded: {n_leak} excluded-type edge(s) still active. " + "Model-level pair_exclude_types is a graph-BUILD transform " + "(decision #18) -- apply it at graph build (Python builders / " + "C++ applyPairExclusion); this seam does not re-apply it." + ) + def _finalize_atomic_ret( self, ret_dict: dict, atom_mask: Array, atype: Array ) -> dict: diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index e4592c2087..721a4f5895 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -79,8 +79,15 @@ def test_self_consistency( md0.reinit_pair_exclude(pair_excl) md1 = DPAtomicModel.deserialize(md0.serialize()) - ret0 = md0.forward_common_atomic(self.coord_ext, self.atype_ext, self.nlist) - ret1 = md1.forward_common_atomic(self.coord_ext, self.atype_ext, self.nlist) + # forward_common_atomic consumes a pre-excluded nlist (decision + # #18/A4). md0 and md1 share the same pair_excl (deserialized), so + # fold it into the nlist once; the serialize/deserialize consistency + # check is unaffected by which (identical) nlist both consume. + nlist = self.nlist + if md0.pair_excl is not None: + nlist = apply_pair_exclusion_nlist(nlist, self.atype_ext, md0.pair_excl) + ret0 = md0.forward_common_atomic(self.coord_ext, self.atype_ext, nlist) + ret1 = md1.forward_common_atomic(self.coord_ext, self.atype_ext, nlist) np.testing.assert_allclose(ret0["energy"], ret1["energy"]) diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index fa5e09e163..34b5c27b04 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -168,11 +168,13 @@ def test_model_pair_exclude_types_graph_matches_dense(): def test_model_pair_exclude_applied_at_build_not_in_lower(): - """Seam contract (decision #18): model-level pair_exclude is a graph-BUILD - transform. The graph lower must NOT re-apply it — it consumes whatever - ``edge_mask`` the builder produced. Feeding a NON-excluded graph to the lower - on a model that HAS ``pair_exclude_types`` yields the SAME result as a model - with no exclusion; the exclusion only takes effect when applied at build. + """Seam contract + fail-safe (decision #18): model-level pair_exclude is a + graph-BUILD transform; the graph lower does NOT re-apply it. Because the + consume-time backstop was removed, the lower would otherwise be fail-OPEN + (a non-excluded input silently INCLUDES excluded pairs). To keep that from + being silent, the lower guards its input (eager): feeding a NON-excluded + graph to a model that HAS ``pair_exclude_types`` RAISES. Applying exclusion + at BUILD is the correct path and changes the lower's output. """ rng = np.random.default_rng(4) nloc = 6 @@ -193,21 +195,13 @@ def test_model_pair_exclude_applied_at_build_not_in_lower(): "edge_vec": ng_raw.edge_vec, "edge_mask": ng_raw.edge_mask, } - out_with_excl_model = model.call_lower_graph(**kw) - # clear the model-level exclusion; the lower output must be UNCHANGED, proving - # the lower never consulted ``pair_excl`` (exclusion is not applied here). - model.atomic_model.reinit_pair_exclude([]) - assert model.atomic_model.pair_excl is None - out_no_excl_model = model.call_lower_graph(**kw) - np.testing.assert_allclose( - np.asarray(out_with_excl_model["energy_redu"]), - np.asarray(out_no_excl_model["energy_redu"]), - rtol=1e-12, - atol=1e-12, - ) + # Fail-safe: the lower refuses a non-excluded graph (contract boundary) — + # the excluded (0, 1) pairs are within rcut here, so at least one leaks. + with pytest.raises(AssertionError, match="NOT pair-excluded"): + model.call_lower_graph(**kw) - # Positive control: applying exclusion at BUILD (excluded edge_mask) DOES - # change the same lower's output. + # Positive control: applying exclusion at BUILD (excluded edge_mask) passes + # the guard AND changes the output vs the same lower with no exclusion. ng_excl = build_neighbor_graph( coord, atype, box, model.get_rcut(), pair_excl=PairExcludeMask(2, [(0, 1)]) ) @@ -218,6 +212,11 @@ def test_model_pair_exclude_applied_at_build_not_in_lower(): edge_vec=ng_excl.edge_vec, edge_mask=ng_excl.edge_mask, ) + # No-exclusion reference: clearing pair_excl makes the raw graph valid again + # (guard skipped when pair_excl is None), so the lower runs on it. + model.atomic_model.reinit_pair_exclude([]) + assert model.atomic_model.pair_excl is None + out_no_excl_model = model.call_lower_graph(**kw) assert not np.allclose( np.asarray(out_built_excl["energy_redu"]), np.asarray(out_no_excl_model["energy_redu"]), @@ -227,12 +226,13 @@ def test_model_pair_exclude_applied_at_build_not_in_lower(): def test_model_pair_exclude_applied_at_build_not_in_dense_lower(): - """Dense-route seam contract (decision #18/A4, mirror of the graph test): - model-level pair_exclude is a nlist-BUILD transform. The dense lower - (``call_lower``) must NOT re-apply it — it consumes whatever nlist the - builder produced. Feeding a RAW nlist to the lower on a model that HAS - ``pair_exclude_types`` yields the SAME result as a model with no exclusion; - the exclusion only takes effect when folded in at build. + """Dense-route seam contract + fail-safe (decision #18/A4, mirror of the + graph test): model-level pair_exclude is a nlist-BUILD transform; the dense + lower (``call_lower``) does NOT re-apply it. With the consume-time backstop + removed, a non-excluded nlist would silently INCLUDE excluded pairs, so the + lower guards its input (eager): feeding a RAW nlist to a model that HAS + ``pair_exclude_types`` RAISES. Folding exclusion in at BUILD is correct and + changes the output. """ from deepmd.dpmodel.utils.nlist import ( apply_pair_exclusion_nlist, @@ -258,25 +258,20 @@ def test_model_pair_exclude_applied_at_build_not_in_dense_lower(): mixed_types=True, box=box, ) - # EnergyModel.call_lower translates keys: reduced energy -> "energy" - out_with_excl_model = model.call_lower(coord_ext, atype_ext, nlist, mapping) - # clear the model-level exclusion; the lower output must be UNCHANGED, - # proving the lower never consulted ``pair_excl``. - model.atomic_model.reinit_pair_exclude([]) - assert model.atomic_model.pair_excl is None - out_no_excl_model = model.call_lower(coord_ext, atype_ext, nlist, mapping) - np.testing.assert_allclose( - np.asarray(out_with_excl_model["energy"]), - np.asarray(out_no_excl_model["energy"]), - rtol=1e-12, - atol=1e-12, - ) + # Fail-safe: the lower refuses a non-excluded nlist (contract boundary). + with pytest.raises(AssertionError, match="NOT pair-excluded"): + model.call_lower(coord_ext, atype_ext, nlist, mapping) - # Positive control: folding the exclusion in at BUILD changes the output. + # Positive control: folding the exclusion in at BUILD passes the guard AND + # changes the output vs the same lower with no exclusion. nlist_excl = apply_pair_exclusion_nlist( nlist, atype_ext, PairExcludeMask(2, [(0, 1)]) ) out_built_excl = model.call_lower(coord_ext, atype_ext, nlist_excl, mapping) + # No-exclusion reference: clearing pair_excl makes the raw nlist valid again. + model.atomic_model.reinit_pair_exclude([]) + assert model.atomic_model.pair_excl is None + out_no_excl_model = model.call_lower(coord_ext, atype_ext, nlist, mapping) assert not np.allclose( np.asarray(out_built_excl["energy"]), np.asarray(out_no_excl_model["energy"]), From 5c5889c212b4a1d135e6808a49e6cd98901138f2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 8 Jul 2026 19:52:59 +0800 Subject: [PATCH 46/63] fix(jax): apply pair exclusion at the jax-md call_lower ingestion seam Close the jax-md leak iProzd flagged: _eval_with_jax_md_neighbor built the lower nlist from the JAX-MD neighbor list without model-level pair_exclude_types, then called call_lower -- which no longer re-applies it (decision #18/A4) -- so excluded pairs were silently included (fail-open). Fold apply_pair_exclusion_nlist in at this ingestion seam, mirroring DeepEval's nlist path. Guarded on the model actually carrying pair_excl, so the mock-model jax-md tests are unaffected. Test (runs without jax_md via the DenseNeighbor path): an excluded (0,1) pair drops the only edge -> energy 0 vs 0.04 without exclusion. --- deepmd/jax/jax_md/__init__.py | 12 ++++++++++++ source/tests/jax/test_jax_md.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/deepmd/jax/jax_md/__init__.py b/deepmd/jax/jax_md/__init__.py index 45080dbf8c..13f23d212b 100644 --- a/deepmd/jax/jax_md/__init__.py +++ b/deepmd/jax/jax_md/__init__.py @@ -323,6 +323,18 @@ def _eval_with_jax_md_neighbor( displacement_fn, displacement_kwargs, ) + # Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision + # #18/A4): ``call_lower`` consumes a pre-excluded nlist and no longer + # re-applies it. The JAX-MD neighbor list is built without exclusion, so + # fold it in at this ingestion seam -- otherwise excluded pairs would be + # silently included (fail-open). + pair_excl = getattr(getattr(model, "atomic_model", None), "pair_excl", None) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) return model.call_lower( extended_coord, extended_atype, diff --git a/source/tests/jax/test_jax_md.py b/source/tests/jax/test_jax_md.py index ce4b7288c6..557da898ac 100644 --- a/source/tests/jax/test_jax_md.py +++ b/source/tests/jax/test_jax_md.py @@ -109,6 +109,36 @@ def test_dense_neighbor_uses_jax_md_displacement_convention(): np.testing.assert_allclose(potential(coord, neighbor=neighbor), 0.04, atol=1e-12) +def test_dense_neighbor_applies_model_pair_exclusion(): + """Model-level ``pair_exclude_types`` must be folded into the JAX-MD nlist + at the ``call_lower`` ingestion seam (decision #18/A4): the lower no longer + re-applies it, so without this the JAX-MD path would silently INCLUDE + excluded pairs (fail-open). Types are (0, 1); excluding that pair drops the + only edge, so the energy is 0 instead of the 0.04 the same geometry gives + without exclusion. + """ + from types import ( + SimpleNamespace, + ) + + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + class ExclEdgeModel(EdgeModel): + atomic_model = SimpleNamespace(pair_excl=PairExcludeMask(2, [(0, 1)])) + + disp = lambda ra, rb: (ra - rb) - 10.0 * jnp.round((ra - rb) / 10.0) # noqa: E731 + coord = jnp.asarray([[0.1, 0.0, 0.0], [9.9, 0.0, 0.0]]) + neighbor = DenseNeighbor(jnp.asarray([[1], [0]], dtype=jnp.int32)) + + excl = energy_fn(ExclEdgeModel(), [0, 1], displacement_fn=disp) + np.testing.assert_allclose(excl(coord, neighbor=neighbor), 0.0, atol=1e-12) + # control: identical geometry WITHOUT exclusion keeps the edge (0.04). + noexcl = energy_fn(EdgeModel(), [0, 1], displacement_fn=disp) + np.testing.assert_allclose(noexcl(coord, neighbor=neighbor), 0.04, atol=1e-12) + + def test_dense_neighbor_rejects_scalar_metric(): potential = energy_fn( EdgeModel(), From a95a9581da0c9043fdc3699cd2452544af9e2619 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 9 Jul 2026 18:55:48 +0800 Subject: [PATCH 47/63] test(infer): derive deeppot_dpa3_pairexcl_mpi.pt2 by metadata patch, no recompile The pair-exclusion DPA3 MPI fixture was a second full inductor compile (~2 with- comm compiles, expensive). Model-level pair_exclude_types is applied at the C++ ingestion seam (from metadata.json) and the Python DeepEval build seam (from the serialized model.json), NOT baked into the exported graph -- so the compiled AOTI artifact (incl. the nested with-comm .pt2) is byte-identical to deeppot_dpa3_mpi.pt2. Derive the pairexcl fixture by copying that archive and patching pair_exclude_types into BOTH JSON blobs (metadata.json for C++, model.json for Python) so the two paths agree. Verified via DeepEval: the derived model is bit-identical to a full recompile (E=2.5311 vs baseline 2.4916) and differs from the no-exclusion baseline. deeppot_dpa3_mpi.pt2 stays the exact no-exclusion baseline for the reference and active-vs-baseline tests. --- source/tests/infer/gen_dpa3.py | 54 +++++++++++++++++++--------------- 1 file changed, 31 insertions(+), 23 deletions(-) diff --git a/source/tests/infer/gen_dpa3.py b/source/tests/infer/gen_dpa3.py index 3bbf7fbc8b..d2fcf7a859 100644 --- a/source/tests/infer/gen_dpa3.py +++ b/source/tests/infer/gen_dpa3.py @@ -111,32 +111,40 @@ def main(): pt2_mpi_path, copy.deepcopy(data_mpi), do_atomic_virial=True ) - # Multi-rank variant WITH model-level ``pair_exclude_types`` — same weights - # as ``deeppot_dpa3_mpi.pt2`` above (pair_exclude_types is a build-time - # nlist mask, NOT a descriptor weight), so that model is the exact - # no-exclusion baseline. Exercises the C++ dense multi-rank - # (``run_model_with_comm``) pair-exclusion seam: DPA3 is message-passing, so - # use_loc_mapping=False routes multi-rank through ``run_model_with_comm``, - # where model-level exclusion must be applied at the C++ ingestion seam - # (the exported lower no longer bakes it in; decision #18/A4). + # Multi-rank variant WITH model-level ``pair_exclude_types`` — derived + # CHEAPLY from ``deeppot_dpa3_mpi.pt2`` by patching the exclusion list into + # the archive, with NO second inductor compile. Model-level exclusion is a + # BUILD-time transform (decision #18/A4) applied at the C++ ingestion seam + # (``applyPairExclusionNlist``, from ``metadata.json``) and the Python + # DeepEval build seam (from the serialized ``model.json``); it is NOT baked + # into the exported graph, so the compiled AOTI artifact (incl. the nested + # with-comm ``.pt2``) is byte-identical to the baseline. Only the two JSON + # blobs that carry ``pair_exclude_types`` differ -- patch both so the C++ + # and Python paths agree. ``deeppot_dpa3_mpi.pt2`` is the exact no-exclusion + # baseline. (Verified bit-identical to a full recompile via DeepEval.) # See test_lammps_dpa3_pt2.py::test_pair_deepmd_mpi_dpa3_pairexcl_*. - config_pairexcl_mpi = copy.deepcopy(config_mpi) - config_pairexcl_mpi["pair_exclude_types"] = [[0, 1]] - model_pairexcl_mpi = get_model(copy.deepcopy(config_pairexcl_mpi)) - data_pairexcl_mpi = { - "model": model_pairexcl_mpi.serialize(), - "model_def_script": config_pairexcl_mpi, - "backend": "dpmodel", - "software": "deepmd-kit", - "version": "3.0.0", - } + import json + import zipfile + pt2_pairexcl_mpi_path = os.path.join(base_dir, "deeppot_dpa3_pairexcl_mpi.pt2") - print(f"Exporting to {pt2_pairexcl_mpi_path} ...") # noqa: T201 - pt_expt_deserialize_to_file( - pt2_pairexcl_mpi_path, - copy.deepcopy(data_pairexcl_mpi), - do_atomic_virial=True, + print( # noqa: T201 + f"Deriving {pt2_pairexcl_mpi_path} from {pt2_mpi_path} " + "(pair_exclude_types patch, no recompile) ..." ) + pair_exclude_types = [[0, 1]] + with zipfile.ZipFile(pt2_mpi_path) as zin: + entries = [(info, zin.read(info.filename)) for info in zin.infolist()] + with zipfile.ZipFile(pt2_pairexcl_mpi_path, "w") as zout: + for info, blob in entries: + if info.filename == "model/extra/metadata.json": + meta = json.loads(blob) + meta["pair_exclude_types"] = pair_exclude_types + blob = json.dumps(meta).encode() + elif info.filename == "model/extra/model.json": + mdl = json.loads(blob) + mdl["model"]["pair_exclude_types"] = pair_exclude_types + blob = json.dumps(mdl).encode() + zout.writestr(info, blob) # Float32 multi-rank variant — same architecture as the float64 # MPI fixture but with ``precision: float32``. Used by From 6e01b62a7712d540eb9115d3f6186b9a65ca3bd9 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 9 Jul 2026 19:45:04 +0800 Subject: [PATCH 48/63] test(infer): derive graph pair-exclude fixtures by patch, share helper Apply the metadata-patch fixture derivation (a95a9581d) to gen_dpa1_pairexcl.py and factor the logic into gen_common.derive_pair_exclude_pt2 (used by both gen_dpa3.py and gen_dpa1_pairexcl.py). gen_dpa1_pairexcl.py generated three inductor-compiled models (_none/_graph/_nlist). _graph has the same weights and lower_kind as the _none graph baseline and differs only in pair_exclude_types, so it is now DERIVED from _none by patching the archive -- no third compile. _nlist is a different lower_kind (different exported graph) so it stays a compile. 3 compiles -> 2. Verified via DeepEval: the derived _graph is bit-identical to a full recompile (E active vs the _none baseline). gen_dpa3.py refactored to use the shared helper. --- source/tests/infer/gen_common.py | 42 +++++++++++++++++++++++ source/tests/infer/gen_dpa1_pairexcl.py | 45 ++++++++++++++++++------- source/tests/infer/gen_dpa3.py | 33 ++++-------------- 3 files changed, 80 insertions(+), 40 deletions(-) diff --git a/source/tests/infer/gen_common.py b/source/tests/infer/gen_common.py index fea29fc63e..4ea27feb32 100644 --- a/source/tests/infer/gen_common.py +++ b/source/tests/infer/gen_common.py @@ -2,8 +2,50 @@ """Common utilities shared by model generation scripts (gen_*.py).""" import glob +import json import os import sys +import zipfile + + +def derive_pair_exclude_pt2(src_pt2, dst_pt2, exclude_types): + """Derive a model-level ``pair_exclude_types`` variant of a ``.pt2`` archive. + + Parameters + ---------- + src_pt2 : str + Path to the no-exclusion baseline ``.pt2``. Must have the same weights + and ``lower_kind`` as the desired variant (only the exclusion list may + differ); a graph-lower baseline cannot derive an nlist-lower variant. + dst_pt2 : str + Path to write the derived ``.pt2``. + exclude_types : list[list[int]] + The ``pair_exclude_types`` list to inject, e.g. ``[[0, 1]]``. + + Notes + ----- + Model-level ``pair_exclude_types`` is a build-time transform (decision + #18/A4) applied at the C++ ingestion seam (from ``metadata.json``) and the + Python DeepEval build seam (from the serialized ``model.json``); it is not + baked into the exported graph, so the compiled AOTI artifact (including any + nested with-comm ``.pt2``) is byte-identical to the baseline. Only the two + JSON blobs that carry the list are patched -- both, so the C++ and Python + inference paths agree; patching one alone makes them disagree. This avoids + a second inductor compile and is bit-identical to a full recompile. + """ + with zipfile.ZipFile(src_pt2) as zin: + entries = [(info, zin.read(info.filename)) for info in zin.infolist()] + with zipfile.ZipFile(dst_pt2, "w") as zout: + for info, blob in entries: + if info.filename == "model/extra/metadata.json": + meta = json.loads(blob) + meta["pair_exclude_types"] = exclude_types + blob = json.dumps(meta).encode() + elif info.filename == "model/extra/model.json": + mdl = json.loads(blob) + mdl["model"]["pair_exclude_types"] = exclude_types + blob = json.dumps(mdl).encode() + zout.writestr(info, blob) def ensure_inductor_compiler(): diff --git a/source/tests/infer/gen_dpa1_pairexcl.py b/source/tests/infer/gen_dpa1_pairexcl.py index 8c2291db4d..452a7b4e86 100644 --- a/source/tests/infer/gen_dpa1_pairexcl.py +++ b/source/tests/infer/gen_dpa1_pairexcl.py @@ -3,12 +3,18 @@ """Generate DPA1 test models carrying model-level ``pair_exclude_types``. Produces two graph-eligible DPA1(attn_layer=0) models with identical weights, -one exported through each C++ ingestion route, plus a no-exclusion baseline: +one per C++ ingestion route, plus a no-exclusion baseline: - deeppot_dpa1_pairexcl_graph.pt2 (lower_kind="graph", pair_exclude=[[0,1]]) - deeppot_dpa1_pairexcl_nlist.pt2 (lower_kind="nlist", pair_exclude=[[0,1]]) - deeppot_dpa1_pairexcl_none.pt2 (lower_kind="graph", NO exclusion) +Only two inductor compiles are needed: the graph baseline (``_none``) and the +nlist route (``_nlist``). ``_graph`` has the same weights and lower_kind as +``_none`` and differs only in the exclusion list, so it is DERIVED from +``_none`` by patching the archive (``gen_common.derive_pair_exclude_pt2``) -- +no third compile. + The pair models exercise the C++ pair-exclusion ownership (decision #18/A4: exclusion is a BUILD-time transform on BOTH routes) plus the ``pair_exclude_types`` metadata round-trip in ``DeepPotPTExpt::init``: @@ -45,6 +51,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) from gen_common import ( + derive_pair_exclude_pt2, ensure_inductor_compiler, load_custom_ops, write_expected_ref, @@ -134,19 +141,31 @@ def build_data(exclude_types): # ---- Pair-exclusion models (graph + dense routes) ---- exclude_types = [[0, 1]] - data_e = build_data(exclude_types) - for lower_kind, tag in (("graph", "graph"), ("nlist", "nlist")): + # graph route: SAME weights + graph lower as the no-exclusion baseline + # above, only the exclusion list differs -> derive by patching the archive, + # NO second inductor compile. See gen_common.derive_pair_exclude_pt2. + graph_pt2 = os.path.join(base_dir, "deeppot_dpa1_pairexcl_graph.pt2") + print( # noqa: T201 + f"Deriving {graph_pt2} from {none_pt2} " + "(pair_exclude_types patch, no recompile) ..." + ) + derive_pair_exclude_pt2(none_pt2, graph_pt2, exclude_types) + # nlist route: a DIFFERENT lower_kind (a different exported graph) that the + # graph baseline cannot supply -> compile it. + nlist_pt2 = os.path.join(base_dir, "deeppot_dpa1_pairexcl_nlist.pt2") + print( # noqa: T201 + f"Exporting to {nlist_pt2} (lower_kind='nlist', " + f"pair_exclude_types={exclude_types}) ..." + ) + pt_expt_deserialize_to_file( + nlist_pt2, + copy.deepcopy(build_data(exclude_types)), + do_atomic_virial=True, + lower_kind="nlist", + ) + # Eval + activeness check + reference sidecar for both routes. + for tag in ("graph", "nlist"): pt2_path = os.path.join(base_dir, f"deeppot_dpa1_pairexcl_{tag}.pt2") - print( # noqa: T201 - f"Exporting to {pt2_path} (lower_kind='{lower_kind}', " - f"pair_exclude_types={exclude_types}) ..." - ) - pt_expt_deserialize_to_file( - pt2_path, - copy.deepcopy(data_e), - do_atomic_virial=True, - lower_kind=lower_kind, - ) dp_e = DeepPot(pt2_path) e_e1, f_e1, v_e1, ae_e1, av_e1 = dp_e.eval(coord, box, atype, atomic=True) e_enp, f_enp, v_enp, ae_enp, av_enp = dp_e.eval(coord, None, atype, atomic=True) diff --git a/source/tests/infer/gen_dpa3.py b/source/tests/infer/gen_dpa3.py index d2fcf7a859..feb064eb37 100644 --- a/source/tests/infer/gen_dpa3.py +++ b/source/tests/infer/gen_dpa3.py @@ -17,6 +17,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) from gen_common import ( + derive_pair_exclude_pt2, ensure_inductor_compiler, load_custom_ops, write_expected_ref, @@ -113,38 +114,16 @@ def main(): # Multi-rank variant WITH model-level ``pair_exclude_types`` — derived # CHEAPLY from ``deeppot_dpa3_mpi.pt2`` by patching the exclusion list into - # the archive, with NO second inductor compile. Model-level exclusion is a - # BUILD-time transform (decision #18/A4) applied at the C++ ingestion seam - # (``applyPairExclusionNlist``, from ``metadata.json``) and the Python - # DeepEval build seam (from the serialized ``model.json``); it is NOT baked - # into the exported graph, so the compiled AOTI artifact (incl. the nested - # with-comm ``.pt2``) is byte-identical to the baseline. Only the two JSON - # blobs that carry ``pair_exclude_types`` differ -- patch both so the C++ - # and Python paths agree. ``deeppot_dpa3_mpi.pt2`` is the exact no-exclusion - # baseline. (Verified bit-identical to a full recompile via DeepEval.) - # See test_lammps_dpa3_pt2.py::test_pair_deepmd_mpi_dpa3_pairexcl_*. - import json - import zipfile - + # the archive, with NO second inductor compile (same weights + lower_kind, + # only the exclusion list differs). ``deeppot_dpa3_mpi.pt2`` is the exact + # no-exclusion baseline. See ``derive_pair_exclude_pt2`` for why this is + # sound, and test_lammps_dpa3_pt2.py::test_pair_deepmd_mpi_dpa3_pairexcl_*. pt2_pairexcl_mpi_path = os.path.join(base_dir, "deeppot_dpa3_pairexcl_mpi.pt2") print( # noqa: T201 f"Deriving {pt2_pairexcl_mpi_path} from {pt2_mpi_path} " "(pair_exclude_types patch, no recompile) ..." ) - pair_exclude_types = [[0, 1]] - with zipfile.ZipFile(pt2_mpi_path) as zin: - entries = [(info, zin.read(info.filename)) for info in zin.infolist()] - with zipfile.ZipFile(pt2_pairexcl_mpi_path, "w") as zout: - for info, blob in entries: - if info.filename == "model/extra/metadata.json": - meta = json.loads(blob) - meta["pair_exclude_types"] = pair_exclude_types - blob = json.dumps(meta).encode() - elif info.filename == "model/extra/model.json": - mdl = json.loads(blob) - mdl["model"]["pair_exclude_types"] = pair_exclude_types - blob = json.dumps(mdl).encode() - zout.writestr(info, blob) + derive_pair_exclude_pt2(pt2_mpi_path, pt2_pairexcl_mpi_path, [[0, 1]]) # Float32 multi-rank variant — same architecture as the float64 # MPI fixture but with ``precision: float32``. Used by From 9f9da9c2686a1e92c6b227782ba2c4f487756884 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 9 Jul 2026 20:13:20 +0800 Subject: [PATCH 49/63] docs: numpydoc-align docstrings added in this PR Audit of the PR's added docstrings against the package numpydoc convention (CLAUDE.md): public functions were already compliant; fix the private/helper idiom mismatches so they match their files' sibling methods. - base_atomic_model._assert_{nlist,graph}_pair_excluded: add Parameters + Raises sections (siblings like _finalize_atomic_ret use full numpydoc). - env_mat_stat._graph_env_mat: add Parameters + Returns (siblings iter/__call__ use numpydoc). - jax2tf/make_model.model_call_from_call_lower: revert docstring to its original one-liner (the function documents no other params); the pair_excl rationale is already an inline comment at the application site, so the single-param prose was redundant + inconsistent. - gen_common.derive_pair_exclude_pt2: concise summary + typed Parameters, move the decision/verification cross-refs into a Notes section (not floating prose). No Google-style/See-Also/RTD-breaking issues found. --- .../dpmodel/atomic_model/base_atomic_model.py | 29 +++++++++++++++++++ deepmd/dpmodel/utils/env_mat_stat.py | 16 ++++++++++ deepmd/jax/jax2tf/make_model.py | 8 +---- 3 files changed, 46 insertions(+), 7 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index b16dd31a05..ee406c11ae 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -392,6 +392,20 @@ def _assert_nlist_pair_excluded(self, nlist: Array, extended_atype: Array) -> No ``torch.export`` / jax ``jit`` (where it cannot be traced) and in compiled production. The C++ / exported-``.pt2`` ingestion paths are covered by their own ingestion-site regression tests instead. + + Parameters + ---------- + nlist + neighbor list, shape: nf x nloc x nsel; ``-1`` marks empty slots. + extended_atype + extended atom types, shape: nf x nall. + + Raises + ------ + AssertionError + if a real (``>= 0``) neighbour carries an excluded type pair, i.e. + the build seam was skipped. Only when ``self.pair_excl`` is set and + ``nlist`` is a numpy array. """ if self.pair_excl is None or not array_api_compat.is_numpy_array(nlist): return @@ -418,6 +432,21 @@ def _assert_graph_pair_excluded(self, graph: "NeighborGraph", atype: Array) -> N A pre-excluded graph has ``edge_mask == False`` on every excluded edge, so any *active* edge whose type pair is excluded is a leak. Eager (numpy) only, for the same reasons. + + Parameters + ---------- + graph + neighbor graph reaching the graph seam; only ``edge_mask`` and + ``edge_index`` are inspected. + atype + flat local node types, shape: N (clamped ``>= 0``). + + Raises + ------ + AssertionError + if an active edge (``edge_mask`` True) carries an excluded type + pair. Only when ``self.pair_excl`` is set and ``graph.edge_mask`` + is a numpy array. """ if self.pair_excl is None or not array_api_compat.is_numpy_array( graph.edge_mask diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 2fde517421..543b918b42 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -182,6 +182,22 @@ def _graph_env_mat( ``-1`` in the pre-excluded ``nlist``) carry ``edge_mask=False`` and are zeroed -- so the ``(E, 4)`` output reshapes 1:1 back to the dense ``(nf, nloc, nsel, 4)`` env-matrix tensor. + + Parameters + ---------- + extended_coord + extended coordinates, shape: nf x (nall x 3). + extended_atype + extended atom types, shape: nf x nall. + mapping + extended-to-local index mapping, shape: nf x nall. + nlist + pre-excluded neighbor list, shape: nf x nloc x nsel. + + Returns + ------- + env_mat + the environment matrix, shape: nf x nloc x nsel x last_dim. """ xp = array_api_compat.array_namespace(extended_coord, nlist) dev = array_api_compat.device(extended_coord) diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index dbd521dad2..3d569b6a71 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -62,13 +62,7 @@ def model_call_from_call_lower( do_atomic_virial: bool = False, pair_excl: "PairExcludeMask | None" = None, ) -> dict[str, tf.Tensor]: - """Return model prediction from lower interface. - - ``pair_excl`` is the model-level pair-type exclusion mask. Exclusion is a - nlist-BUILD transform (decision #18/A4): it is folded into the nlist here, - in the traced TF wrapper, because the lower JAX model consumes a - pre-excluded nlist and never re-applies it. - """ + """Return model prediction from lower interface.""" atype_shape = tf.shape(atype) nframes, nloc = atype_shape[0], atype_shape[1] cc, bb, fp, ap = coord, box, fparam, aparam From 6535cbdb827b589ac70b0bf183dc3c57c7206fc7 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 10 Jul 2026 23:19:37 +0800 Subject: [PATCH 50/63] fix(tf2): apply model-level pair_exclude on the live/eager upper path The eager TF2 model (call_common in dp_model.py) built the nlist via model_call_from_call_lower but never passed pair_excl, so live/training inference silently included excluded type pairs (the SavedModel export already folds it in). Thread pair_excl through, matching the base dpmodel upper call and the SavedModel path (decision #18/A4). --- deepmd/tf2/model/dp_model.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/deepmd/tf2/model/dp_model.py b/deepmd/tf2/model/dp_model.py index 87a939a890..28c8531a62 100644 --- a/deepmd/tf2/model/dp_model.py +++ b/deepmd/tf2/model/dp_model.py @@ -107,6 +107,11 @@ def call_common( coord_corr_for_virial=to_tensorflow_array(coord_corr_for_virial), charge_spin=cs, neighbor_list=neighbor_list, + # Model-level pair exclusion is a nlist-BUILD transform + # (decision #18/A4): fold it into the freshly built nlist here so + # the live/eager TF2 upper path matches the SavedModel export and + # the other backends. Identity when nothing is excluded. + pair_excl=getattr(self.atomic_model, "pair_excl", None), pass_lower_kwargs=True, ) return self._output_type_cast(model_predict, input_prec) From 4c15f8c341761dabe856a4c1d7fc9de47f9e24b7 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 10 Jul 2026 23:19:38 +0800 Subject: [PATCH 51/63] test(pt_expt): exercise graph route in dpa1 exportable exclude test test_exportable exported without mapping, so exclude_types only covered the dense fallback, not the graph-native apply_pair_exclusion path. Export both routes (with and without mapping) so the exclusion is traced on the graph lower too (CodeRabbit review). --- source/tests/pt_expt/descriptor/test_dpa1.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index 9ab2cc7b6b..8fd4e77598 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -139,12 +139,18 @@ def test_exportable(self, idt, prec, excl_types) -> None: dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device) dd0.se_atten.stddev = torch.tensor(dstd, dtype=dtype, device=self.device) dd0 = dd0.eval() - inputs = ( - torch.tensor(self.coord_ext, dtype=dtype, device=self.device), - torch.tensor(self.atype_ext, dtype=int, device=self.device), - torch.tensor(self.nlist, dtype=int, device=self.device), - ) - torch.export.export(dd0, inputs) + coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device) + atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device) + nlist = torch.tensor(self.nlist, dtype=int, device=self.device) + # Dense route: no mapping -> legacy dense exclusion mask. + torch.export.export(dd0, (coord_ext, atype_ext, nlist)) + # Graph route: passing ``mapping`` selects the graph lower, so the + # exclude_types case exercises the graph-native ``apply_pair_exclusion`` + # (``build_edge_exclude_mask``) path, not just the dense fallback. The + # mixin's nall(4) > nloc(3) requires a real mapping (identity would + # index out of range). + mapping = torch.tensor(self.mapping, dtype=int, device=self.device) + torch.export.export(dd0, (coord_ext, atype_ext, nlist, mapping)) @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision def test_compressed_forward(self, prec) -> None: From 33318f4c63cfe4ed8530de6e9cca6db33ff5d87b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 10 Jul 2026 23:20:32 +0800 Subject: [PATCH 52/63] feat(jax): apply model-level pair_exclude at the DeepPotJAX nlist seam C++ DeepPotJAX (which consumes both jax and tf2 .savedmodel via LAMMPS InputNlist) passed the raw nlist to the exported call_lower_* with no exclusion, so message-passing/direct-nlist inference silently included excluded type pairs (fail-open). The exported call_lower_* consumes a pre-excluded nlist by design (decision #18/A4), so fold exclusion in at the ingestion seam: - export a flat get_pair_exclude_types getter from both jax2tf and tf2 serialization (models exported before it baked exclusion into the graph, so a missing getter => identity is correct); - DeepPotJAX reads it at init and builds the keep table once; - filter the LAMMPS nlist vector (torch-free twin of applyPairExclusionNlist) before tensor construction. Hoist buildPairExcludeTable to the torch-free common.h so the libtorch apply-helpers (commonPT.h) and the TF-C-API DeepPotJAX share one builder. --- deepmd/jax/jax2tf/serialization.py | 15 ++++++++++ deepmd/tf2/utils/serialization.py | 14 +++++++++ source/api_cc/include/DeepPotJAX.h | 5 ++++ source/api_cc/include/common.h | 48 ++++++++++++++++++++++++++++++ source/api_cc/include/commonPT.h | 41 ------------------------- source/api_cc/src/DeepPotJAX.cc | 44 +++++++++++++++++++++++++++ 6 files changed, 126 insertions(+), 41 deletions(-) diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index 61db06a096..ae3901be5c 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -294,6 +294,21 @@ def get_sel() -> tf.Tensor: tf_model.get_sel = get_sel + @tf.function + def get_pair_exclude_types() -> tf.Tensor: + # Model-level pair_exclude_types, exported FLAT [ti0, tj0, ti1, tj1, + # ...] so the C++ ingestion seam (DeepPotJAX) can fold exclusion into + # the LAMMPS nlist before the traced call_lower_* consumes it + # (decision #18/A4). The upper ``call`` already pre-excludes its + # freshly built nlist. Guard atomic_model: test doubles may lack it. + pet = getattr( + getattr(model, "atomic_model", None), "pair_exclude_types", [] + ) + flat = [int(t) for pair in (pet or []) for t in pair] + return tf.constant(flat, dtype=tf.int64) + + tf_model.get_pair_exclude_types = get_pair_exclude_types + @tf.function def get_model_def_script() -> tf.Tensor: return tf.constant( diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index 431d87cc47..e4b0bb5549 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -509,6 +509,20 @@ def get_sel() -> tf.Tensor: tf_model.get_sel = get_sel + @tf.function + def get_pair_exclude_types() -> tf.Tensor: + # Model-level pair_exclude_types, exported FLAT [ti0, tj0, ti1, tj1, ...] + # so the C++ ingestion seam (DeepPotJAX, which also consumes the tf2 + # ``.savedmodel``) can fold exclusion into the LAMMPS nlist before the + # traced call_lower_* consumes it (decision #18/A4). The compiled ``call`` + # already pre-excludes its freshly built nlist. Guard atomic_model: test + # doubles may lack it. + pet = getattr(getattr(model, "atomic_model", None), "pair_exclude_types", []) + flat = [int(t) for pair in (pet or []) for t in pair] + return tf.constant(flat, dtype=tf.int64) + + tf_model.get_pair_exclude_types = get_pair_exclude_types + @tf.function def get_model_def_script() -> tf.Tensor: return tf.constant( diff --git a/source/api_cc/include/DeepPotJAX.h b/source/api_cc/include/DeepPotJAX.h index 9469118968..af6ea3ada8 100644 --- a/source/api_cc/include/DeepPotJAX.h +++ b/source/api_cc/include/DeepPotJAX.h @@ -201,6 +201,11 @@ class DeepPotJAX : public DeepPotBackend { bool do_message_passing; // has default fparam bool has_default_fparam_; + // Model-level pair-exclusion keep table (flat (ntypes+1)^2), built once in + // init from the exported ``get_pair_exclude_types``. Empty => no exclusion. + // The exported ``call_lower_*`` consumes a pre-excluded nlist (decision + // #18/A4), so this ingestion seam folds exclusion into the LAMMPS nlist. + std::vector pair_exclude_table_; // whether SavedModel execution goes through XLA and benefits from shape // padding; true for JAX/jax2tf XlaCallModule and TF2 jit_compile exports bool uses_xla_compilation_ = false; diff --git a/source/api_cc/include/common.h b/source/api_cc/include/common.h index 42bfb8a670..9ef7a54752 100644 --- a/source/api_cc/include/common.h +++ b/source/api_cc/include/common.h @@ -2,7 +2,9 @@ #pragma once #include +#include #include +#include #include #include "AtomMap.h" @@ -272,4 +274,50 @@ void fold_back(std::vector& out, } } } + +/** + * @brief Build the flat ``(ntypes+1)^2`` pair-type keep table. + * + * Inference-path mirror of the Python ``PairExcludeMask`` constructor + * (``deepmd/dpmodel/utils/exclude_mask.py``). The table is row-major over + * ``[tj][ti]`` (flat index ``tj * (ntypes+1) + ti``); an entry is ``0`` when + * the ordered pair ``(ti, tj)`` is excluded and ``1`` otherwise. Both ``(ti, + * tj)`` and ``(tj, ti)`` are inserted into the exclude set, so the table is + * symmetric. Type ``ntypes`` is the reserved virtual-atom row/column. + * + * Returns an empty vector when ``exclude_types`` is empty, so callers can treat + * an empty table as "no exclusion" (identity) just like the Python + * ``pair_excl is None`` early-exit. + * + * This lives in the backend-agnostic ``common.h`` (not ``commonPT.h``) so both + * the libtorch apply-helpers (``applyPairExclusion*`` in ``commonPT.h``) and + * the torch-free TF-C-API ``DeepPotJAX`` ingestion seam share one canonical + * table builder. + * + * @param ntypes Number of real atom types. + * @param exclude_types List of excluded ``(ti, tj)`` type pairs. + */ +inline std::vector buildPairExcludeTable( + const int ntypes, const std::vector>& exclude_types) { + if (exclude_types.empty()) { + return {}; + } + const int n1 = ntypes + 1; + std::set> excl; + for (const auto& tt : exclude_types) { + excl.insert({tt.first, tt.second}); + excl.insert({tt.second, tt.first}); + } + // type_mask[tj][ti] == 0 iff (ti, tj) is excluded (mirrors the Python + // list comprehension in PairExcludeMask.__init__, reshape(-1)). + std::vector type_mask(static_cast(n1) * n1, 1); + for (int tj = 0; tj < n1; ++tj) { + for (int ti = 0; ti < n1; ++ti) { + if (excl.count({ti, tj})) { + type_mask[static_cast(tj) * n1 + ti] = 0; + } + } + } + return type_mask; +} } // namespace deepmd diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index fd88d933ac..eec3fdfbd2 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -464,47 +464,6 @@ inline GraphTensorPack buildGraphTensors( return pack; } -/** - * @brief Build the flat ``(ntypes+1)^2`` pair-type keep table. - * - * Inference-path mirror of the Python ``PairExcludeMask`` constructor - * (``deepmd/dpmodel/utils/exclude_mask.py``). The table is row-major over - * ``[tj][ti]`` (flat index ``tj * (ntypes+1) + ti``); an entry is ``0`` when - * the ordered pair ``(ti, tj)`` is excluded and ``1`` otherwise. Both ``(ti, - * tj)`` and ``(tj, ti)`` are inserted into the exclude set, so the table is - * symmetric. Type ``ntypes`` is the reserved virtual-atom row/column. - * - * Returns an empty vector when ``exclude_types`` is empty, so callers can treat - * an empty table as "no exclusion" (identity) just like the Python - * ``pair_excl is None`` early-exit. - * - * @param ntypes Number of real atom types. - * @param exclude_types List of excluded ``(ti, tj)`` type pairs. - */ -inline std::vector buildPairExcludeTable( - const int ntypes, const std::vector>& exclude_types) { - if (exclude_types.empty()) { - return {}; - } - const int n1 = ntypes + 1; - std::set> excl; - for (const auto& tt : exclude_types) { - excl.insert({tt.first, tt.second}); - excl.insert({tt.second, tt.first}); - } - // type_mask[tj][ti] == 0 iff (ti, tj) is excluded (mirrors the Python - // list comprehension in PairExcludeMask.__init__, reshape(-1)). - std::vector type_mask(static_cast(n1) * n1, 1); - for (int tj = 0; tj < n1; ++tj) { - for (int ti = 0; ti < n1; ++ti) { - if (excl.count({ti, tj})) { - type_mask[static_cast(tj) * n1 + ti] = 0; - } - } - } - return type_mask; -} - /** * @brief Graph pair-type exclusion: AND the per-edge keep-mask into * ``edge_mask``. diff --git a/source/api_cc/src/DeepPotJAX.cc b/source/api_cc/src/DeepPotJAX.cc index 676c67a55b..ebf92b8519 100644 --- a/source/api_cc/src/DeepPotJAX.cc +++ b/source/api_cc/src/DeepPotJAX.cc @@ -15,6 +15,7 @@ #include #include #include +#include #include #include "common.h" @@ -600,6 +601,23 @@ void deepmd::DeepPotJAX::init(const std::string& model, } catch (tf_function_not_found& e) { has_default_fparam_ = false; } + try { + // Model-level pair_exclude_types, exported flat [ti0, tj0, ti1, tj1, ...]. + // Fold exclusion into the LAMMPS nlist at ingestion (decision #18/A4); the + // exported call_lower_* consumes a pre-excluded nlist. Models exported + // before this getter existed baked exclusion into the graph, so a missing + // getter (=> empty table => identity here) is correct for them. + std::vector pet_flat = get_vector( + ctx, "get_pair_exclude_types", func_vector, device, status); + std::vector> pair_exclude_types; + for (size_t ii = 0; ii + 1 < pet_flat.size(); ii += 2) { + pair_exclude_types.emplace_back(static_cast(pet_flat[ii]), + static_cast(pet_flat[ii + 1])); + } + pair_exclude_table_ = buildPairExcludeTable(ntypes, pair_exclude_types); + } catch (tf_function_not_found& e) { + pair_exclude_table_.clear(); + } inited = true; } @@ -865,6 +883,32 @@ void deepmd::DeepPotJAX::compute(std::vector& ener, } } } + // Model-level pair exclusion (decision #18/A4): erase excluded-type + // neighbours to -1 before the exported call_lower_* consumes the nlist. This + // is the torch-free twin of applyPairExclusionNlist (commonPT.h), operating + // on the plain C++ nlist vector because DeepPotJAX uses the TF C API. Empty + // table => identity. center type ti = atype[ii]; neighbour type tj = + // atype[nb]; drop when keep-table[ti*(ntypes+1)+tj] == 0 (center*(n1)+nbr, + // matching applyPairExclusionNlist / buildPairExcludeTable). + if (!pair_exclude_table_.empty()) { + const int n1 = ntypes + 1; + for (int ii = 0; ii < nloc_real; ii++) { + const int ti = atype[ii]; + for (int jj = 0; jj < static_cast(max_size); jj++) { + const int64_t nb = nlist[ii * max_size + jj]; + if (nb < 0) { + continue; + } + const int tj = atype[nb]; + if (tj < 0) { + continue; // padding slot; nothing to exclude + } + if (pair_exclude_table_[static_cast(ti) * n1 + tj] == 0) { + nlist[ii * max_size + jj] = -1; + } + } + } + } input_list[2] = add_input(op, nlist, nlist_shape, data_tensor[2], status); // mapping; for now, set it to -1, assume it is not used std::vector mapping_shape = {nframes, nall_model}; From 4846ec7837df5e08f1a12e3f227dada4485d6a24 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 10 Jul 2026 23:48:50 +0800 Subject: [PATCH 53/63] refactor(api_cc): extract torch-free nlist pair-exclude filter + unit-test it Move the DeepPotJAX inline nlist exclusion loop into a shared torch-free helper applyPairExcludeNlistVec (common.h), the plain-vector twin of the libtorch applyPairExclusionNlist (commonPT.h). Add test_pair_exclude_table.cc pinning the keep-table convention and the vector filter (empty table == identity, both-direction exclusion, empty-slot/virtual-type handling). The test is torch-free so it runs in every C++ build, giving the DeepPotJAX seam logic direct coverage without a TF build. --- source/api_cc/include/common.h | 55 +++++++++++++ source/api_cc/src/DeepPotJAX.cc | 30 ++----- .../api_cc/tests/test_pair_exclude_table.cc | 81 +++++++++++++++++++ 3 files changed, 141 insertions(+), 25 deletions(-) create mode 100644 source/api_cc/tests/test_pair_exclude_table.cc diff --git a/source/api_cc/include/common.h b/source/api_cc/include/common.h index 9ef7a54752..283205a028 100644 --- a/source/api_cc/include/common.h +++ b/source/api_cc/include/common.h @@ -1,6 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #pragma once +#include #include #include #include @@ -320,4 +321,58 @@ inline std::vector buildPairExcludeTable( } return type_mask; } + +/** + * @brief Dense-nlist pair-type exclusion on a plain flat ``int64`` neighbour + * list: erase excluded-type neighbours to ``-1`` in place. + * + * Torch-free twin of ``applyPairExclusionNlist`` (``commonPT.h``, which works + * on + * ``torch::Tensor``); used by the TF-C-API ``DeepPotJAX`` ingestion seam, which + * cannot depend on libtorch. Same keep-table convention as + * ``buildPairExcludeTable`` / ``applyPairExclusionNlist``: keep index + * ``center_type * (ntypes+1) + neighbour_type``. An empty ``type_mask_table`` + * is identity (no-op), mirroring the ``pair_excl is None`` early-exit. + * + * Exclusion is a BUILD-time transform (decision #18/A4): the exported + * ``call_lower_*`` consumes a pre-excluded nlist and never re-applies it, so + * this is the single application site for the JAX/tf2 SavedModel LAMMPS path. + * + * @param nlist Flat ``nloc * max_size`` neighbour list (extended-space indices; + * ``-1`` == empty slot), modified in place. + * @param atype Extended atom types; ``atype[i] >= 0`` for real atoms. Indexed + * by the centre id ``ii < nloc`` and by each neighbour id in ``nlist``. + * @param type_mask_table Flat ``(ntypes+1)^2`` keep table from + * ``buildPairExcludeTable``. Empty => identity. + * @param ntypes Number of real atom types. + * @param nloc Number of local (centre) atoms. + * @param max_size Neighbours per centre (row stride of ``nlist``). + */ +inline void applyPairExcludeNlistVec(std::vector& nlist, + const std::vector& atype, + const std::vector& type_mask_table, + const int ntypes, + const int nloc, + const int max_size) { + if (type_mask_table.empty()) { + return; + } + const int n1 = ntypes + 1; + for (int ii = 0; ii < nloc; ++ii) { + const int ti = atype[ii]; // centre type + for (int jj = 0; jj < max_size; ++jj) { + const std::int64_t nb = nlist[static_cast(ii) * max_size + jj]; + if (nb < 0) { + continue; // empty slot + } + const int tj = atype[static_cast(nb)]; // neighbour type + if (tj < 0) { + continue; // padding atom; nothing to exclude + } + if (type_mask_table[static_cast(ti) * n1 + tj] == 0) { + nlist[static_cast(ii) * max_size + jj] = -1; + } + } + } +} } // namespace deepmd diff --git a/source/api_cc/src/DeepPotJAX.cc b/source/api_cc/src/DeepPotJAX.cc index ebf92b8519..ff1e07af4c 100644 --- a/source/api_cc/src/DeepPotJAX.cc +++ b/source/api_cc/src/DeepPotJAX.cc @@ -884,31 +884,11 @@ void deepmd::DeepPotJAX::compute(std::vector& ener, } } // Model-level pair exclusion (decision #18/A4): erase excluded-type - // neighbours to -1 before the exported call_lower_* consumes the nlist. This - // is the torch-free twin of applyPairExclusionNlist (commonPT.h), operating - // on the plain C++ nlist vector because DeepPotJAX uses the TF C API. Empty - // table => identity. center type ti = atype[ii]; neighbour type tj = - // atype[nb]; drop when keep-table[ti*(ntypes+1)+tj] == 0 (center*(n1)+nbr, - // matching applyPairExclusionNlist / buildPairExcludeTable). - if (!pair_exclude_table_.empty()) { - const int n1 = ntypes + 1; - for (int ii = 0; ii < nloc_real; ii++) { - const int ti = atype[ii]; - for (int jj = 0; jj < static_cast(max_size); jj++) { - const int64_t nb = nlist[ii * max_size + jj]; - if (nb < 0) { - continue; - } - const int tj = atype[nb]; - if (tj < 0) { - continue; // padding slot; nothing to exclude - } - if (pair_exclude_table_[static_cast(ti) * n1 + tj] == 0) { - nlist[ii * max_size + jj] = -1; - } - } - } - } + // neighbours to -1 before the exported call_lower_* consumes the nlist. + // applyPairExcludeNlistVec (common.h) is the torch-free twin of + // applyPairExclusionNlist (commonPT.h); empty table => identity. + applyPairExcludeNlistVec(nlist, atype, pair_exclude_table_, ntypes, nloc_real, + static_cast(max_size)); input_list[2] = add_input(op, nlist, nlist_shape, data_tensor[2], status); // mapping; for now, set it to -1, assume it is not used std::vector mapping_shape = {nframes, nall_model}; diff --git a/source/api_cc/tests/test_pair_exclude_table.cc b/source/api_cc/tests/test_pair_exclude_table.cc new file mode 100644 index 0000000000..329dbb417a --- /dev/null +++ b/source/api_cc/tests/test_pair_exclude_table.cc @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Torch-free unit tests for the shared pair-exclusion helpers in common.h: +// - buildPairExcludeTable (canonical keep table, also used by the libtorch +// applyPairExclusion* helpers in commonPT.h) +// - applyPairExcludeNlistVec (plain-vector nlist filter used by the TF-C-API +// DeepPotJAX ingestion seam, which cannot depend on libtorch) +// These run in every build (no TF/PT needed) and pin the exclusion convention +// (keep index center_type*(ntypes+1)+neighbour_type; empty table == identity) +// that the compiled backends rely on. +#include + +#include +#include +#include + +#include "common.h" + +namespace { + +TEST(PairExcludeTable, empty_exclude_is_identity_table) { + // No excluded pairs -> empty table -> callers treat as "no exclusion". + const auto table = deepmd::buildPairExcludeTable(2, {}); + EXPECT_TRUE(table.empty()); +} + +TEST(PairExcludeTable, symmetric_and_correct_for_one_pair) { + // ntypes=2, exclude the (0,1) pair. Table is flat (ntypes+1)^2 = 9. + const int ntypes = 2; + const int n1 = ntypes + 1; + const auto table = + deepmd::buildPairExcludeTable(ntypes, {std::make_pair(0, 1)}); + ASSERT_EQ(table.size(), static_cast(n1) * n1); + // Excluded (both directions): center 0 <-> neighbour 1. + EXPECT_EQ(table[0 * n1 + 1], 0); // center type 0, neighbour type 1 + EXPECT_EQ(table[1 * n1 + 0], 0); // center type 1, neighbour type 0 + // Kept: same-type and virtual-type (ntypes) rows/cols. + EXPECT_EQ(table[0 * n1 + 0], 1); + EXPECT_EQ(table[1 * n1 + 1], 1); + EXPECT_EQ(table[0 * n1 + 2], 1); // virtual (empty) neighbour type + EXPECT_EQ(table[2 * n1 + 1], 1); +} + +TEST(PairExcludeNlistVec, erases_excluded_neighbours_both_directions) { + const int ntypes = 2; + const auto table = + deepmd::buildPairExcludeTable(ntypes, {std::make_pair(0, 1)}); + // 4 extended atoms; 2 are local centres. Types O,H,H,O. + const std::vector atype = {0, 1, 1, 0}; + const int nloc = 2; + const int max_size = 3; + // centre 0 (type 0): neighbours 1(H),2(H),3(O) -> drop 1,2 keep 3 + // centre 1 (type 1): neighbours 0(O),2(H),3(O) -> drop 0,3 keep 2 + std::vector nlist = {1, 2, 3, 0, 2, 3}; + deepmd::applyPairExcludeNlistVec(nlist, atype, table, ntypes, nloc, max_size); + const std::vector expected = {-1, -1, 3, -1, 2, -1}; + EXPECT_EQ(nlist, expected); +} + +TEST(PairExcludeNlistVec, empty_table_is_identity) { + const std::vector atype = {0, 1, 1, 0}; + std::vector nlist = {1, 2, 3, 0, 2, 3}; + const std::vector before = nlist; + deepmd::applyPairExcludeNlistVec(nlist, atype, /*type_mask_table=*/{}, 2, + /*nloc=*/2, /*max_size=*/3); + EXPECT_EQ(nlist, before); // no-op +} + +TEST(PairExcludeNlistVec, keeps_empty_slots_and_nonexcluded_pairs) { + const int ntypes = 2; + const auto table = + deepmd::buildPairExcludeTable(ntypes, {std::make_pair(0, 1)}); + // centre 0 (type 0): neighbours 3(O), -1(empty), 3(O) -> all kept + const std::vector atype = {0, 1, 1, 0}; + std::vector nlist = {3, -1, 3}; + const std::vector before = nlist; + deepmd::applyPairExcludeNlistVec(nlist, atype, table, ntypes, /*nloc=*/1, + /*max_size=*/3); + EXPECT_EQ(nlist, before); // O-O kept, -1 stays -1 +} + +} // namespace From c5b53451a391547cdc798c2fd1a5104b4a9f164b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 00:02:57 +0800 Subject: [PATCH 54/63] test(consistent): stop skipping tf2 for pair/atom exclusion in test_ener The eager tf2 backend (#5598) skipped tf2 whenever pair_exclude_types or atom_exclude_types was non-empty, which hid that dp_model.call_common dropped pair_excl on the live path. Now that the build seam threads exclusion through the eager upper path, un-skip tf2 so the existing end-to-end consistency test exercises exclusion and guards against the drop (jax was already un-skipped and covered). Applies to both TestEner (upper) and TestEnerLower (lower). --- source/tests/consistent/model/test_ener.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index 9522122d91..b972e5e679 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -176,10 +176,11 @@ def skip_tf(self): @property def skip_tf2(self) -> bool: - return not INSTALLED_TF2 or ( - self.data["pair_exclude_types"] != [] - or self.data["atom_exclude_types"] != [] - ) + # tf2 folds model-level pair_exclude_types into the nlist at BUILD time + # (decision #18/A4): the eager upper path (dp_model.call_common) passes + # pair_excl through, so exclusion IS exercised here — do not skip it (it + # is the regression guard for that path). + return not INSTALLED_TF2 @property def skip_jax(self) -> bool: @@ -411,10 +412,9 @@ def skip_tf(self) -> bool: @property def skip_tf2(self) -> bool: - return not INSTALLED_TF2 or ( - self.data["pair_exclude_types"] != [] - or self.data["atom_exclude_types"] != [] - ) + # See TestEner.skip_tf2: tf2 supports model-level pair_exclude_types via + # the build seam, so exclusion is exercised (regression guard). + return not INSTALLED_TF2 @property def skip_jax(self) -> bool: From 4acf4de0879dec7fcd6d1e58d8699eecc17f14b4 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 00:43:51 +0800 Subject: [PATCH 55/63] test(api_cc): end-to-end DeepPotJAX InputNlist pair-exclusion gtest Add test_deeppot_sea_pairexcl_jax.cc: loads jax .savedmodel + tf2 .savedmodeltf pairexcl variants (identical weights to the baseline, exclusion injected by inject_pair_exclude.py in convert-models.sh) and drives the LAMMPS InputNlist path. Asserts (1) excluded energy differs from baseline (exclusion active, not dropped) and (2) the InputNlist/lower path agrees with the coord-only/upper path (excludes the same pairs). Guards the DeepPotJAX integration that the torch-free unit test cannot. Runtime-skips when the JAX backend isn't built, so it compiles everywhere and runs under the TF-enabled CI build. --- .../tests/test_deeppot_sea_pairexcl_jax.cc | 145 ++++++++++++++++++ source/tests/infer/convert-models.sh | 9 ++ source/tests/infer/inject_pair_exclude.py | 38 +++++ 3 files changed, 192 insertions(+) create mode 100644 source/api_cc/tests/test_deeppot_sea_pairexcl_jax.cc create mode 100644 source/tests/infer/inject_pair_exclude.py diff --git a/source/api_cc/tests/test_deeppot_sea_pairexcl_jax.cc b/source/api_cc/tests/test_deeppot_sea_pairexcl_jax.cc new file mode 100644 index 0000000000..b483a98d1d --- /dev/null +++ b/source/api_cc/tests/test_deeppot_sea_pairexcl_jax.cc @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// End-to-end test of the C++ DeepPotJAX model-level pair-exclusion ingestion +// seam on the LAMMPS InputNlist path. +// +// DeepPotJAX consumes BOTH the jax2tf ``.savedmodel`` and the tf2 +// ``.savedmodeltf`` SavedModel flavours through the TensorFlow C API. +// Model-level pair_exclude_types is a BUILD-time transform (decision #18/A4): +// the exported +// ``call_lower_*`` consumes a pre-excluded nlist and never re-applies it, so +// the C++ ingestion seam is the SINGLE application site on this path. It reads +// the exported ``get_pair_exclude_types`` getter at init and folds exclusion +// into the LAMMPS nlist before ``call_lower_*`` runs. +// +// ``deeppot_sea_pairexcl.{savedmodel,savedmodeltf}`` carry the SAME weights as +// the no-exclusion baseline ``deeppot_sea.{savedmodel,savedmodeltf}`` (derived +// by injecting only ``pair_exclude_types`` — see convert-models.sh), so the +// ONLY difference is the exclusion. Two assertions make this load-bearing: +// 1. excluded energy differs from the baseline (exclusion is genuinely active +// on the InputNlist path, not silently dropped); +// 2. the InputNlist (lower) path agrees with the coord-only (upper) path, +// which +// builds and pre-excludes its own nlist -- i.e. the seam excludes the SAME +// pairs the upper path does. +// Without the seam, (1) collapses to equality and (2) diverges. +#include + +#include +#include +#include + +#include "DeepPot.h" +#include "deeppot_universal_test_common.h" +#include "neighbor_list.h" + +namespace { +using deepmd_test::universal::Backend; +using deepmd_test::universal::backend_enabled; +using deepmd_test::universal::path_exists; + +constexpr const char* kJaxBase = "../../tests/infer/deeppot_sea.savedmodel"; +constexpr const char* kJaxExcl = + "../../tests/infer/deeppot_sea_pairexcl.savedmodel"; +constexpr const char* kTf2Base = "../../tests/infer/deeppot_sea.savedmodeltf"; +constexpr const char* kTf2Excl = + "../../tests/infer/deeppot_sea_pairexcl.savedmodeltf"; + +// Full (all-pairs) neighbour list: every atom neighbours every other atom. +std::vector> make_full_nlist(const int natoms) { + std::vector> nlist_data(natoms); + for (int ii = 0; ii < natoms; ++ii) { + for (int jj = 0; jj < natoms; ++jj) { + if (ii != jj) { + nlist_data[ii].push_back(jj); + } + } + } + return nlist_data; +} + +// Evaluate the total energy on the NoPBC LAMMPS InputNlist path (the DeepPotJAX +// ``call_lower_*`` route that folds in model-level exclusion). +double eval_input_nlist(const char* model, + const std::vector& coord, + const std::vector& atype) { + deepmd::DeepPot dp; + dp.init(model); + const int natoms = static_cast(atype.size()); + auto nlist_data = make_full_nlist(natoms); + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, ilist.data(), numneigh.data(), + firstneigh.data()); + deepmd::convert_nlist(inlist, nlist_data); + double energy = 0.0; + std::vector force, virial; + const std::vector box; // empty => NoPBC + dp.compute(energy, force, virial, coord, atype, box, 0, inlist, 0); + return energy; +} + +class TestDeepPotSeaPairExclJax : public ::testing::Test { + protected: + // Two compact O-H clusters; every atom is well within the model cutoff, so + // there are O-H pairs to exclude and the exclusion visibly changes the + // energy. + std::vector coord = { + 0.00, 0.00, 0.00, // O + 0.90, 0.00, 0.00, // H + -0.30, 0.90, 0.00, // H + 2.50, 0.00, 0.00, // O + 3.40, 0.00, 0.00, // H + 2.20, 0.90, 0.00, // H + }; + std::vector atype = {0, 1, 1, 0, 1, 1}; + + void SetUp() override { + if (!backend_enabled(Backend::JAX)) { + GTEST_SKIP() << "JAX backend support is not enabled."; + } + if (!path_exists(kJaxExcl) || !path_exists(kJaxBase)) { + GTEST_SKIP() << "jax SavedModel pairexcl fixtures are not available."; + } + } +}; + +// jax2tf .savedmodel: exclusion is applied on the InputNlist (lower) path. +TEST_F(TestDeepPotSeaPairExclJax, jax_savedmodel_inputnlist_excludes) { + const double e_base = eval_input_nlist(kJaxBase, coord, atype); + const double e_excl = eval_input_nlist(kJaxExcl, coord, atype); + EXPECT_TRUE(std::isfinite(e_base)); + EXPECT_TRUE(std::isfinite(e_excl)); + EXPECT_GT(std::fabs(e_excl - e_base), 1e-6) + << "model-level pair_exclude_types is silently dropped on the DeepPotJAX " + "InputNlist path (excluded energy == baseline)"; +} + +// The InputNlist (lower) route must exclude the SAME pairs as the coord-only +// (upper) route, which builds and pre-excludes its own nlist. +TEST_F(TestDeepPotSeaPairExclJax, jax_inputnlist_matches_coord_only) { + const double e_nlist = eval_input_nlist(kJaxExcl, coord, atype); + deepmd::DeepPot dp; + dp.init(kJaxExcl); + double e_coord = 0.0; + std::vector force, virial; + const std::vector box; // empty => NoPBC, model builds its own nlist + dp.compute(e_coord, force, virial, coord, atype, box); + EXPECT_LT(std::fabs(e_nlist - e_coord), 1e-6) + << "InputNlist-path exclusion disagrees with the coord-only path"; +} + +// tf2 .savedmodeltf goes through the SAME C++ DeepPotJAX loader/seam. +TEST_F(TestDeepPotSeaPairExclJax, tf2_savedmodeltf_inputnlist_excludes) { + if (!path_exists(kTf2Excl) || !path_exists(kTf2Base)) { + GTEST_SKIP() << "tf2 SavedModel pairexcl fixtures are not available."; + } + const double e_base = eval_input_nlist(kTf2Base, coord, atype); + const double e_excl = eval_input_nlist(kTf2Excl, coord, atype); + EXPECT_TRUE(std::isfinite(e_base)); + EXPECT_TRUE(std::isfinite(e_excl)); + EXPECT_GT(std::fabs(e_excl - e_base), 1e-6) + << "model-level pair_exclude_types is silently dropped on the DeepPotJAX " + "tf2 SavedModel InputNlist path"; +} + +} // namespace diff --git a/source/tests/infer/convert-models.sh b/source/tests/infer/convert-models.sh index f99f8d6478..f5d3204260 100755 --- a/source/tests/infer/convert-models.sh +++ b/source/tests/infer/convert-models.sh @@ -11,3 +11,12 @@ dp convert-backend ${SCRIPT_PATH}/deeppot_sea.yaml ${SCRIPT_PATH}/deeppot_sea.sa dp convert-backend ${SCRIPT_PATH}/deeppot_dpa.yaml ${SCRIPT_PATH}/deeppot_dpa.savedmodel dp convert-backend ${SCRIPT_PATH}/deeppot_sea.yaml ${SCRIPT_PATH}/deeppot_sea.savedmodeltf dp convert-backend ${SCRIPT_PATH}/deeppot_dpa.yaml ${SCRIPT_PATH}/deeppot_dpa.savedmodeltf + +# Model-level pair_exclude_types variant of deeppot_sea (identical weights, only +# the exclusion differs) for the DeepPotJAX InputNlist ingestion-seam test. The +# C++ seam must fold exclusion into the LAMMPS nlist before the exported +# call_lower_* runs (decision #18/A4). Both SavedModel flavours (jax +# .savedmodel + tf2 .savedmodeltf) are consumed by the C++ DeepPotJAX loader. +python ${SCRIPT_PATH}/inject_pair_exclude.py ${SCRIPT_PATH}/deeppot_sea.yaml ${SCRIPT_PATH}/deeppot_sea_pairexcl.yaml +dp convert-backend ${SCRIPT_PATH}/deeppot_sea_pairexcl.yaml ${SCRIPT_PATH}/deeppot_sea_pairexcl.savedmodel +dp convert-backend ${SCRIPT_PATH}/deeppot_sea_pairexcl.yaml ${SCRIPT_PATH}/deeppot_sea_pairexcl.savedmodeltf diff --git a/source/tests/infer/inject_pair_exclude.py b/source/tests/infer/inject_pair_exclude.py new file mode 100644 index 0000000000..e522fa291e --- /dev/null +++ b/source/tests/infer/inject_pair_exclude.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Inject ``pair_exclude_types: [[0, 1]]`` into a serialized model YAML. + +Used by ``convert-models.sh`` to derive a model-level ``pair_exclude_types`` +variant of an existing ``.yaml`` model for the DeepPotJAX pair-exclusion +ingestion-seam test. The injection is a byte-preserving text replacement of the +single ``pair_exclude_types: []`` line, so the derived model has WEIGHTS +IDENTICAL to the baseline -- the only difference is the exclusion, which is +exactly what the test compares (excluded vs baseline). + +Usage +----- + python inject_pair_exclude.py +""" + +import sys + +_ANCHOR = "\n pair_exclude_types: []\n" +_INJECTED = "\n pair_exclude_types:\n - - 0\n - 1\n" + + +def main() -> None: + src, dst = sys.argv[1], sys.argv[2] + with open(src) as fp: + text = fp.read() + if _ANCHOR not in text: + raise SystemExit( + f"anchor ' pair_exclude_types: []' not found in {src}; " + "cannot inject exclusion" + ) + text = text.replace(_ANCHOR, _INJECTED, 1) + with open(dst, "w") as fp: + fp.write(text) + + +if __name__ == "__main__": + main() From 79733df9471103a6c5f4933f845080bdd4665b19 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 09:48:30 +0800 Subject: [PATCH 56/63] fix(tf2): guard atomic_model when threading pair_excl in call_common The live-path pair_excl fix accessed self.atomic_model directly, which raises AttributeError on tf2 model test doubles that lack it (broke test_tf2_dp_model_call_common_uses_tf2_helper). Guard atomic_model too, matching the SavedModel-export pattern. --- deepmd/tf2/model/dp_model.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/deepmd/tf2/model/dp_model.py b/deepmd/tf2/model/dp_model.py index 28c8531a62..6276a7c5e4 100644 --- a/deepmd/tf2/model/dp_model.py +++ b/deepmd/tf2/model/dp_model.py @@ -110,8 +110,11 @@ def call_common( # Model-level pair exclusion is a nlist-BUILD transform # (decision #18/A4): fold it into the freshly built nlist here so # the live/eager TF2 upper path matches the SavedModel export and - # the other backends. Identity when nothing is excluded. - pair_excl=getattr(self.atomic_model, "pair_excl", None), + # the other backends. Identity when nothing is excluded. Guard + # atomic_model too: test doubles may lack it. + pair_excl=getattr( + getattr(self, "atomic_model", None), "pair_excl", None + ), pass_lower_kwargs=True, ) return self._output_type_cast(model_predict, input_prec) From 779b2778b59342e6d463b3a1309307fd1cbf057a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 09:48:32 +0800 Subject: [PATCH 57/63] fix(dpmodel): apply model-level pair_exclude on the returned quartet, not in build() Passing pair_excl=... to builder.build() unconditionally broke custom NeighborList strategies whose published build(coord, atype, box, rcut, sel, return_mode=...) signature has no pair_excl param (TypeError on every model). Apply apply_pair_exclusion_nlist centrally on the returned nlist instead, so custom strategies keep the published signature (njzjz-bot review). --- deepmd/dpmodel/model/make_model.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 84175549d1..21d3557d79 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -131,8 +131,19 @@ def model_call_from_call_lower( del coord, box, fparam, aparam builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() extended_coord, extended_atype, nlist, mapping = builder.build( - cc, atype, bb, rcut, sel, pair_excl=pair_excl + cc, atype, bb, rcut, sel ) + if pair_excl is not None: + # Model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): apply it centrally on the returned quartet rather than inside + # ``builder.build(...)``, so custom NeighborList strategies keep the + # published ``build(coord, atype, box, rcut, sel, return_mode=...)`` + # signature (passing ``pair_excl`` would raise TypeError on them). + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) extended_coord = extended_coord.reshape(nframes, -1, 3) if coord_corr_for_virial is not None: xp = array_api_compat.array_namespace(coord_corr_for_virial) From 9bdda5d6345b40a206c81f12d7cde6de661eee3f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 09:48:34 +0800 Subject: [PATCH 58/63] fix(tf2): fold pair_exclude into prepare_lower_inputs for compiled training The compiled training path (trainer._make_compiled_prepare_lower_batch -> prepare_lower_inputs) fed the lower directly without applying model-level pair_exclude_types, so compiled/XLA training kept excluded neighbors. Thread pair_excl into prepare_lower_inputs as the single owner (removing the now redundant application in model_call_from_call_lower), and pass model.atomic_model.pair_excl from the trainer. Add a non-vacuity regression test that the prepared nlist erases excluded cross-type pairs (njzjz-bot review). --- deepmd/tf2/make_model.py | 23 +++++++++++------ deepmd/tf2/train/trainer.py | 7 ++++++ source/tests/tf2/test_training.py | 41 +++++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/deepmd/tf2/make_model.py b/deepmd/tf2/make_model.py index e08890247e..6d835b6a27 100644 --- a/deepmd/tf2/make_model.py +++ b/deepmd/tf2/make_model.py @@ -131,14 +131,11 @@ def model_call_from_call_lower( coord_corr_for_virial=coord_corr_for_virial, charge_spin=charge_spin, neighbor_list=neighbor_list, + # Model-level pair exclusion is folded into the nlist inside + # prepare_lower_inputs (single owner), so the compiled-training prepare + # step gets the same pre-excluded nlist as this upper call. + pair_excl=pair_excl, ) - if pair_excl is not None: - # Model-level pair exclusion is a nlist-BUILD transform (decision - # #18/A4): fold it into the freshly built nlist here (already an - # ndtensorflow array, so no wrap/unwrap) via the canonical dpmodel - # helper, before the lower consumes it. Identity when nothing is - # excluded. - nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) lower_kwargs: dict[str, Any] = {"fparam": fp, "aparam": ap} if pass_lower_kwargs: if nlist_is_formatted: @@ -183,6 +180,7 @@ def prepare_lower_inputs( coord_corr_for_virial: Array | None = None, charge_spin: Array | None = None, neighbor_list: NeighborList | None = None, + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[ Array, Array, @@ -194,7 +192,14 @@ def prepare_lower_inputs( Array | None, bool, ]: - """Build lower-interface tensors outside the train-step compiler boundary.""" + """Build lower-interface tensors outside the train-step compiler boundary. + + Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision + #18/A4): when ``pair_excl`` is provided it is folded into the freshly built + nlist here, so EVERY caller (the eager/compiled upper call and the compiled + training prepare step) gets a pre-excluded nlist and the lower never + re-applies it. + """ cc = to_tensorflow_array(coord) atype = to_tensorflow_array(atype) bb = to_tensorflow_array(box) @@ -281,6 +286,8 @@ def no_pbc() -> tuple[Array, Array, Array, Array]: extended_coord_corr = None if uses_native_nlist_builder and not mixed_types: nlist = nlist_distinguish_types(nlist, extended_atype, sel) + if pair_excl is not None: + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) return ( extended_coord, extended_atype, diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index ada601b077..095b9c9e18 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -887,6 +887,13 @@ def compiled_prepare_lower_batch( fparam=fp, aparam=ap, charge_spin=cs, + # Model-level pair exclusion is a nlist-BUILD transform + # (decision #18/A4): the compiled lower consumes a pre-excluded + # nlist, so fold exclusion in here at the compiled-training + # prepare seam. Guard atomic_model for test doubles. + pair_excl=getattr( + getattr(model, "atomic_model", None), "pair_excl", None + ), ) return compiled_prepare_lower_batch diff --git a/source/tests/tf2/test_training.py b/source/tests/tf2/test_training.py index 9d7941ade1..6e693e7742 100644 --- a/source/tests/tf2/test_training.py +++ b/source/tests/tf2/test_training.py @@ -481,6 +481,47 @@ def fake_helper(**kwargs: Any) -> dict[str, Any]: assert isinstance(to_tf_tensor(captured["coord"]), tf.Tensor) +def test_prepare_lower_inputs_folds_in_pair_exclusion() -> None: + """prepare_lower_inputs pre-excludes the nlist (the compiled-training seam). + + The compiled training path (``_make_compiled_prepare_lower_batch`` -> + ``prepare_lower_inputs``) feeds the lower directly, and the lower no longer + re-applies model-level ``pair_exclude_types`` (decision #18/A4). A non-vacuous + check: a type-[0, 1] pair within the cutoff yields a real cross-type + neighbour with NO exclusion, and that neighbour must be erased to -1 once the + (0, 1) exclusion is folded in here. + """ + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + from deepmd.tf2.make_model import ( + prepare_lower_inputs, + ) + + coord = tf.constant([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]], dtype=tf.float64) + atype = tf.constant([[0, 1]], dtype=tf.int32) + kwargs = { + "rcut": 6.0, + "sel": [10, 10], + "mixed_types": False, + "coord": coord, + "atype": atype, + "box": None, + "fparam": None, + "aparam": None, + } + + nlist_base = np.asarray(prepare_lower_inputs(**kwargs)[2]) + # non-vacuous: the type-0 and type-1 atoms ARE within the cutoff, so without + # exclusion each has the other as a real (>= 0) neighbour. + assert (nlist_base >= 0).any() + + pair_excl = PairExcludeMask(2, [[0, 1]]) + nlist_excl = np.asarray(prepare_lower_inputs(pair_excl=pair_excl, **kwargs)[2]) + # the only neighbours are cross-type (0, 1) pairs, so exclusion erases them. + assert (nlist_excl == -1).all() + + def test_tf2_dp_model_call_common_lower_does_not_forward_do_deriv_c() -> None: dp_model_module = importlib.import_module("deepmd.tf2.model.dp_model") captured: dict[str, Any] = {} From 8988c05594273052655d6f473208cac0a7fb5bad Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 09:52:19 +0800 Subject: [PATCH 59/63] refactor(dpmodel): builder owns pair_exclude; pass kwarg only when set Rework 01784ee5f: applying apply_pair_exclusion_nlist centrally in make_model put the transform outside the builder, inconsistent with the graph path where build_neighbor_graph owns exclusion. Restore builder ownership and pass the pair_excl keyword only when non-None: models without exclusion keep the old published build() signature (custom strategies unaffected, njzjz-bot's compatibility case), while a legacy builder combined with a non-empty exclusion fails loudly with TypeError instead of silently including excluded pairs (fail-closed). --- deepmd/dpmodel/model/make_model.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 21d3557d79..a609ee6dcd 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -130,20 +130,22 @@ def model_call_from_call_lower( cc, bb, fp, ap = coord, box, fparam, aparam del coord, box, fparam, aparam builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() - extended_coord, extended_atype, nlist, mapping = builder.build( - cc, atype, bb, rcut, sel - ) - if pair_excl is not None: + if pair_excl is None: + # Old published build() signature: custom NeighborList strategies that + # predate the pair_excl keyword keep working for every model without + # exclusion. + extended_coord, extended_atype, nlist, mapping = builder.build( + cc, atype, bb, rcut, sel + ) + else: # Model-level pair exclusion is a nlist-BUILD transform (decision - # #18/A4): apply it centrally on the returned quartet rather than inside - # ``builder.build(...)``, so custom NeighborList strategies keep the - # published ``build(coord, atype, box, rcut, sel, return_mode=...)`` - # signature (passing ``pair_excl`` would raise TypeError on them). - from deepmd.dpmodel.utils.nlist import ( - apply_pair_exclusion_nlist, + # #18/A4): the BUILDER owns it (mirroring build_neighbor_graph on the + # graph path), so the lower always consumes a pre-excluded nlist. A + # legacy custom builder without the pair_excl keyword fails loudly here + # (TypeError) instead of silently including excluded pairs. + extended_coord, extended_atype, nlist, mapping = builder.build( + cc, atype, bb, rcut, sel, pair_excl=pair_excl ) - - nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) extended_coord = extended_coord.reshape(nframes, -1, 3) if coord_corr_for_virial is not None: xp = array_api_compat.array_namespace(coord_corr_for_virial) From a947793888cd88f879b1d02195cdc40ddbb2fcd1 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 09:55:20 +0800 Subject: [PATCH 60/63] style(dpmodel): merge pair_excl build branches via conditional kwargs --- deepmd/dpmodel/model/make_model.py | 27 +++++++++++---------------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index a609ee6dcd..f31f7e70e1 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -130,22 +130,17 @@ def model_call_from_call_lower( cc, bb, fp, ap = coord, box, fparam, aparam del coord, box, fparam, aparam builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() - if pair_excl is None: - # Old published build() signature: custom NeighborList strategies that - # predate the pair_excl keyword keep working for every model without - # exclusion. - extended_coord, extended_atype, nlist, mapping = builder.build( - cc, atype, bb, rcut, sel - ) - else: - # Model-level pair exclusion is a nlist-BUILD transform (decision - # #18/A4): the BUILDER owns it (mirroring build_neighbor_graph on the - # graph path), so the lower always consumes a pre-excluded nlist. A - # legacy custom builder without the pair_excl keyword fails loudly here - # (TypeError) instead of silently including excluded pairs. - extended_coord, extended_atype, nlist, mapping = builder.build( - cc, atype, bb, rcut, sel, pair_excl=pair_excl - ) + # Model-level pair exclusion is a nlist-BUILD transform (decision #18/A4): + # the BUILDER owns it (mirroring build_neighbor_graph on the graph path), so + # the lower always consumes a pre-excluded nlist. The keyword is passed only + # when set: custom NeighborList strategies predating it keep working for + # every model without exclusion, while a legacy builder combined with a + # non-empty exclusion fails loudly (TypeError) instead of silently + # including excluded pairs. + excl_kwargs = {} if pair_excl is None else {"pair_excl": pair_excl} + extended_coord, extended_atype, nlist, mapping = builder.build( + cc, atype, bb, rcut, sel, **excl_kwargs + ) extended_coord = extended_coord.reshape(nframes, -1, 3) if coord_corr_for_virial is not None: xp = array_api_compat.array_namespace(coord_corr_for_virial) From 4856eed5dffc7a9d1d076d71ad138989543cb6bf Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 10:02:07 +0800 Subject: [PATCH 61/63] refactor(tf2): move make_model.py under tf2/model/; builder owns pair_excl Move deepmd/tf2/make_model.py to deepmd/tf2/model/make_model.py, matching the dpmodel/pt/pt_expt/pd layout (model/make_model.py); added at top level by the eager-tf2 PR. Update all importers. Also align the pair-exclusion convention with dpmodel/pt_expt make_model: the nlist BUILDER owns exclusion. The custom-strategy branch passes pair_excl into neighbor_list.build() (keyword only when set, so legacy strategies keep working without exclusion and fail loudly with it), and the native inline builder applies apply_pair_exclusion_nlist at build time, instead of a post-hoc application at the end of prepare_lower_inputs. --- deepmd/tf2/model/dp_model.py | 6 +++--- deepmd/tf2/{ => model}/make_model.py | 14 +++++++++++--- deepmd/tf2/train/trainer.py | 2 +- deepmd/tf2/utils/serialization.py | 6 +++--- source/tests/tf2/test_training.py | 6 +++--- 5 files changed, 21 insertions(+), 13 deletions(-) rename deepmd/tf2/{ => model}/make_model.py (92%) diff --git a/deepmd/tf2/model/dp_model.py b/deepmd/tf2/model/dp_model.py index 6276a7c5e4..4b9258a05c 100644 --- a/deepmd/tf2/model/dp_model.py +++ b/deepmd/tf2/model/dp_model.py @@ -23,12 +23,12 @@ tf, xp, ) -from deepmd.tf2.make_model import ( - model_call_from_call_lower as tf2_model_call_from_call_lower, -) from deepmd.tf2.model.base_model import ( forward_common_atomic, ) +from deepmd.tf2.model.make_model import ( + model_call_from_call_lower as tf2_model_call_from_call_lower, +) from deepmd.tf2.utils.jit import ( default_jit_compile, ) diff --git a/deepmd/tf2/make_model.py b/deepmd/tf2/model/make_model.py similarity index 92% rename from deepmd/tf2/make_model.py rename to deepmd/tf2/model/make_model.py index 6d835b6a27..24c0c3ebd6 100644 --- a/deepmd/tf2/make_model.py +++ b/deepmd/tf2/model/make_model.py @@ -249,8 +249,13 @@ def no_pbc() -> tuple[Array, Array, Array, Array]: uses_native_nlist_builder = neighbor_list is None if neighbor_list is not None: + # The BUILDER owns model-level pair exclusion (same convention as + # dpmodel/pt_expt make_model). The keyword is passed only when set, so + # legacy custom strategies keep working without exclusion and fail + # loudly (TypeError) with it instead of silently including pairs. + excl_kwargs = {} if pair_excl is None else {"pair_excl": pair_excl} extended_coord, extended_atype, nlist, mapping = neighbor_list.build( - cc, atype, bb, rcut, sel + cc, atype, bb, rcut, sel, **excl_kwargs ) else: has_pbc = _box_has_pbc(bb) @@ -274,6 +279,11 @@ def no_pbc() -> tuple[Array, Array, Array, Array]: extended_atype = to_tensorflow_array(extended_atype_tensor) nlist = to_tensorflow_array(nlist_tensor) mapping = to_tensorflow_array(mapping_tensor) + if pair_excl is not None: + # Native inline builder: exclude at BUILD time, mirroring + # DefaultNeighborList.build on the dpmodel path (the custom-builder + # branch above already excluded inside build()). + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) if coord_corr is not None: coord_corr = xp.reshape(coord_corr, (nframes, nloc, 3)) @@ -286,8 +296,6 @@ def no_pbc() -> tuple[Array, Array, Array, Array]: extended_coord_corr = None if uses_native_nlist_builder and not mixed_types: nlist = nlist_distinguish_types(nlist, extended_atype, sel) - if pair_excl is not None: - nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) return ( extended_coord, extended_atype, diff --git a/deepmd/tf2/train/trainer.py b/deepmd/tf2/train/trainer.py index 095b9c9e18..f49d21b9f0 100644 --- a/deepmd/tf2/train/trainer.py +++ b/deepmd/tf2/train/trainer.py @@ -60,7 +60,7 @@ from deepmd.tf2.env import ( tf, ) -from deepmd.tf2.make_model import ( +from deepmd.tf2.model.make_model import ( prepare_lower_inputs, ) from deepmd.tf2.model.model import ( diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index e4b0bb5549..baedd5ba87 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -24,12 +24,12 @@ from deepmd.tf2.common import ( unwrap_value, ) -from deepmd.tf2.make_model import ( - model_call_from_call_lower, -) from deepmd.tf2.model.base_model import ( BaseModel, ) +from deepmd.tf2.model.make_model import ( + model_call_from_call_lower, +) from deepmd.tf2.model.model import ( get_model, ) diff --git a/source/tests/tf2/test_training.py b/source/tests/tf2/test_training.py index 6e693e7742..dc6324684c 100644 --- a/source/tests/tf2/test_training.py +++ b/source/tests/tf2/test_training.py @@ -248,7 +248,7 @@ def test_forward_common_atomic_can_skip_virial_derivative() -> None: def test_model_call_from_call_lower_uses_tf2_native_communicate( monkeypatch: pytest.MonkeyPatch, ) -> None: - make_model_module = importlib.import_module("deepmd.tf2.make_model") + make_model_module = importlib.import_module("deepmd.tf2.model.make_model") captured: dict[str, Any] = {} def fake_communicate( @@ -341,7 +341,7 @@ def call_lower( def test_model_call_from_call_lower_reshapes_coord_corr_for_mapping( monkeypatch: pytest.MonkeyPatch, ) -> None: - make_model_module = importlib.import_module("deepmd.tf2.make_model") + make_model_module = importlib.import_module("deepmd.tf2.model.make_model") captured: dict[str, Any] = {} def fake_communicate( @@ -494,7 +494,7 @@ def test_prepare_lower_inputs_folds_in_pair_exclusion() -> None: from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, ) - from deepmd.tf2.make_model import ( + from deepmd.tf2.model.make_model import ( prepare_lower_inputs, ) From 6363eea53edb493edbe4b8346452488513be6662 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 10:25:33 +0800 Subject: [PATCH 62/63] docs(pt_expt): note why deep_eval passes pair_excl directly to its builder _nlist_builder is always the in-repo VesinNeighborList (never user-injected), so the direct keyword is correct here; the conditional-kwargs idiom in dpmodel/tf2 make_model exists only for third-party NeighborList strategies injected via the public neighbor_list= parameter. --- deepmd/pt_expt/infer/deep_eval.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index f198d4bbe4..2f348e9e10 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1005,6 +1005,11 @@ def _build_nlist_native( # ``eval_descriptor`` calling the descriptor directly -- receive the # type-distinguished nlist a non-mixed-type descriptor expects. The # main eval path is unaffected (its ``format_nlist`` re-formats). + # The BUILDER owns exclusion. ``_nlist_builder`` is always our own + # VesinNeighborList (never user-injected), so the keyword is passed + # directly -- the conditional-kwargs idiom in dpmodel/tf2 make_model + # exists only for third-party NeighborList strategies injected via + # the public ``neighbor_list=`` parameter. extended_coord, extended_atype, nlist, mapping = self._nlist_builder.build( coords, atom_types, cells, rcut, sel, pair_excl=pair_excl ) From a33ff7905c68c6f6d908c512448a8921a61a0bae Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 10:28:15 +0800 Subject: [PATCH 63/63] refactor: pass pair_excl to build() directly; pair_excl is part of the contract Drop the conditional **excl_kwargs indirection in dpmodel/tf2 make_model and pass pair_excl=pair_excl uniformly (matching pt_expt deep_eval and the graph builders). pair_excl is part of the NeighborList.build() contract (base-class signature); a custom strategy predating it fails loudly with TypeError instead of silently including excluded pairs. --- deepmd/dpmodel/model/make_model.py | 11 ++++------- deepmd/pt_expt/infer/deep_eval.py | 5 ----- deepmd/tf2/model/make_model.py | 9 ++++----- 3 files changed, 8 insertions(+), 17 deletions(-) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index f31f7e70e1..3a5d19f83a 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -132,14 +132,11 @@ def model_call_from_call_lower( builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() # Model-level pair exclusion is a nlist-BUILD transform (decision #18/A4): # the BUILDER owns it (mirroring build_neighbor_graph on the graph path), so - # the lower always consumes a pre-excluded nlist. The keyword is passed only - # when set: custom NeighborList strategies predating it keep working for - # every model without exclusion, while a legacy builder combined with a - # non-empty exclusion fails loudly (TypeError) instead of silently - # including excluded pairs. - excl_kwargs = {} if pair_excl is None else {"pair_excl": pair_excl} + # the lower always consumes a pre-excluded nlist. ``pair_excl`` is part of + # the NeighborList.build() contract; a custom strategy predating it fails + # loudly (TypeError) instead of silently including excluded pairs. extended_coord, extended_atype, nlist, mapping = builder.build( - cc, atype, bb, rcut, sel, **excl_kwargs + cc, atype, bb, rcut, sel, pair_excl=pair_excl ) extended_coord = extended_coord.reshape(nframes, -1, 3) if coord_corr_for_virial is not None: diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 2f348e9e10..f198d4bbe4 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -1005,11 +1005,6 @@ def _build_nlist_native( # ``eval_descriptor`` calling the descriptor directly -- receive the # type-distinguished nlist a non-mixed-type descriptor expects. The # main eval path is unaffected (its ``format_nlist`` re-formats). - # The BUILDER owns exclusion. ``_nlist_builder`` is always our own - # VesinNeighborList (never user-injected), so the keyword is passed - # directly -- the conditional-kwargs idiom in dpmodel/tf2 make_model - # exists only for third-party NeighborList strategies injected via - # the public ``neighbor_list=`` parameter. extended_coord, extended_atype, nlist, mapping = self._nlist_builder.build( coords, atom_types, cells, rcut, sel, pair_excl=pair_excl ) diff --git a/deepmd/tf2/model/make_model.py b/deepmd/tf2/model/make_model.py index 24c0c3ebd6..a5e1d74565 100644 --- a/deepmd/tf2/model/make_model.py +++ b/deepmd/tf2/model/make_model.py @@ -250,12 +250,11 @@ def no_pbc() -> tuple[Array, Array, Array, Array]: uses_native_nlist_builder = neighbor_list is None if neighbor_list is not None: # The BUILDER owns model-level pair exclusion (same convention as - # dpmodel/pt_expt make_model). The keyword is passed only when set, so - # legacy custom strategies keep working without exclusion and fail - # loudly (TypeError) with it instead of silently including pairs. - excl_kwargs = {} if pair_excl is None else {"pair_excl": pair_excl} + # dpmodel/pt_expt make_model). ``pair_excl`` is part of the + # NeighborList.build() contract; a custom strategy predating it fails + # loudly (TypeError) instead of silently including excluded pairs. extended_coord, extended_atype, nlist, mapping = neighbor_list.build( - cc, atype, bb, rcut, sel, **excl_kwargs + cc, atype, bb, rcut, sel, pair_excl=pair_excl ) else: has_pbc = _box_has_pbc(bb)