refactor(dense): apply model-level pair_exclude at nlist BUILD, completing A4

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

Author: Han Wang

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,

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

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,

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

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,

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

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:

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

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

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)