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,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)
-        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.
 
         ext_atom_mask = self.make_atom_mask(extended_atype)
         ret_dict = self.forward_atomic(
@@ -326,11 +330,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 +360,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))
-        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.
         ret_dict = self.forward_atomic_graph(
             graph,
             atype_clamped,
@@ -673,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/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,20 +435,34 @@ 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``
         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`).
         """
-        return (
-            self.se_atten.tebd_input_mode == "concat"
-            and not self.se_atten.exclude_types
-        )
+        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
@@ -575,9 +592,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:
@@ -626,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,
@@ -659,9 +666,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 +679,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
 
@@ -680,8 +694,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
         ----------
@@ -1377,7 +1391,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():
@@ -1749,10 +1767,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,21 +1779,22 @@ 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)
         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.

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 @@
     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 @@
         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 @@
     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:
@@ -313,9 +322,8 @@
 
                 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
@@ -378,6 +386,8 @@
                     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
@@ -461,10 +471,20 @@
                     "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 "