refactor: pair exclude_types as canonical NeighborGraph transform; dpa1 graph path supports exclude_types (decision #18)

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 (
@@ -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,10 +297,14 @@ 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)
+        # 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.  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(
@@ -326,11 +333,11 @@ 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
-        gated by ``uses_graph_lower()==False``.
+        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
         ----------
@@ -356,14 +363,13 @@ 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),
-            )
+        # 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.
+        # 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,92 @@ 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.
+
+        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
+        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.
+
+        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
+        ):
+            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:
@@ -673,6 +765,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/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."""
@@ -432,26 +435,42 @@ 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 OR strip type-embedding. Remaining ineligible configs
-        (``exclude_types``, and compressed descriptors) fall back to the legacy
-        dense path, so those models keep working unchanged.
+        with concat OR strip type-embedding.  ``exclude_types`` is fully
+        supported via
+        :func:`~deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`.
+        Compressed descriptors are the remaining ineligible config and fall
+        back to the legacy dense path, so those models keep working unchanged.
 
         Eligibility does NOT imply numerical interchangeability with the
         dense route for every config: with ``smooth_type_embedding=True``
         the carry-all graph attention is sel-independent by design and
         differs from the dense lower by up to ~1e-4 (see the Notes of
         :meth:`call_graph`).
         """
+        if self._graph_lower_disabled:
+            return False
         # compressed descriptors have no graph kernel (geo/tebd tabulation is
         # dense-only); keep them on the legacy dense path.
         if self.compress:
             return False
-        # exclude_types stays dense (graph exclusion is owned elsewhere); strip is
-        # now graph-eligible (per-edge factorized embedding, no neighbor coupling).
-        return (
-            self.se_atten.tebd_input_mode in ("concat", "strip")
-            and not self.se_atten.exclude_types
-        )
+        # strip is graph-eligible (per-edge factorized embedding, no neighbor
+        # coupling); exclude_types is graph-native via ``apply_pair_exclusion``
+        # (owned at this seam). Only compression / the disable flag force dense.
+        return self.se_atten.tebd_input_mode in ("concat", "strip")
+
+    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
@@ -581,9 +600,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 (exclude_types, compressed descriptors) 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 (compressed descriptors) 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:
@@ -632,26 +651,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,
@@ -665,9 +674,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,
@@ -677,6 +687,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
 
@@ -687,7 +703,7 @@ def _call_dense(
         nlist: Array,
     ) -> Array:
         """Legacy dense descriptor body (the ineligible ``call`` path:
-        compressed descriptors, exclude_types, or the no-mapping ghost case).
+        compressed descriptors or the no-mapping ghost case).
 
         Parameters
         ----------
@@ -1383,7 +1399,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():
@@ -1756,9 +1776,10 @@ def call_graph(
         -----
         Known limitations:
         - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense;
-        - ``exclude_types`` is not yet supported and raises (lands in a later PR).
+        - ``exclude_types`` is applied graph-natively via ``apply_pair_exclusion``.
         """
         from deepmd.dpmodel.utils.neighbor_graph import (
+            apply_pair_exclusion,
             edge_env_mat,
             segment_sum,
         )
@@ -1767,21 +1788,22 @@ def call_graph(
             raise NotImplementedError(
                 f"graph path does not support tebd_input_mode={self.tebd_input_mode!r}"
             )
-        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)
         dev = array_api_compat.device(graph.edge_vec)
         # 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.