feat(dpmodel): introduce NeighborContract for graph-native preprocessing

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -10,6 +10,9 @@
 )
 
 if TYPE_CHECKING:
+    from deepmd.dpmodel.utils.neighbor_contract import (
+        NeighborContract,
+    )
     from deepmd.dpmodel.utils.neighbor_graph import (
         NeighborGraph,
     )
@@ -199,6 +202,21 @@ def uses_graph_lower(self) -> bool:
         """
         return False
 
+    def get_neighbor_contract(self) -> "NeighborContract":
+        """Return the neighbor representation contract of this atomic model.
+
+        Default matches :meth:`uses_graph_lower`: graph-capable models publish
+        a capacity-free graph contract; dense models publish ``get_sel()``.
+        Compositions override to merge or short-circuit child contracts.
+        """
+        from deepmd.dpmodel.utils.neighbor_contract import (
+            NeighborContract,
+        )
+
+        if self.uses_graph_lower():
+            return NeighborContract.graph()
+        return NeighborContract.dense(self.get_sel(), requires_capacity=False)
+
     def supports_native_spin(self) -> bool:
         """Returns whether this atomic model consumes a per-atom spin input.
 

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -8,6 +8,9 @@
 )
 
 if TYPE_CHECKING:
+    from deepmd.dpmodel.utils.neighbor_contract import (
+        NeighborContract,
+    )
     from deepmd.dpmodel.utils.neighbor_graph import (
         NeighborGraph,
     )
@@ -257,6 +260,19 @@ def get_sel(self) -> list[int]:
         """Get the neighbor selection."""
         return self.descriptor.get_sel()
 
+    def get_neighbor_contract(self) -> "NeighborContract":
+        """Return the neighbor representation contract of this atomic model."""
+        getter = getattr(self.descriptor, "get_neighbor_contract", None)
+        if callable(getter):
+            return getter()
+        from deepmd.dpmodel.utils.neighbor_contract import (
+            NeighborContract,
+        )
+
+        if self.uses_graph_lower():
+            return NeighborContract.graph()
+        return NeighborContract.dense(self.get_sel(), requires_capacity=False)
+
     def set_case_embd(self, case_idx: int) -> None:
         """
         Set the case embedding of this atomic model by the given case_idx,

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -3,9 +3,15 @@
     Callable,
 )
 from typing import (
+    TYPE_CHECKING,
     Any,
 )
 
+if TYPE_CHECKING:
+    from deepmd.dpmodel.utils.neighbor_contract import (
+        NeighborContract,
+    )
+
 import array_api_compat
 import numpy as np
 
@@ -325,6 +331,25 @@ def compression_needs_min_nbor_dist(self) -> bool:
         """
         return any(m.compression_needs_min_nbor_dist() for m in self.models)
 
+    def get_neighbor_contract(self) -> "NeighborContract":
+        """Merge child neighbor contracts into one consistent contract.
+
+        Mixed dense compositions (e.g. standard DP+ZBL with ``PairTab``) do
+        not share one representation; publish the composition's dense
+        ``get_sel()`` instead of asking incompatible children to merge.
+        """
+        from deepmd.dpmodel.utils.neighbor_contract import (
+            NeighborContract,
+        )
+
+        if not self.uses_graph_lower():
+            return NeighborContract.dense(self.get_sel(), requires_capacity=False)
+        contracts = [m.get_neighbor_contract() for m in self.models]
+        contract = contracts[0]
+        for other in contracts[1:]:
+            contract = contract.merge(other)
+        return contract
+
     def uses_graph_lower(self) -> bool:
         """Graph-capable iff EVERY child supports the graph lower.
 

deepmd/dpmodel/descriptor/dpa1.py

@@ -40,6 +40,12 @@
 from deepmd.dpmodel.utils.env_mat_stat import (
     EnvMatStatSe,
 )
+from deepmd.dpmodel.utils.neighbor_contract import (
+    NeighborContract,
+    ensure_construction_sel,
+    graph_eligible_tebd_mode,
+    is_auto_sel,
+)
 from deepmd.dpmodel.utils.network import (
     LayerNorm,
     NativeLayer,
@@ -629,6 +635,49 @@ def uses_compact_edge_pairs(self) -> bool:
         """
         return self.se_atten.attn_layer > 0
 
+    @classmethod
+    def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract:
+        """Resolve DPA1 neighbor contract from config.
+
+        Graph-eligible tebd modes publish a graph representation. Capacity
+        discovery is skipped only when ``sel`` is omitted: ``auto`` / explicit
+        ``sel`` still run the legacy update path so dual-path mean/std and
+        ``nnei`` normalization stay unchanged until the companion sel-decoupling
+        work lands.
+        """
+        sel = local_jdata.get("sel")
+        if not graph_eligible_tebd_mode(local_jdata.get("tebd_input_mode")):
+            return NeighborContract.from_legacy_sel(sel)
+        if sel is None:
+            # Omitted sel: graph-native config that skips capacity discovery.
+            return NeighborContract.graph()
+        if is_auto_sel(sel):
+            # Keep auto-sel discovery so existing dual-path nnei/mean-std
+            # behavior is unchanged until companion issue #5824.
+            return NeighborContract(
+                representation="graph",
+                requires_capacity=True,
+                capacity=None,
+            )
+        # Explicit capacity: still graph execution, no discovery needed.
+        return NeighborContract(
+            representation="graph",
+            requires_capacity=False,
+            capacity=None,
+        )
+
+    @classmethod
+    def prepare_jdata_for_neighbor_contract(
+        cls,
+        local_jdata: dict,
+        contract: NeighborContract,
+    ) -> dict:
+        """Inject a construction ``sel`` only when capacity discovery is skipped."""
+        if contract.is_graph and not contract.requires_capacity:
+            if local_jdata.get("sel") is None:
+                return ensure_construction_sel(local_jdata)
+        return dict(local_jdata)
+
     def disable_graph_lower(self) -> None:
         """Force the legacy dense lower for this descriptor.
 

deepmd/dpmodel/descriptor/dpa2.py

@@ -26,6 +26,11 @@
     EnvMat,
     NetworkCollection,
 )
+from deepmd.dpmodel.utils.neighbor_contract import (
+    GRAPH_NATIVE_CONSTRUCTION_SEL,
+    NeighborContract,
+    is_auto_sel,
+)
 from deepmd.dpmodel.utils.network import (
     Identity,
     NativeLayer,
@@ -753,6 +758,108 @@ def uses_graph_lower(self) -> bool:
             return False
         return self.repinit.tebd_input_mode in ("concat", "strip")
 
+    @classmethod
+    def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract:
+        """Match :meth:`uses_graph_lower` eligibility from config alone.
+
+        DPA2 stores capacities under ``repinit.nsel`` / ``repformer.nsel``.
+        Graph-eligible configs keep auto-nsel discovery (as DPA1 does for
+        ``sel: auto``) so dual-path mean/std and nnei normalization stay
+        intact until companion sel-decoupling work lands.
+
+        Dense discovery is required when *any* block still needs capacity
+        (repinit / repformer / three_body), not only when repinit is auto.
+        """
+        repinit = local_jdata.get("repinit") or {}
+        repformer = local_jdata.get("repformer") or {}
+        use_three_body = bool(
+            local_jdata.get("use_three_body", False)
+            or repinit.get("use_three_body", False)
+        )
+
+        def _block_sel(block: dict) -> object:
+            if "nsel" in block:
+                return block.get("nsel")
+            return block.get("sel")
+
+        def _needs_discovery(sel: object) -> bool:
+            return sel is None or is_auto_sel(sel)
+
+        def _any_block_needs_discovery() -> bool:
+            needs = _needs_discovery(_block_sel(repinit)) or _needs_discovery(
+                _block_sel(repformer)
+            )
+            if use_three_body:
+                needs = needs or _needs_discovery(repinit.get("three_body_sel"))
+            return needs
+
+        def _dense_from_blocks() -> NeighborContract:
+            # Outer capacity tracks repinit; OR requires_capacity across blocks
+            # so an explicit repinit cannot skip discovery for auto repformer /
+            # three_body_sel.
+            base = NeighborContract.from_legacy_sel(_block_sel(repinit))
+            if _any_block_needs_discovery():
+                return NeighborContract.dense(base.capacity, requires_capacity=True)
+            return base
+
+        if use_three_body:
+            return _dense_from_blocks()
+        if not repinit.get("set_davg_zero", True):
+            return _dense_from_blocks()
+        if not repformer.get("set_davg_zero", True):
+            return _dense_from_blocks()
+        tebd = repinit.get("tebd_input_mode", "concat")
+        if tebd not in ("concat", "strip"):
+            return _dense_from_blocks()
+
+        sels = [_block_sel(repinit), _block_sel(repformer)]
+        if all(sel is None for sel in sels):
+            return NeighborContract.graph()
+        if any(is_auto_sel(sel) for sel in sels):
+            return NeighborContract.graph(requires_capacity=True)
+        return NeighborContract.graph(requires_capacity=False)
+
+    @classmethod
+    def prepare_jdata_for_neighbor_contract(
+        cls,
+        local_jdata: dict,
+        contract: NeighborContract,
+    ) -> dict:
+        """Rewrite block ``nsel`` placeholders; never inject a top-level ``sel``.
+
+        Construction placeholders must satisfy
+        ``repinit.nsel > repformer.nsel`` (:meth:`DescrptDPA2.__init__`).
+        """
+        out = dict(local_jdata)
+        if not (contract.is_graph and not contract.requires_capacity):
+            return out
+
+        def _numeric_nsel(sel: object) -> int | None:
+            if sel is None or is_auto_sel(sel):
+                return None
+            if isinstance(sel, (list, tuple)):
+                return int(sel[0]) if sel else None
+            return int(sel)
+
+        repinit = dict(out.get("repinit") or {})
+        repformer = dict(out.get("repformer") or {})
+        repinit_nsel = repinit.get("nsel", repinit.get("sel"))
+        repformer_nsel = repformer.get("nsel", repformer.get("sel"))
+        base = int(GRAPH_NATIVE_CONSTRUCTION_SEL)
+
+        if repformer_nsel is None or is_auto_sel(repformer_nsel):
+            repformer["nsel"] = base
+            out["repformer"] = repformer
+            repformer_value = base
+        else:
+            repformer_value = _numeric_nsel(repformer_nsel)
+
+        if repinit_nsel is None or is_auto_sel(repinit_nsel):
+            # Keep the DescrptDPA2 ordering invariant: repinit > repformer.
+            repinit["nsel"] = max(base + 1, (repformer_value or base) + 1)
+            out["repinit"] = repinit
+        return out
+
     def graph_type_embedding_table(self) -> Array:
         """Full type-embedding table consumed by the graph-route forward.
 

deepmd/dpmodel/descriptor/dpa4.py

@@ -63,6 +63,11 @@
 from deepmd.dpmodel.utils.exclude_mask import (
     PairExcludeMask,
 )
+from deepmd.dpmodel.utils.neighbor_contract import (
+    NeighborContract,
+    ensure_construction_sel,
+    is_auto_sel,
+)
 from deepmd.dpmodel.utils.neighbor_graph import (
     apply_pair_exclusion,
     frame_id_from_n_node,
@@ -2447,6 +2452,29 @@ def uses_graph_lower(self) -> bool:
         """
         return not self._graph_lower_disabled
 
+    @classmethod
+    def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract:
+        """DPA4 defaults to the carry-all graph lower.
+
+        ``sel: auto`` still requires capacity discovery so update_sel can
+        materialize a numeric width for mean/std buffers.
+        """
+        if is_auto_sel(local_jdata.get("sel")):
+            return NeighborContract.graph(requires_capacity=True)
+        return NeighborContract.graph()
+
+    @classmethod
+    def prepare_jdata_for_neighbor_contract(
+        cls,
+        local_jdata: dict,
+        contract: NeighborContract,
+    ) -> dict:
+        """Inject a construction ``sel`` only when capacity discovery is skipped."""
+        if contract.is_graph and not contract.requires_capacity:
+            if local_jdata.get("sel") is None or is_auto_sel(local_jdata.get("sel")):
+                return ensure_construction_sel(local_jdata)
+        return dict(local_jdata)
+
     def uses_compact_edge_pairs(self) -> bool:
         """DPA4 attention is a per-edge scatter softmax; no pair axis."""
         return False

deepmd/dpmodel/descriptor/dpa4c.py

@@ -49,6 +49,9 @@
 from deepmd.dpmodel.utils import (
     PairExcludeMask,
 )
+from deepmd.dpmodel.utils.neighbor_contract import (
+    NeighborContract,
+)
 from deepmd.dpmodel.utils.network import (
     NativeLayer,
 )
@@ -2030,6 +2033,22 @@ def uses_graph_lower(self) -> bool:
         """Return whether graph-native lowering is supported."""
         return True
 
+    @classmethod
+    def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract:
+        """DPA4C is always carry-all graph-native and has no capacity."""
+        del local_jdata
+        return NeighborContract.graph()
+
+    @classmethod
+    def prepare_jdata_for_neighbor_contract(
+        cls,
+        local_jdata: dict,
+        contract: NeighborContract,
+    ) -> dict:
+        """No ``sel`` field to rewrite."""
+        del contract
+        return dict(local_jdata)
+
     def disable_graph_lower(self) -> None:
         """Reject the dense lower, which this descriptor cannot serve.