From fde4d9fd765fd0677059e8add36d9127f27ac024 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:16:04 +0800 Subject: [PATCH 01/12] feat(dpmodel): introduce NeighborContract for graph-native preprocessing Add a backend-neutral NeighborContract under deepmd.dpmodel and expose it through descriptor/model plugins so training can decide capacity discovery before construction. Graph-native configs skip update_sel; dense models keep the legacy path. Export metadata carries a versioned contract and omits dummy sel/nnei for graph models, with legacy adapters for old archives. Entrypoints call BaseModel.prepare_neighbors instead of unconditional update_sel. DPA1 omits capacity discovery only when sel is absent; auto sel still discovers so dual-path nnei/mean-std behavior stays intact until #5824. Refs #5825 #5755 --- .../dpmodel/atomic_model/dp_atomic_model.py | 13 + .../atomic_model/linear_atomic_model.py | 9 + deepmd/dpmodel/descriptor/dpa1.py | 50 ++++ deepmd/dpmodel/descriptor/dpa2.py | 22 ++ deepmd/dpmodel/descriptor/dpa4.py | 10 + deepmd/dpmodel/descriptor/dpa4c.py | 20 ++ deepmd/dpmodel/descriptor/hybrid.py | 44 +++ .../descriptor/make_base_descriptor.py | 52 ++++ deepmd/dpmodel/model/base_model.py | 30 ++ deepmd/dpmodel/model/dp_linear_model.py | 47 ++-- deepmd/dpmodel/model/dp_model.py | 52 +++- deepmd/dpmodel/model/dp_zbl_model.py | 33 +-- deepmd/dpmodel/model/make_model.py | 4 + deepmd/dpmodel/utils/__init__.py | 4 + deepmd/dpmodel/utils/neighbor_contract.py | 261 ++++++++++++++++++ deepmd/pt/entrypoints/main.py | 4 +- deepmd/pt/model/model/dp_model.py | 41 ++- deepmd/pt_expt/entrypoints/main.py | 4 +- deepmd/pt_expt/infer/deep_eval.py | 43 ++- deepmd/pt_expt/model/dp_linear_model.py | 38 ++- deepmd/pt_expt/utils/serialization.py | 51 +++- .../common/dpmodel/test_neighbor_contract.py | 194 +++++++++++++ 22 files changed, 938 insertions(+), 88 deletions(-) create mode 100644 deepmd/dpmodel/utils/neighbor_contract.py create mode 100644 source/tests/common/dpmodel/test_neighbor_contract.py diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index ffd33b5bd2..40133b7747 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -257,6 +257,19 @@ def get_sel(self) -> list[int]: """Get the neighbor selection.""" return self.descriptor.get_sel() + def get_neighbor_contract(self): + """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, diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 5d879043af..3c50985f9d 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -325,6 +325,15 @@ 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): + """Merge child neighbor contracts into one consistent contract.""" + 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. diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 1781728951..d0c9dc83c6 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -55,6 +55,12 @@ remap_atype_to_padding, take_type_embedding, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, + ensure_construction_sel, + graph_eligible_tebd_mode, + is_auto_sel, +) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -629,6 +635,50 @@ 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. diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 55566fee54..cac8d108c4 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -41,6 +41,9 @@ TypeEmbedNet, take_type_embedding, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -753,6 +756,25 @@ 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.""" + if local_jdata.get("use_three_body", False): + return NeighborContract.from_legacy_sel( + (local_jdata.get("repinit") or {}).get("sel", local_jdata.get("sel")) + ) + repinit = local_jdata.get("repinit") or {} + repformer = local_jdata.get("repformer") or {} + if not repinit.get("set_davg_zero", True): + return NeighborContract.from_legacy_sel(repinit.get("sel")) + if not repformer.get("set_davg_zero", True): + return NeighborContract.from_legacy_sel(repinit.get("sel")) + tebd = repinit.get("tebd_input_mode", "concat") + if tebd in ("concat", "strip"): + return NeighborContract.graph() + return NeighborContract.from_legacy_sel(repinit.get("sel")) + def graph_type_embedding_table(self) -> Array: """Full type-embedding table consumed by the graph-route forward. diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 59c027bbf9..7410961106 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -71,6 +71,9 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -2447,6 +2450,13 @@ 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.""" + del local_jdata + return NeighborContract.graph() + def uses_compact_edge_pairs(self) -> bool: """DPA4 attention is a per-edge scatter softmax; no pair axis.""" return False diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index e38d8374da..8de461f4eb 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -55,6 +55,9 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -2030,6 +2033,23 @@ 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. diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 796678f3bd..1d25085cb3 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -26,6 +26,9 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -419,6 +422,47 @@ def call( return out_descriptor, out_gr, out_g2, out_h2, out_sw @classmethod + + @classmethod + def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: + """Merge child descriptor contracts; require a single representation.""" + from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, + ) + + children = local_jdata.get("list") or [] + if not children: + raise ValueError("hybrid descriptor config has an empty child list") + contract = BaseDescriptor.neighbor_contract_from_jdata(children[0]) + for child in children[1:]: + contract = contract.merge(BaseDescriptor.neighbor_contract_from_jdata(child)) + return contract + + @classmethod + def prepare_jdata_for_neighbor_contract( + cls, + local_jdata: dict, + contract: NeighborContract, + ) -> dict: + """Prepare each hybrid child under the merged contract.""" + from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, + ) + + out = dict(local_jdata) + prepared = [] + for child in local_jdata.get("list") or []: + child_contract = BaseDescriptor.neighbor_contract_from_jdata(child) + # Merged contract already enforced representation equality. + prepared.append( + BaseDescriptor.prepare_jdata_for_neighbor_contract( + child, child_contract + ) + ) + out["list"] = prepared + del contract + return out + def update_sel( cls, train_data: DeepmdDataSystem, diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index b0598e9735..7ed7a7a7f4 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -23,6 +23,11 @@ from deepmd.utils.path import ( DPPath, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, + dense_contract_from_sel_field, + ensure_construction_sel, +) from deepmd.utils.plugin import ( PluginVariant, make_plugin_registry, @@ -248,6 +253,53 @@ def uses_compact_edge_pairs(self) -> bool: """ return False + def get_neighbor_contract(self) -> NeighborContract: + """Return the neighbor representation contract of this descriptor. + + Graph-eligible descriptors publish a capacity-free graph contract. + Dense descriptors publish their ``get_sel()`` capacity. + """ + if self.uses_graph_lower(): + return NeighborContract.graph() + return NeighborContract.dense(self.get_sel(), requires_capacity=False) + + @classmethod + def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: + """Derive a neighbor contract from descriptor config before construction. + + Dispatches through the descriptor plugin registry. The default is a + dense contract driven by the ``sel`` field (including auto-sel). + Graph-native descriptors override to return + :meth:`NeighborContract.graph`. + """ + if cls is BD: + cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__)) + return cls.neighbor_contract_from_jdata(local_jdata) + return dense_contract_from_sel_field(local_jdata.get("sel")) + + @classmethod + def prepare_jdata_for_neighbor_contract( + cls, + local_jdata: dict, + contract: NeighborContract, + ) -> dict: + """Normalize descriptor config under a resolved neighbor contract. + + Graph contracts replace missing/auto ``sel`` with a construction + placeholder so models that still size mean/std buffers can build + without neighbor-capacity discovery. Dense contracts leave the + config for ``update_sel``. + """ + if cls is BD: + cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__)) + return cls.prepare_jdata_for_neighbor_contract(local_jdata, contract) + if contract.is_graph: + # Default assumes a sel-bearing descriptor. Graph-native + # descriptors without sel (DPA4C) override this method. + return ensure_construction_sel(local_jdata) + return dict(local_jdata) + + def dense_lower_supports_comm(self) -> bool: """Whether the DENSE (nlist) lower implements comm_dict exchange. diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 394b124f2f..8c9f6c4924 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -256,6 +256,36 @@ def update_sel( cls = cls.get_class_by_type(model_type) return cls.update_sel(train_data, type_map, local_jdata) + @classmethod + def prepare_neighbors( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float | None]: + """Prepare neighbor config under the model neighbor contract. + + Dispatches through the model plugin registry. Prefers a concrete + ``prepare_neighbors`` defined above this base class; otherwise + falls back to :meth:`update_sel`. + """ + model_type = local_jdata.get("type", "standard") + if model_type == "standard": + model_type = local_jdata.get("fitting", {}).get("type", "ener") + concrete = cls.get_class_by_type(model_type) + prepare = None + for klass in concrete.__mro__: + if klass is BaseBaseModel: + break + if "prepare_neighbors" in klass.__dict__: + prepare = klass.__dict__["prepare_neighbors"] + break + if prepare is not None: + return prepare.__get__(None, concrete)( + train_data, type_map, local_jdata + ) + return concrete.update_sel(train_data, type_map, local_jdata) + @abstractmethod def get_observed_type_list(self) -> list[str]: """Get observed types (elements) of the model during data statistics. diff --git a/deepmd/dpmodel/model/dp_linear_model.py b/deepmd/dpmodel/model/dp_linear_model.py index 2255edbede..c435fa619f 100644 --- a/deepmd/dpmodel/model/dp_linear_model.py +++ b/deepmd/dpmodel/model/dp_linear_model.py @@ -50,46 +50,41 @@ def __init__( DPLinearModel_.__init__(self, *args, **kwargs) @classmethod - def update_sel( + def prepare_neighbors( cls, train_data: DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict, ) -> tuple[dict, float | None]: - """Update the selection and perform neighbor statistics. - - Updates each learned child in place, skipping analytical - (``inner_potential``) and pair-table children, and aggregates the - minimum neighbor distance (twin of the pt_expt implementation). - - Parameters - ---------- - train_data : DeepmdDataSystem - data used to do neighbor statistics - type_map : list[str], optional - The name of each type of atoms - local_jdata : dict - The local data refer to the current class - - Returns - ------- - dict - The updated local data - float - The minimum distance between two atoms - """ + """Prepare each learned child under its neighbor contract.""" local_jdata_cpy = local_jdata.copy() type_map = local_jdata_cpy["type_map"] min_nbor_dist = None for idx, sub_model in enumerate(local_jdata_cpy["models"]): if sub_model.get("type") == "inner_potential": - # analytical child: no descriptor, no selection to update continue if "tab_file" not in sub_model: - sub_model, temp_min = DPModelCommon.update_sel( + sub_model, temp_min = DPModelCommon.prepare_neighbors( train_data, type_map, local_jdata_cpy["models"][idx] ) local_jdata_cpy["models"][idx] = sub_model - if min_nbor_dist is None or temp_min <= min_nbor_dist: + if min_nbor_dist is None or ( + temp_min is not None and temp_min <= min_nbor_dist + ): min_nbor_dist = temp_min return local_jdata_cpy, min_nbor_dist + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float | None]: + """Update the selection and perform neighbor statistics. + + Updates each learned child in place, skipping analytical + (``inner_potential``) and pair-table children, and aggregates the + minimum neighbor distance (twin of the pt_expt implementation). + """ + return cls.prepare_neighbors(train_data, type_map, local_jdata) diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 5a1a253596..894c5666bb 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -7,6 +7,12 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) +from deepmd.dpmodel.utils.update_sel import ( + UpdateSel, +) from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -20,6 +26,44 @@ class DPModelCommon: neighbor selection updates and fitting network access. """ + @classmethod + def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: + """Resolve the model neighbor contract from config before construction.""" + return BaseDescriptor.neighbor_contract_from_jdata(local_jdata["descriptor"]) + + @classmethod + def prepare_neighbors( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float | None]: + """Prepare neighbor config under the model neighbor contract. + + This is the single entrypoint preparation hook: it resolves the + descriptor/model contract recursively, rewrites graph-native configs so + they do not require capacity discovery, and only runs the legacy + ``update_sel`` path when the contract still needs a dense capacity. + """ + local_jdata_cpy = dict(local_jdata) + descriptor_jdata = dict(local_jdata_cpy["descriptor"]) + contract = BaseDescriptor.neighbor_contract_from_jdata(descriptor_jdata) + descriptor_jdata = BaseDescriptor.prepare_jdata_for_neighbor_contract( + descriptor_jdata, contract + ) + local_jdata_cpy["descriptor"] = descriptor_jdata + if not contract.requires_capacity: + # Capacity discovery is skipped. Still report min neighbor distance + # when the descriptor exposes a cutoff (used by compression bounds). + rcut = descriptor_jdata.get("rcut") + if rcut is None: + return local_jdata_cpy, None + return local_jdata_cpy, float(UpdateSel().get_min_nbor_dist(train_data)) + local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( + train_data, type_map, descriptor_jdata + ) + return local_jdata_cpy, min_nbor_dist + @classmethod def update_sel( cls, @@ -45,11 +89,9 @@ def update_sel( float The minimum distance between two atoms """ - local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( - train_data, type_map, local_jdata["descriptor"] - ) - return local_jdata_cpy, min_nbor_dist + # Prefer the contract-aware preparation hook so callers that still + # invoke update_sel inherit graph-native capacity skipping. + return cls.prepare_neighbors(train_data, type_map, local_jdata) def get_fitting_net(self) -> BaseFitting: """Get the fitting network.""" diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py index d4ff3ed800..4f48f25bcb 100644 --- a/deepmd/dpmodel/model/dp_zbl_model.py +++ b/deepmd/dpmodel/model/dp_zbl_model.py @@ -151,32 +151,25 @@ def translated_output_def(self) -> dict[str, Any]: return output_def @classmethod - def update_sel( + def prepare_neighbors( cls, train_data: DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict, ) -> tuple[dict, float | None]: - """Update the selection and perform neighbor statistics. - - Parameters - ---------- - train_data : DeepmdDataSystem - data used to do neighbor statistics - type_map : list[str], optional - The name of each type of atoms - local_jdata : dict - The local data refer to the current class - - Returns - ------- - dict - The updated local data - float - The minimum distance between two atoms - """ + """Prepare the learned DP child under its neighbor contract.""" local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["dpmodel"], min_nbor_dist = DPModelCommon.update_sel( + local_jdata_cpy["dpmodel"], min_nbor_dist = DPModelCommon.prepare_neighbors( train_data, type_map, local_jdata["dpmodel"] ) return local_jdata_cpy, min_nbor_dist + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float | None]: + """Update the selection and perform neighbor statistics.""" + return cls.prepare_neighbors(train_data, type_map, local_jdata) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b74fb6b697..b5f6f61a63 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -1220,6 +1220,10 @@ def get_sel(self) -> list[int]: """Returns the number of selected atoms for each type.""" return self.atomic_model.get_sel() + def get_neighbor_contract(self): + """Return the neighbor representation contract of this model.""" + return self.atomic_model.get_neighbor_contract() + def mixed_types(self) -> bool: """If true, the model 1. assumes total number of atoms aligned across frames; diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 1f22690d1a..9af058a07a 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -32,6 +32,9 @@ segment_mean, segment_sum, ) +from .neighbor_contract import ( + NeighborContract, +) from .neighbor_list import ( NeighborList, ) @@ -91,6 +94,7 @@ "NativeLayer", "NativeNet", "NeighborGraph", + "NeighborContract", "NeighborList", "NetworkCollection", "PairExcludeMask", diff --git a/deepmd/dpmodel/utils/neighbor_contract.py b/deepmd/dpmodel/utils/neighbor_contract.py new file mode 100644 index 0000000000..0aa281c489 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_contract.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-neutral neighbor representation contract. + +The contract separates *how* a model consumes neighbors (dense fixed-capacity +lists vs. carry-all graphs) from the legacy ``sel`` / ``get_sel()`` ABI. +Descriptors and models expose it through plugins so preprocessing can decide +whether neighbor-capacity discovery is required before construction. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + Any, + Literal, +) + +NeighborRepresentation = Literal["dense", "graph"] + +# Construction stand-in for graph-native descriptors that still size mean/std +# buffers with ``nnei`` until the companion sel-decoupling work lands. +# Width 1 matches the intended graph-native per-type statistic shape. +GRAPH_NATIVE_CONSTRUCTION_SEL = 1 + +NEIGHBOR_CONTRACT_VERSION = 1 + + +@dataclass(frozen=True, slots=True) +class NeighborContract: + r"""Versioned neighbor-representation capability. + + Parameters + ---------- + version + Schema version of this contract record. + representation + ``\"dense\"`` for fixed-capacity neighbor lists; ``\"graph\"`` for + carry-all ``NeighborGraph`` execution. + requires_capacity + Whether neighbor-capacity discovery (``update_sel`` / auto-sel) must + run before construction. Graph-native models set this to ``False``. + capacity + Explicit per-type (or mixed-type single) neighbor capacities when the + dense ABI needs them. ``None`` for graph-native models - metadata must + not invent a dummy capacity. + """ + + version: int = NEIGHBOR_CONTRACT_VERSION + representation: NeighborRepresentation = "dense" + requires_capacity: bool = True + capacity: tuple[int, ...] | None = None + + @classmethod + def dense( + cls, + capacity: list[int] | tuple[int, ...] | int | None = None, + *, + requires_capacity: bool | None = None, + ) -> NeighborContract: + """Build a dense-list contract. + + An explicit integer capacity means discovery is already settled. + ``None`` (or an auto-sel marker handled by callers) means discovery is + still required. + """ + normalized = _normalize_capacity(capacity) + if requires_capacity is None: + requires_capacity = normalized is None + return cls( + version=NEIGHBOR_CONTRACT_VERSION, + representation="dense", + requires_capacity=bool(requires_capacity), + capacity=normalized, + ) + + @classmethod + def graph(cls) -> NeighborContract: + """Build a carry-all graph contract with no neighbor capacity.""" + return cls( + version=NEIGHBOR_CONTRACT_VERSION, + representation="graph", + requires_capacity=False, + capacity=None, + ) + + @classmethod + def from_legacy_sel( + cls, + sel: list[int] | tuple[int, ...] | int | str | None, + ) -> NeighborContract: + """Reconstruct a dense contract from a legacy ``sel`` value.""" + if sel is None or _is_auto_sel(sel): + return cls.dense(None, requires_capacity=True) + return cls.dense(sel, requires_capacity=False) + + @classmethod + def from_metadata(cls, metadata: dict[str, Any]) -> NeighborContract: + """Read a contract from export/inference metadata. + + New archives carry ``neighbor_contract``. Legacy archives only have + ``sel`` / ``nnei`` and are treated as dense. + """ + raw = metadata.get("neighbor_contract") + if raw is not None: + return cls.from_dict(raw) + if "sel" in metadata: + return cls.from_legacy_sel(metadata["sel"]) + if "nnei" in metadata: + return cls.dense(int(metadata["nnei"]), requires_capacity=False) + return cls.dense(None, requires_capacity=True) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> NeighborContract: + """Deserialize a contract dictionary.""" + representation = data.get("representation", "dense") + if representation not in ("dense", "graph"): + raise ValueError( + f"Unsupported neighbor representation {representation!r}; " + "expected 'dense' or 'graph'." + ) + capacity = data.get("capacity") + return cls( + version=int(data.get("version", NEIGHBOR_CONTRACT_VERSION)), + representation=representation, + requires_capacity=bool( + data.get( + "requires_capacity", + representation == "dense" and capacity is None, + ) + ), + capacity=_normalize_capacity(capacity), + ) + + def to_dict(self) -> dict[str, Any]: + """Serialize for export metadata.""" + return { + "version": int(self.version), + "representation": self.representation, + "requires_capacity": bool(self.requires_capacity), + "capacity": list(self.capacity) if self.capacity is not None else None, + } + + @property + def is_graph(self) -> bool: + """Whether this contract describes a carry-all graph model.""" + return self.representation == "graph" + + @property + def is_dense(self) -> bool: + """Whether this contract describes a fixed-capacity dense model.""" + return self.representation == "dense" + + def legacy_sel(self) -> list[int]: + """Return a dense ``sel`` list for legacy adapters. + + Graph contracts have no capacity; callers that still need a list must + not invent one silently — this raises instead. + """ + if self.capacity is None: + raise ValueError( + "NeighborContract has no capacity to expose as legacy sel " + f"(representation={self.representation!r}, " + f"requires_capacity={self.requires_capacity})." + ) + return list(self.capacity) + + def merge(self, other: NeighborContract) -> NeighborContract: + """Combine contracts from composed descriptors/models. + + Composites must agree on representation. Dense capacities take the + element-wise maximum when both sides publish one; discovery remains + required if either side still needs it. + """ + if self.representation != other.representation: + raise ValueError( + "Incompatible neighbor contracts in a composed model: " + f"{self.representation!r} vs {other.representation!r}. " + "Every child must share one representation (dense or graph)." + ) + if self.is_graph: + return NeighborContract.graph() + capacity = _merge_capacities(self.capacity, other.capacity) + requires_capacity = self.requires_capacity or other.requires_capacity + if capacity is not None: + requires_capacity = False + return NeighborContract.dense(capacity, requires_capacity=requires_capacity) + + +def is_auto_sel(sel: Any) -> bool: + """Return whether ``sel`` is an auto-selection marker.""" + return _is_auto_sel(sel) + + +def _is_auto_sel(sel: Any) -> bool: + if not isinstance(sel, str): + return False + return sel.split(":", 1)[0] == "auto" + + +def _normalize_capacity( + capacity: list[int] | tuple[int, ...] | int | None, +) -> tuple[int, ...] | None: + if capacity is None: + return None + if isinstance(capacity, int): + return (int(capacity),) + return tuple(int(x) for x in capacity) + + +def _merge_capacities( + left: tuple[int, ...] | None, + right: tuple[int, ...] | None, +) -> tuple[int, ...] | None: + if left is None and right is None: + return None + if left is None: + return right + if right is None: + return left + if len(left) == 1 and len(right) != 1: + left = left * len(right) + if len(right) == 1 and len(left) != 1: + right = right * len(left) + if len(left) != len(right): + raise ValueError( + "Incompatible dense neighbor capacities in a composed model: " + f"{list(left)} vs {list(right)}." + ) + return tuple(max(a, b) for a, b in zip(left, right, strict=True)) + + +def dense_contract_from_sel_field(sel: Any) -> NeighborContract: + """Interpret a descriptor ``sel`` config field as a dense contract.""" + return NeighborContract.from_legacy_sel(sel) + + +def graph_eligible_tebd_mode(tebd_input_mode: str | None) -> bool: + """Whether a DPA1/DPA2-style tebd mode is graph-lower eligible.""" + return (tebd_input_mode or "concat") in ("concat", "strip") + + +def ensure_construction_sel( + local_jdata: dict[str, Any], + *, + placeholder: int = GRAPH_NATIVE_CONSTRUCTION_SEL, +) -> dict[str, Any]: + """Replace missing/auto ``sel`` with a construction placeholder. + + Used for graph-native descriptors that still allocate mean/std buffers + with an ``nnei`` axis. The placeholder is not a discovered capacity and + must not appear in exported graph metadata. + """ + out = dict(local_jdata) + sel = out.get("sel") + if sel is None or _is_auto_sel(sel): + out["sel"] = int(placeholder) + return out diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index da6b890d13..d2024bd10c 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -411,7 +411,7 @@ def train( train_data = get_data( config["training"]["training_data"], 0, type_map, None ) - config["model"], min_nbor_dist = BaseModel.update_sel( + config["model"], min_nbor_dist = BaseModel.prepare_neighbors( train_data, type_map, config["model"] ) else: @@ -439,7 +439,7 @@ def train( None, ) config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( - BaseModel.update_sel( + BaseModel.prepare_neighbors( train_data, type_map, config["model"]["model_dict"][model_item] ) ) diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 264c01e8ee..3237bd4e71 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -1,8 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.pt.model.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.pt.utils.update_sel import ( + UpdateSel, +) from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -11,6 +17,35 @@ class DPModelCommon: """A base class to implement common methods for all the Models.""" + @classmethod + def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: + """Resolve the model neighbor contract from config before construction.""" + return BaseDescriptor.neighbor_contract_from_jdata(local_jdata["descriptor"]) + + @classmethod + def prepare_neighbors( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float | None]: + """Prepare neighbor config under the model neighbor contract.""" + local_jdata_cpy = dict(local_jdata) + descriptor_jdata = dict(local_jdata_cpy["descriptor"]) + contract = BaseDescriptor.neighbor_contract_from_jdata(descriptor_jdata) + descriptor_jdata = BaseDescriptor.prepare_jdata_for_neighbor_contract( + descriptor_jdata, contract + ) + local_jdata_cpy["descriptor"] = descriptor_jdata + if not contract.requires_capacity: + if descriptor_jdata.get("rcut") is None: + return local_jdata_cpy, None + return local_jdata_cpy, float(UpdateSel().get_min_nbor_dist(train_data)) + local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( + train_data, type_map, descriptor_jdata + ) + return local_jdata_cpy, min_nbor_dist + @classmethod def update_sel( cls, @@ -36,11 +71,7 @@ def update_sel( float The minimum distance between two atoms """ - local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( - train_data, type_map, local_jdata["descriptor"] - ) - return local_jdata_cpy, min_nbor_dist + return cls.prepare_neighbors(train_data, type_map, local_jdata) # sadly, use -> BaseFitting here will not make torchscript happy def get_fitting_net(self): # noqa: ANN201 diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 7a381b4938..b6ccf1590d 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -361,7 +361,7 @@ def update_neighbor_stat( *, multi_task: bool, ) -> tuple[dict[str, Any], float | dict[str, float | None] | None]: - """Update pt_expt descriptor selections from neighbor statistics.""" + """Prepare neighbor contracts and update dense selections when required.""" log.info( "Calculate neighbor statistics... " "(add --skip-neighbor-stat to skip this step)" @@ -376,7 +376,7 @@ def update_neighbor_stat( train_data = _get_neighbor_stat_data( dict(task_config.training_data_params), type_map ) - updated_model_params, task_min_nbor_dist = BaseModel.update_sel( + updated_model_params, task_min_nbor_dist = BaseModel.prepare_neighbors( train_data, type_map, dict(task_config.model_params) ) min_nbor_dist[task_config.key] = task_min_nbor_dist diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 94c1e1f635..f78de912e3 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -73,6 +73,9 @@ from deepmd.utils.charge_state import ( CHARGE_STATE_TABLE_RANGES, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.utils.pt_checkpoint import ( detect_pt_checkpoint_backend, ) @@ -405,7 +408,18 @@ def _init_from_model_json(self, model_json_str: str) -> None: self._rcut = self._dpmodel.get_rcut() self._type_map = self._dpmodel.get_type_map() - self._sel = list(self._dpmodel.get_sel()) + self._neighbor_contract = ( + self._dpmodel.get_neighbor_contract() + if hasattr(self._dpmodel, "get_neighbor_contract") + else NeighborContract.dense( + self._dpmodel.get_sel(), requires_capacity=False + ) + ) + self._sel = ( + [] + if self._neighbor_contract.is_graph + else list(self._dpmodel.get_sel()) + ) self._mixed_types = bool(self._dpmodel.mixed_types()) if self._is_spin: spin_fitting_defs = self._dpmodel.model_output_def().def_outp.get_data() @@ -436,7 +450,14 @@ def _init_from_metadata(self) -> None: self._is_spin = bool(self.metadata.get("is_spin", False)) self._rcut = float(self.metadata["rcut"]) self._type_map = list(self.metadata["type_map"]) - self._sel = [int(s) for s in self.metadata["sel"]] + self._neighbor_contract = NeighborContract.from_metadata(self.metadata) + if self._neighbor_contract.is_graph: + # Graph-native archives omit sel; keep an empty legacy list. + self._sel = [] + elif "sel" in self.metadata: + self._sel = [int(s) for s in self.metadata["sel"]] + else: + self._sel = self._neighbor_contract.legacy_sel() self._mixed_types = bool(self.metadata["mixed_types"]) fitting_defs = [] @@ -670,7 +691,16 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: ) self._rcut = model.get_rcut() self._type_map = model.get_type_map() - self._sel = list(model.get_sel()) + self._neighbor_contract = ( + model.get_neighbor_contract() + if hasattr(model, "get_neighbor_contract") + else NeighborContract.dense(model.get_sel(), requires_capacity=False) + ) + self._sel = ( + [] + if self._neighbor_contract.is_graph + else list(model.get_sel()) + ) self._mixed_types = bool(model.mixed_types()) if self._is_spin: self._model_output_def = ModelOutputDef( @@ -700,7 +730,12 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: # (LinearEnergyAtomicModel) have no single descriptor to reach for "ntypes": len(model.get_type_map()), "rcut": model.get_rcut(), - "sel": model.get_sel(), + "neighbor_contract": self._neighbor_contract.to_dict(), + **( + {} + if self._neighbor_contract.is_graph + else {"sel": model.get_sel(), "nnei": sum(model.get_sel())} + ), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), "dim_chg_spin": model.get_dim_chg_spin(), diff --git a/deepmd/pt_expt/model/dp_linear_model.py b/deepmd/pt_expt/model/dp_linear_model.py index 85cc927f4e..f63b3ea68b 100644 --- a/deepmd/pt_expt/model/dp_linear_model.py +++ b/deepmd/pt_expt/model/dp_linear_model.py @@ -196,42 +196,36 @@ def fn( return traced @classmethod - def update_sel( + def prepare_neighbors( cls, train_data: DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict, ) -> tuple[dict, float | None]: - """Update the selection and perform neighbor statistics. - - Parameters - ---------- - train_data : DeepmdDataSystem - data used to do neighbor statistics - type_map : list[str], optional - The name of each type of atoms - local_jdata : dict - The local data refer to the current class - - Returns - ------- - dict - The updated local data - float - The minimum distance between two atoms - """ + """Prepare each learned child under its neighbor contract.""" local_jdata_cpy = local_jdata.copy() type_map = local_jdata_cpy["type_map"] min_nbor_dist = None for idx, sub_model in enumerate(local_jdata_cpy["models"]): if sub_model.get("type") == "inner_potential": - # analytical child: no descriptor, no selection to update continue if "tab_file" not in sub_model: - sub_model, temp_min = DPModelCommon.update_sel( + sub_model, temp_min = DPModelCommon.prepare_neighbors( train_data, type_map, local_jdata_cpy["models"][idx] ) local_jdata_cpy["models"][idx] = sub_model - if min_nbor_dist is None or temp_min <= min_nbor_dist: + if min_nbor_dist is None or ( + temp_min is not None and temp_min <= min_nbor_dist + ): min_nbor_dist = temp_min return local_jdata_cpy, min_nbor_dist + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: list[str] | None, + local_jdata: dict, + ) -> tuple[dict, float | None]: + """Update the selection and perform neighbor statistics.""" + return cls.prepare_neighbors(train_data, type_map, local_jdata) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index aeb11d798a..f65e8f18d2 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -32,6 +32,9 @@ from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.utils.charge_state import ( CHARGE_STATE_TABLE_RANGES, ) @@ -1120,6 +1123,34 @@ def _spin_scheme(model_type: str | None) -> str | None: return None + +def _neighbor_contract_for_model(model: torch.nn.Module) -> NeighborContract: + """Resolve the export neighbor contract from a live model.""" + getter = getattr(model, "get_neighbor_contract", None) + if callable(getter): + return getter() + atomic = getattr(model, "atomic_model", None) + if atomic is not None: + getter = getattr(atomic, "get_neighbor_contract", None) + if callable(getter): + return getter() + descriptor = getattr(atomic, "descriptor", None) + if descriptor is not None: + getter = getattr(descriptor, "get_neighbor_contract", None) + if callable(getter): + return getter() + uses = getattr(descriptor, "uses_graph_lower", None) + if callable(uses) and bool(uses()): + return NeighborContract.graph() + uses = getattr(model, "uses_graph_lower", None) + if callable(uses) and bool(uses()): + return NeighborContract.graph() + try: + return NeighborContract.dense(model.get_sel(), requires_capacity=False) + except Exception: + return NeighborContract.graph() + + def _collect_metadata( model: torch.nn.Module, spin_scheme: str | None = None, @@ -1167,11 +1198,27 @@ def _collect_metadata( "intensive": vdef.intensive, } ) + neighbor_contract = _neighbor_contract_for_model(model) meta = { "type_map": model.get_type_map(), "rcut": model.get_rcut(), - "sel": model.get_sel(), - "nnei": sum(model.get_sel()), + "neighbor_contract": neighbor_contract.to_dict(), + # Legacy dense fields: omitted for graph-native models so metadata does + # not invent a dummy capacity. Dense models keep sel/nnei for old readers. + **( + {} + if neighbor_contract.is_graph + else { + "sel": list(neighbor_contract.capacity) + if neighbor_contract.capacity is not None + else list(model.get_sel()), + "nnei": ( + sum(neighbor_contract.capacity) + if neighbor_contract.capacity is not None + else sum(model.get_sel()) + ), + } + ), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), "dim_chg_spin": model.get_dim_chg_spin(), diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py new file mode 100644 index 0000000000..b4b3e3987f --- /dev/null +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""NeighborContract unit tests.""" + +from __future__ import annotations + +import unittest +from unittest import mock + +from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.dpmodel.descriptor.dpa1 import ( + DescrptDPA1, +) +from deepmd.dpmodel.descriptor.dpa4c import ( + DescrptDPA4C, +) +from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, +) +from deepmd.dpmodel.utils.neighbor_contract import ( + GRAPH_NATIVE_CONSTRUCTION_SEL, + NeighborContract, + ensure_construction_sel, +) + + +class TestNeighborContract(unittest.TestCase): + def test_graph_has_no_capacity(self) -> None: + contract = NeighborContract.graph() + self.assertTrue(contract.is_graph) + self.assertFalse(contract.requires_capacity) + self.assertIsNone(contract.capacity) + with self.assertRaises(ValueError): + contract.legacy_sel() + + def test_dense_auto_requires_capacity(self) -> None: + contract = NeighborContract.from_legacy_sel("auto") + self.assertTrue(contract.is_dense) + self.assertTrue(contract.requires_capacity) + self.assertIsNone(contract.capacity) + + def test_dense_explicit_capacity(self) -> None: + contract = NeighborContract.from_legacy_sel([12, 24]) + self.assertEqual(contract.legacy_sel(), [12, 24]) + self.assertFalse(contract.requires_capacity) + + def test_metadata_roundtrip_and_legacy_adapter(self) -> None: + graph = NeighborContract.graph() + restored = NeighborContract.from_dict(graph.to_dict()) + self.assertEqual(restored, graph) + + legacy = NeighborContract.from_metadata({"sel": [8, 16], "nnei": 24}) + self.assertTrue(legacy.is_dense) + self.assertEqual(legacy.legacy_sel(), [8, 16]) + + from_contract = NeighborContract.from_metadata( + {"neighbor_contract": graph.to_dict()} + ) + self.assertTrue(from_contract.is_graph) + + def test_merge_rejects_mixed_representation(self) -> None: + with self.assertRaises(ValueError): + NeighborContract.graph().merge(NeighborContract.dense([4])) + + def test_ensure_construction_sel(self) -> None: + prepared = ensure_construction_sel({"type": "dpa1", "sel": "auto"}) + self.assertEqual(prepared["sel"], GRAPH_NATIVE_CONSTRUCTION_SEL) + prepared = ensure_construction_sel({"type": "dpa1"}) + self.assertEqual(prepared["sel"], GRAPH_NATIVE_CONSTRUCTION_SEL) + + +class TestDescriptorContracts(unittest.TestCase): + def test_dpa1_graph_from_jdata_without_sel(self) -> None: + jdata = { + "type": "dpa1", + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "concat", + } + contract = DescrptDPA1.neighbor_contract_from_jdata(jdata) + self.assertTrue(contract.is_graph) + self.assertFalse(contract.requires_capacity) + prepared = DescrptDPA1.prepare_jdata_for_neighbor_contract(jdata, contract) + self.assertEqual(prepared["sel"], GRAPH_NATIVE_CONSTRUCTION_SEL) + + def test_dpa1_dense_when_tebd_ineligible(self) -> None: + # Force a non-eligible tebd mode string that from_legacy still densifies. + # Unknown modes are not graph-eligible. + jdata = {"type": "dpa1", "sel": "auto", "tebd_input_mode": "other"} + contract = DescrptDPA1.neighbor_contract_from_jdata(jdata) + self.assertTrue(contract.is_dense) + self.assertTrue(contract.requires_capacity) + + def test_dpa4c_always_graph(self) -> None: + contract = DescrptDPA4C.neighbor_contract_from_jdata( + {"type": "dpa4c", "rcut": 6.0} + ) + self.assertTrue(contract.is_graph) + prepared = DescrptDPA4C.prepare_jdata_for_neighbor_contract( + {"type": "dpa4c", "rcut": 6.0}, contract + ) + self.assertNotIn("sel", prepared) + + def test_plugin_dispatch(self) -> None: + contract = BaseDescriptor.neighbor_contract_from_jdata( + { + "type": "se_atten_v2", + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "strip", + } + ) + self.assertTrue(contract.is_graph) + self.assertFalse(contract.requires_capacity) + + def test_dpa1_auto_sel_still_requires_capacity(self) -> None: + contract = DescrptDPA1.neighbor_contract_from_jdata( + { + "type": "dpa1", + "tebd_input_mode": "concat", + "sel": "auto", + } + ) + self.assertTrue(contract.is_graph) + self.assertTrue(contract.requires_capacity) + + +class TestPrepareNeighborsSkipsUpdateSel(unittest.TestCase): + def test_graph_dpa1_does_not_call_update_sel(self) -> None: + jdata = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa1", + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "concat", + }, + "fitting": {"type": "ener"}, + } + train_data = mock.MagicMock() + with ( + mock.patch( + "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel" + ) as update_sel, + mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, + ): + update_sel_cls.return_value.get_min_nbor_dist.return_value = 0.5 + updated, min_dist = DPModelCommon.prepare_neighbors( + train_data, ["O", "H"], jdata + ) + update_sel.assert_not_called() + self.assertEqual( + updated["descriptor"]["sel"], GRAPH_NATIVE_CONSTRUCTION_SEL + ) + self.assertEqual(min_dist, 0.5) + + def test_dense_still_calls_update_sel(self) -> None: + jdata = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "rcut": 6.0, + "rcut_smth": 0.5, + "sel": "auto", + "neuron": [2, 4, 8], + }, + "fitting": {"type": "ener"}, + } + train_data = mock.MagicMock() + with ( + mock.patch( + "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel", + return_value=({"type": "se_e2_a", "sel": [4, 8]}, 0.8), + ) as update_sel, + mock.patch( + "deepmd.dpmodel.model.dp_model.BaseDescriptor.neighbor_contract_from_jdata", + return_value=NeighborContract.dense(None, requires_capacity=True), + ), + mock.patch( + "deepmd.dpmodel.model.dp_model.BaseDescriptor.prepare_jdata_for_neighbor_contract", + side_effect=lambda j, c: dict(j), + ), + ): + updated, min_dist = DPModelCommon.prepare_neighbors( + train_data, ["O", "H"], jdata + ) + update_sel.assert_called_once() + self.assertEqual(updated["descriptor"]["sel"], [4, 8]) + self.assertEqual(min_dist, 0.8) + + +if __name__ == "__main__": + unittest.main() From 6f9a86023601d3bbb3c016c40307c412514d0fba Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:17:18 +0000 Subject: [PATCH 02/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../dpmodel/atomic_model/linear_atomic_model.py | 1 - deepmd/dpmodel/descriptor/dpa1.py | 13 ++++++------- deepmd/dpmodel/descriptor/dpa2.py | 7 +++---- deepmd/dpmodel/descriptor/dpa4.py | 7 +++---- deepmd/dpmodel/descriptor/dpa4c.py | 7 +++---- deepmd/dpmodel/descriptor/hybrid.py | 11 ++++++----- .../dpmodel/descriptor/make_base_descriptor.py | 11 +++++------ deepmd/dpmodel/utils/__init__.py | 8 ++++---- deepmd/pt_expt/infer/deep_eval.py | 16 +++++----------- deepmd/pt_expt/utils/serialization.py | 7 +++---- .../common/dpmodel/test_neighbor_contract.py | 8 ++++++-- 11 files changed, 44 insertions(+), 52 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 3c50985f9d..86f3b08848 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -325,7 +325,6 @@ 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): """Merge child neighbor contracts into one consistent contract.""" contracts = [m.get_neighbor_contract() for m in self.models] diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index d0c9dc83c6..82f07784b8 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/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, @@ -55,12 +61,6 @@ remap_atype_to_padding, take_type_embedding, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, - ensure_construction_sel, - graph_eligible_tebd_mode, - is_auto_sel, -) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -635,7 +635,6 @@ 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. diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index cac8d108c4..a5a00d58ac 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -26,6 +26,9 @@ EnvMat, NetworkCollection, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.network import ( Identity, NativeLayer, @@ -41,9 +44,6 @@ TypeEmbedNet, take_type_embedding, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, -) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -756,7 +756,6 @@ 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.""" diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index 7410961106..ef2e05b63c 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -63,6 +63,9 @@ from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.neighbor_graph import ( apply_pair_exclusion, frame_id_from_n_node, @@ -71,9 +74,6 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, -) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -2450,7 +2450,6 @@ 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.""" diff --git a/deepmd/dpmodel/descriptor/dpa4c.py b/deepmd/dpmodel/descriptor/dpa4c.py index 8de461f4eb..f61f92bb4c 100644 --- a/deepmd/dpmodel/descriptor/dpa4c.py +++ b/deepmd/dpmodel/descriptor/dpa4c.py @@ -49,15 +49,15 @@ from deepmd.dpmodel.utils import ( PairExcludeMask, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.network import ( NativeLayer, ) from deepmd.dpmodel.utils.seed import ( child_seed, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, -) from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) @@ -2033,7 +2033,6 @@ 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.""" diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 1d25085cb3..b0cf5fd06b 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -17,6 +17,9 @@ from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.nlist import ( nlist_distinguish_types, ) @@ -26,9 +29,6 @@ from deepmd.utils.path import ( DPPath, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, -) from deepmd.utils.version import ( check_version_compatibility, ) @@ -422,7 +422,6 @@ def call( return out_descriptor, out_gr, out_g2, out_h2, out_sw @classmethod - @classmethod def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: """Merge child descriptor contracts; require a single representation.""" @@ -435,7 +434,9 @@ def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: raise ValueError("hybrid descriptor config has an empty child list") contract = BaseDescriptor.neighbor_contract_from_jdata(children[0]) for child in children[1:]: - contract = contract.merge(BaseDescriptor.neighbor_contract_from_jdata(child)) + contract = contract.merge( + BaseDescriptor.neighbor_contract_from_jdata(child) + ) return contract @classmethod diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index 7ed7a7a7f4..240495c466 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -17,17 +17,17 @@ from deepmd.dpmodel.array_api import ( Array, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, + dense_contract_from_sel_field, + ensure_construction_sel, +) from deepmd.utils.data_system import ( DeepmdDataSystem, ) from deepmd.utils.path import ( DPPath, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, - dense_contract_from_sel_field, - ensure_construction_sel, -) from deepmd.utils.plugin import ( PluginVariant, make_plugin_registry, @@ -299,7 +299,6 @@ def prepare_jdata_for_neighbor_contract( return ensure_construction_sel(local_jdata) return dict(local_jdata) - def dense_lower_supports_comm(self) -> bool: """Whether the DENSE (nlist) lower implements comm_dict exchange. diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 9af058a07a..554931e3ca 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -18,6 +18,9 @@ is_lmdb, make_neighbor_stat_data, ) +from .neighbor_contract import ( + NeighborContract, +) from .neighbor_graph import ( GraphLayout, NeighborGraph, @@ -32,9 +35,6 @@ segment_mean, segment_sum, ) -from .neighbor_contract import ( - NeighborContract, -) from .neighbor_list import ( NeighborList, ) @@ -93,8 +93,8 @@ "LmdbTestDataNlocView", "NativeLayer", "NativeNet", - "NeighborGraph", "NeighborContract", + "NeighborGraph", "NeighborList", "NetworkCollection", "PairExcludeMask", diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index f78de912e3..84e4a84b41 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -22,6 +22,9 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -73,9 +76,6 @@ from deepmd.utils.charge_state import ( CHARGE_STATE_TABLE_RANGES, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, -) from deepmd.utils.pt_checkpoint import ( detect_pt_checkpoint_backend, ) @@ -416,9 +416,7 @@ def _init_from_model_json(self, model_json_str: str) -> None: ) ) self._sel = ( - [] - if self._neighbor_contract.is_graph - else list(self._dpmodel.get_sel()) + [] if self._neighbor_contract.is_graph else list(self._dpmodel.get_sel()) ) self._mixed_types = bool(self._dpmodel.mixed_types()) if self._is_spin: @@ -696,11 +694,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: if hasattr(model, "get_neighbor_contract") else NeighborContract.dense(model.get_sel(), requires_capacity=False) ) - self._sel = ( - [] - if self._neighbor_contract.is_graph - else list(model.get_sel()) - ) + self._sel = [] if self._neighbor_contract.is_graph else list(model.get_sel()) self._mixed_types = bool(model.mixed_types()) if self._is_spin: self._model_output_def = ModelOutputDef( diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index f65e8f18d2..f99fffd488 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -20,6 +20,9 @@ ) log = logging.getLogger(__name__) +from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, +) from deepmd.dpmodel.utils.region import ( normalize_coord, ) @@ -32,9 +35,6 @@ from deepmd.pt_expt.model.graph_lower import ( graph_edge_dtype, ) -from deepmd.dpmodel.utils.neighbor_contract import ( - NeighborContract, -) from deepmd.utils.charge_state import ( CHARGE_STATE_TABLE_RANGES, ) @@ -1123,7 +1123,6 @@ def _spin_scheme(model_type: str | None) -> str | None: return None - def _neighbor_contract_for_model(model: torch.nn.Module) -> NeighborContract: """Resolve the export neighbor contract from a live model.""" getter = getattr(model, "get_neighbor_contract", None) diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index b4b3e3987f..46b59f32af 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -1,10 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """NeighborContract unit tests.""" -from __future__ import annotations +from __future__ import ( + annotations, +) import unittest -from unittest import mock +from unittest import ( + mock, +) from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, From 80e4886e5eedfccba4e238a18ae601f3657c97e9 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:21:27 +0800 Subject: [PATCH 03/12] fix(dpmodel): restore hybrid update_sel classmethod and merge capacity OR - Restore @classmethod on DescrptHybrid.update_sel (orphan decorator had demoted it to a plain function, breaking BaseDescriptor.update_sel). - Graph merge preserves OR of requires_capacity so sel:auto children still trigger discovery. - Dense merge no longer clears requires_capacity merely because a merged capacity exists; discovery stays if either sibling needs it. - Skip-capacity prepare_neighbors always computes min_nbor_dist (covers nested cutoffs such as DPA2 / hybrid without top-level rcut). - Extend neighbor_contract tests for all four cases. --- deepmd/dpmodel/descriptor/hybrid.py | 2 +- deepmd/dpmodel/model/dp_model.py | 8 +- deepmd/dpmodel/utils/neighbor_contract.py | 14 +- deepmd/pt/model/model/dp_model.py | 5 +- .../common/dpmodel/test_neighbor_contract.py | 129 ++++++++++++++++++ 5 files changed, 145 insertions(+), 13 deletions(-) diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index b0cf5fd06b..3f44c0bb3e 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -421,7 +421,6 @@ def call( out_gr = xp.concat(out_gr, axis=-2) if out_gr else None return out_descriptor, out_gr, out_g2, out_h2, out_sw - @classmethod @classmethod def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: """Merge child descriptor contracts; require a single representation.""" @@ -464,6 +463,7 @@ def prepare_jdata_for_neighbor_contract( del contract return out + @classmethod def update_sel( cls, train_data: DeepmdDataSystem, diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 894c5666bb..35bb262b5b 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -53,11 +53,9 @@ def prepare_neighbors( ) local_jdata_cpy["descriptor"] = descriptor_jdata if not contract.requires_capacity: - # Capacity discovery is skipped. Still report min neighbor distance - # when the descriptor exposes a cutoff (used by compression bounds). - rcut = descriptor_jdata.get("rcut") - if rcut is None: - return local_jdata_cpy, None + # Capacity discovery is skipped, but min neighbor distance is still + # needed for compression bounds. get_min_nbor_dist does not use + # rcut, so nested-cutoff descriptors (DPA2, hybrid, …) are covered. return local_jdata_cpy, float(UpdateSel().get_min_nbor_dist(train_data)) local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( train_data, type_map, descriptor_jdata diff --git a/deepmd/dpmodel/utils/neighbor_contract.py b/deepmd/dpmodel/utils/neighbor_contract.py index 0aa281c489..4ca521aec0 100644 --- a/deepmd/dpmodel/utils/neighbor_contract.py +++ b/deepmd/dpmodel/utils/neighbor_contract.py @@ -78,12 +78,12 @@ def dense( ) @classmethod - def graph(cls) -> NeighborContract: + def graph(cls, *, requires_capacity: bool = False) -> NeighborContract: """Build a carry-all graph contract with no neighbor capacity.""" return cls( version=NEIGHBOR_CONTRACT_VERSION, representation="graph", - requires_capacity=False, + requires_capacity=bool(requires_capacity), capacity=None, ) @@ -182,11 +182,15 @@ def merge(self, other: NeighborContract) -> NeighborContract: "Every child must share one representation (dense or graph)." ) if self.is_graph: - return NeighborContract.graph() + # Preserve OR of requires_capacity so a graph child that still + # needs discovery (e.g. DPA1 sel:auto) is not silently dropped. + return NeighborContract.graph( + requires_capacity=self.requires_capacity or other.requires_capacity + ) capacity = _merge_capacities(self.capacity, other.capacity) + # Discovery remains required if either side still needs it, even when + # the merged capacity is non-None (sibling may still be auto-sel). requires_capacity = self.requires_capacity or other.requires_capacity - if capacity is not None: - requires_capacity = False return NeighborContract.dense(capacity, requires_capacity=requires_capacity) diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 3237bd4e71..d9b8c44c6c 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -38,8 +38,9 @@ def prepare_neighbors( ) local_jdata_cpy["descriptor"] = descriptor_jdata if not contract.requires_capacity: - if descriptor_jdata.get("rcut") is None: - return local_jdata_cpy, None + # Capacity discovery is skipped, but min neighbor distance is still + # needed for compression bounds. get_min_nbor_dist does not use + # rcut, so nested-cutoff descriptors (DPA2, hybrid, …) are covered. return local_jdata_cpy, float(UpdateSel().get_min_nbor_dist(train_data)) local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( train_data, type_map, descriptor_jdata diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index 46b59f32af..2535d8f46c 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -19,6 +19,9 @@ from deepmd.dpmodel.descriptor.dpa4c import ( DescrptDPA4C, ) +from deepmd.dpmodel.descriptor.hybrid import ( + DescrptHybrid, +) from deepmd.dpmodel.model.dp_model import ( DPModelCommon, ) @@ -67,6 +70,32 @@ def test_merge_rejects_mixed_representation(self) -> None: with self.assertRaises(ValueError): NeighborContract.graph().merge(NeighborContract.dense([4])) + def test_graph_merge_preserves_requires_capacity(self) -> None: + # Hybrid: graph without sel + graph sel:auto must keep discovery. + left = NeighborContract.graph() + right = NeighborContract(representation="graph", requires_capacity=True) + merged = left.merge(right) + self.assertTrue(merged.is_graph) + self.assertTrue(merged.requires_capacity) + self.assertIsNone(merged.capacity) + # Symmetric + self.assertTrue(right.merge(left).requires_capacity) + # Both False stays False + self.assertFalse(left.merge(NeighborContract.graph()).requires_capacity) + + def test_dense_merge_preserves_discovery_when_sibling_needs_it(self) -> None: + # Hybrid dense: sel:auto + explicit sel must still require discovery. + auto = NeighborContract.dense(None, requires_capacity=True) + explicit = NeighborContract.dense([8, 16], requires_capacity=False) + merged = auto.merge(explicit) + self.assertTrue(merged.is_dense) + self.assertTrue(merged.requires_capacity) + self.assertEqual(merged.capacity, (8, 16)) + # Explicit + explicit: no discovery + both = NeighborContract.dense([4]).merge(NeighborContract.dense([8])) + self.assertFalse(both.requires_capacity) + self.assertEqual(both.capacity, (8,)) + def test_ensure_construction_sel(self) -> None: prepared = ensure_construction_sel({"type": "dpa1", "sel": "auto"}) self.assertEqual(prepared["sel"], GRAPH_NATIVE_CONSTRUCTION_SEL) @@ -194,5 +223,105 @@ def test_dense_still_calls_update_sel(self) -> None: self.assertEqual(min_dist, 0.8) + +class TestHybridUpdateSelClassmethod(unittest.TestCase): + def test_update_sel_is_classmethod(self) -> None: + self.assertTrue(isinstance(DescrptHybrid.__dict__["update_sel"], classmethod)) + + def test_hybrid_update_sel_dispatch(self) -> None: + jdata = { + "type": "hybrid", + "list": [ + { + "type": "dpa1", + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "concat", + "sel": "auto", + }, + { + "type": "dpa1", + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "concat", + }, + ], + } + train_data = mock.MagicMock() + with mock.patch( + "deepmd.dpmodel.descriptor.hybrid.BaseDescriptor.update_sel", + side_effect=lambda td, tm, child: ( + {**child, "sel": [4] if child.get("sel") == "auto" else child.get("sel", 1)}, + 0.3, + ), + ): + # Must be callable on the class without binding an instance. + updated, min_dist = DescrptHybrid.update_sel(train_data, ["O", "H"], jdata) + self.assertEqual(updated["list"][0]["sel"], [4]) + self.assertEqual(min_dist, 0.3) + + def test_hybrid_graph_merge_keeps_requires_capacity(self) -> None: + jdata = { + "type": "hybrid", + "list": [ + { + "type": "dpa1", + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "concat", + }, + { + "type": "dpa1", + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "concat", + "sel": "auto", + }, + ], + } + contract = DescrptHybrid.neighbor_contract_from_jdata(jdata) + self.assertTrue(contract.is_graph) + self.assertTrue(contract.requires_capacity) + + +class TestPrepareNeighborsMinNborDist(unittest.TestCase): + def test_skip_capacity_returns_min_nbor_dist_without_top_level_rcut(self) -> None: + # DPA2-like nested cutoffs under repinit/repformer only. + jdata = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa2", + "repinit": {"rcut": 6.0, "rcut_smth": 0.5, "nsel": 20}, + "repformer": {"rcut": 4.0, "rcut_smth": 0.5, "nsel": 10}, + }, + "fitting": {"type": "ener"}, + } + train_data = mock.MagicMock() + with ( + mock.patch( + "deepmd.dpmodel.model.dp_model.BaseDescriptor.neighbor_contract_from_jdata", + return_value=NeighborContract.graph(), + ), + mock.patch( + "deepmd.dpmodel.model.dp_model.BaseDescriptor.prepare_jdata_for_neighbor_contract", + side_effect=lambda j, c: dict(j), + ), + mock.patch( + "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel" + ) as update_sel, + mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, + ): + update_sel_cls.return_value.get_min_nbor_dist.return_value = 0.42 + updated, min_dist = DPModelCommon.prepare_neighbors( + train_data, ["O", "H"], jdata + ) + update_sel.assert_not_called() + update_sel_cls.return_value.get_min_nbor_dist.assert_called_once_with( + train_data + ) + self.assertEqual(min_dist, 0.42) + self.assertNotIn("rcut", updated["descriptor"]) + + if __name__ == "__main__": unittest.main() From 704705a56e91da49fc804bdacd830061541437d3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:23:00 +0000 Subject: [PATCH 04/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/dpmodel/test_neighbor_contract.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index 2535d8f46c..0fa56ce175 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -223,7 +223,6 @@ def test_dense_still_calls_update_sel(self) -> None: self.assertEqual(min_dist, 0.8) - class TestHybridUpdateSelClassmethod(unittest.TestCase): def test_update_sel_is_classmethod(self) -> None: self.assertTrue(isinstance(DescrptHybrid.__dict__["update_sel"], classmethod)) @@ -251,7 +250,10 @@ def test_hybrid_update_sel_dispatch(self) -> None: with mock.patch( "deepmd.dpmodel.descriptor.hybrid.BaseDescriptor.update_sel", side_effect=lambda td, tm, child: ( - {**child, "sel": [4] if child.get("sel") == "auto" else child.get("sel", 1)}, + { + **child, + "sel": [4] if child.get("sel") == "auto" else child.get("sel", 1), + }, 0.3, ), ): From d84fd7f317c37a206aa3c6a5af43abf0f0878255 Mon Sep 17 00:00:00 2001 From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:41:26 +0800 Subject: [PATCH 05/12] test(dpmodel): import unittest.mock without dual unittest import Avoid CodeQL dual import-style for unittest in NeighborContract tests. --- deepmd/lib | 1 + .../common/dpmodel/test_neighbor_contract.py | 31 +++++++++---------- 2 files changed, 15 insertions(+), 17 deletions(-) create mode 120000 deepmd/lib diff --git a/deepmd/lib b/deepmd/lib new file mode 120000 index 0000000000..eec07870f2 --- /dev/null +++ b/deepmd/lib @@ -0,0 +1 @@ +/workspace/deepmd-kit/deepmd/lib \ No newline at end of file diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index 0fa56ce175..8e03882194 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -6,9 +6,6 @@ ) import unittest -from unittest import ( - mock, -) from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, @@ -171,12 +168,12 @@ def test_graph_dpa1_does_not_call_update_sel(self) -> None: }, "fitting": {"type": "ener"}, } - train_data = mock.MagicMock() + train_data = unittest.mock.MagicMock() with ( - mock.patch( + unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel" ) as update_sel, - mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, + unittest.mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, ): update_sel_cls.return_value.get_min_nbor_dist.return_value = 0.5 updated, min_dist = DPModelCommon.prepare_neighbors( @@ -200,17 +197,17 @@ def test_dense_still_calls_update_sel(self) -> None: }, "fitting": {"type": "ener"}, } - train_data = mock.MagicMock() + train_data = unittest.mock.MagicMock() with ( - mock.patch( + unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel", return_value=({"type": "se_e2_a", "sel": [4, 8]}, 0.8), ) as update_sel, - mock.patch( + unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.neighbor_contract_from_jdata", return_value=NeighborContract.dense(None, requires_capacity=True), ), - mock.patch( + unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.prepare_jdata_for_neighbor_contract", side_effect=lambda j, c: dict(j), ), @@ -246,8 +243,8 @@ def test_hybrid_update_sel_dispatch(self) -> None: }, ], } - train_data = mock.MagicMock() - with mock.patch( + train_data = unittest.mock.MagicMock() + with unittest.mock.patch( "deepmd.dpmodel.descriptor.hybrid.BaseDescriptor.update_sel", side_effect=lambda td, tm, child: ( { @@ -298,20 +295,20 @@ def test_skip_capacity_returns_min_nbor_dist_without_top_level_rcut(self) -> Non }, "fitting": {"type": "ener"}, } - train_data = mock.MagicMock() + train_data = unittest.mock.MagicMock() with ( - mock.patch( + unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.neighbor_contract_from_jdata", return_value=NeighborContract.graph(), ), - mock.patch( + unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.prepare_jdata_for_neighbor_contract", side_effect=lambda j, c: dict(j), ), - mock.patch( + unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel" ) as update_sel, - mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, + unittest.mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, ): update_sel_cls.return_value.get_min_nbor_dist.return_value = 0.42 updated, min_dist = DPModelCommon.prepare_neighbors( From 263ae7300802cd08e9f5414d0450c39a7f4681f9 Mon Sep 17 00:00:00 2001 From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:41:41 +0800 Subject: [PATCH 06/12] chore: remove accidental deepmd/lib symlink from worktree --- deepmd/lib | 1 - 1 file changed, 1 deletion(-) delete mode 120000 deepmd/lib diff --git a/deepmd/lib b/deepmd/lib deleted file mode 120000 index eec07870f2..0000000000 --- a/deepmd/lib +++ /dev/null @@ -1 +0,0 @@ -/workspace/deepmd-kit/deepmd/lib \ No newline at end of file From 5de8cb64c5b55834b641515c0b0bbbbe27e77b58 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:46:02 +0000 Subject: [PATCH 07/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/dpmodel/test_neighbor_contract.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index 8e03882194..87c12ae534 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -173,7 +173,9 @@ def test_graph_dpa1_does_not_call_update_sel(self) -> None: unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel" ) as update_sel, - unittest.mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, + unittest.mock.patch( + "deepmd.dpmodel.model.dp_model.UpdateSel" + ) as update_sel_cls, ): update_sel_cls.return_value.get_min_nbor_dist.return_value = 0.5 updated, min_dist = DPModelCommon.prepare_neighbors( @@ -308,7 +310,9 @@ def test_skip_capacity_returns_min_nbor_dist_without_top_level_rcut(self) -> Non unittest.mock.patch( "deepmd.dpmodel.model.dp_model.BaseDescriptor.update_sel" ) as update_sel, - unittest.mock.patch("deepmd.dpmodel.model.dp_model.UpdateSel") as update_sel_cls, + unittest.mock.patch( + "deepmd.dpmodel.model.dp_model.UpdateSel" + ) as update_sel_cls, ): update_sel_cls.return_value.get_min_nbor_dist.return_value = 0.42 updated, min_dist = DPModelCommon.prepare_neighbors( From 71ff69aac533de84aafc1552542df215c9b9d690 Mon Sep 17 00:00:00 2001 From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 00:46:29 +0800 Subject: [PATCH 08/12] fix(dpmodel): correct NeighborContract for DPA2/4 and compositions - Short-circuit LinearEnergyAtomicModel to a dense contract when any child is dense-only, and give BaseAtomicModel a default get_neighbor_contract. - Read DPA2 nsel (not sel); keep auto-nsel discovery on the graph path; avoid injecting a top-level sel into DPA2 configs. - Keep DPA4 sel:auto on requires_capacity and inject construction sel only when discovery is skipped. - Stop swallowing get_sel failures as a fake graph contract; retain dense sel for nlist lowers in metadata/DeepEval. --- .../dpmodel/atomic_model/base_atomic_model.py | 15 +++++ .../atomic_model/linear_atomic_model.py | 13 ++++- deepmd/dpmodel/descriptor/dpa2.py | 57 ++++++++++++++++--- deepmd/dpmodel/descriptor/dpa4.py | 23 +++++++- deepmd/pt_expt/infer/deep_eval.py | 20 +++++-- deepmd/pt_expt/utils/serialization.py | 10 ++-- .../common/dpmodel/test_neighbor_contract.py | 47 +++++++++++++++ 7 files changed, 163 insertions(+), 22 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index 3a7922e8bd..b8639e0d3f 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -199,6 +199,21 @@ def uses_graph_lower(self) -> bool: """ return False + def get_neighbor_contract(self): + """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. diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 86f3b08848..b409b3da46 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -326,7 +326,18 @@ 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): - """Merge child neighbor contracts into one consistent contract.""" + """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:]: diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index a5a00d58ac..9f9cdf79f5 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -27,7 +27,9 @@ NetworkCollection, ) from deepmd.dpmodel.utils.neighbor_contract import ( + GRAPH_NATIVE_CONSTRUCTION_SEL, NeighborContract, + is_auto_sel, ) from deepmd.dpmodel.utils.network import ( Identity, @@ -758,21 +760,58 @@ def uses_graph_lower(self) -> bool: @classmethod def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: - """Match :meth:`uses_graph_lower` eligibility from config alone.""" - if local_jdata.get("use_three_body", False): - return NeighborContract.from_legacy_sel( - (local_jdata.get("repinit") or {}).get("sel", local_jdata.get("sel")) - ) + """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. + """ repinit = local_jdata.get("repinit") or {} repformer = local_jdata.get("repformer") or {} + + def _block_sel(block: dict): + if "nsel" in block: + return block.get("nsel") + return block.get("sel") + + def _dense_from_blocks() -> NeighborContract: + return NeighborContract.from_legacy_sel(_block_sel(repinit)) + + if local_jdata.get("use_three_body", False): + return _dense_from_blocks() if not repinit.get("set_davg_zero", True): - return NeighborContract.from_legacy_sel(repinit.get("sel")) + return _dense_from_blocks() if not repformer.get("set_davg_zero", True): - return NeighborContract.from_legacy_sel(repinit.get("sel")) + return _dense_from_blocks() tebd = repinit.get("tebd_input_mode", "concat") - if tebd in ("concat", "strip"): + 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() - return NeighborContract.from_legacy_sel(repinit.get("sel")) + 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``.""" + out = dict(local_jdata) + if not (contract.is_graph and not contract.requires_capacity): + return out + for key in ("repinit", "repformer"): + block = dict(out.get(key) or {}) + nsel = block.get("nsel", block.get("sel")) + if nsel is None or is_auto_sel(nsel): + block["nsel"] = int(GRAPH_NATIVE_CONSTRUCTION_SEL) + out[key] = block + return out def graph_type_embedding_table(self) -> Array: """Full type-embedding table consumed by the graph-route forward. diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index ef2e05b63c..277acd7906 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -65,6 +65,8 @@ ) from deepmd.dpmodel.utils.neighbor_contract import ( NeighborContract, + ensure_construction_sel, + is_auto_sel, ) from deepmd.dpmodel.utils.neighbor_graph import ( apply_pair_exclusion, @@ -2452,10 +2454,27 @@ def uses_graph_lower(self) -> bool: @classmethod def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: - """DPA4 defaults to the carry-all graph lower.""" - del local_jdata + """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 diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 84e4a84b41..d3cfbb29c1 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -415,8 +415,13 @@ def _init_from_model_json(self, model_json_str: str) -> None: self._dpmodel.get_sel(), requires_capacity=False ) ) + uses_graph = bool( + getattr(self._dpmodel, "uses_graph_lower", lambda: False)() + ) self._sel = ( - [] if self._neighbor_contract.is_graph else list(self._dpmodel.get_sel()) + [] + if self._neighbor_contract.is_graph and uses_graph + else list(self._dpmodel.get_sel()) ) self._mixed_types = bool(self._dpmodel.mixed_types()) if self._is_spin: @@ -449,8 +454,9 @@ def _init_from_metadata(self) -> None: self._rcut = float(self.metadata["rcut"]) self._type_map = list(self.metadata["type_map"]) self._neighbor_contract = NeighborContract.from_metadata(self.metadata) - if self._neighbor_contract.is_graph: - # Graph-native archives omit sel; keep an empty legacy list. + lower_kind = self.metadata.get("lower_input_kind", "nlist") + if self._neighbor_contract.is_graph and lower_kind != "nlist": + # Graph-routed archives omit sel; keep an empty legacy list. self._sel = [] elif "sel" in self.metadata: self._sel = [int(s) for s in self.metadata["sel"]] @@ -694,7 +700,11 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: if hasattr(model, "get_neighbor_contract") else NeighborContract.dense(model.get_sel(), requires_capacity=False) ) - self._sel = [] if self._neighbor_contract.is_graph else list(model.get_sel()) + self._sel = ( + [] + if self._neighbor_contract.is_graph and use_graph_lower + else list(model.get_sel()) + ) self._mixed_types = bool(model.mixed_types()) if self._is_spin: self._model_output_def = ModelOutputDef( @@ -727,7 +737,7 @@ def _load_pt(self, model_file: str, head: str | None = None) -> None: "neighbor_contract": self._neighbor_contract.to_dict(), **( {} - if self._neighbor_contract.is_graph + if self._neighbor_contract.is_graph and use_graph_lower else {"sel": model.get_sel(), "nnei": sum(model.get_sel())} ), "dim_fparam": model.get_dim_fparam(), diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index f99fffd488..352f04380b 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1144,10 +1144,7 @@ def _neighbor_contract_for_model(model: torch.nn.Module) -> NeighborContract: uses = getattr(model, "uses_graph_lower", None) if callable(uses) and bool(uses()): return NeighborContract.graph() - try: - return NeighborContract.dense(model.get_sel(), requires_capacity=False) - except Exception: - return NeighborContract.graph() + return NeighborContract.dense(model.get_sel(), requires_capacity=False) def _collect_metadata( @@ -1198,6 +1195,9 @@ def _collect_metadata( } ) neighbor_contract = _neighbor_contract_for_model(model) + # Omit sel only for graph-routed lowers. An nlist lower always needs the + # model's dense capacity even if a child described itself as graph. + omit_legacy_sel = neighbor_contract.is_graph and lower_kind != "nlist" meta = { "type_map": model.get_type_map(), "rcut": model.get_rcut(), @@ -1206,7 +1206,7 @@ def _collect_metadata( # not invent a dummy capacity. Dense models keep sel/nnei for old readers. **( {} - if neighbor_contract.is_graph + if omit_legacy_sel else { "sel": list(neighbor_contract.capacity) if neighbor_contract.capacity is not None diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index 87c12ae534..f32c381781 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -13,6 +13,9 @@ from deepmd.dpmodel.descriptor.dpa1 import ( DescrptDPA1, ) +from deepmd.dpmodel.descriptor.dpa2 import ( + DescrptDPA2, +) from deepmd.dpmodel.descriptor.dpa4c import ( DescrptDPA4C, ) @@ -285,6 +288,50 @@ def test_hybrid_graph_merge_keeps_requires_capacity(self) -> None: self.assertTrue(contract.requires_capacity) + def test_dpa2_auto_nsel_keeps_capacity_discovery(self) -> None: + jdata = { + "type": "dpa2", + "repinit": { + "rcut": 6.0, + "rcut_smth": 0.5, + "nsel": "auto", + "tebd_input_mode": "concat", + "set_davg_zero": True, + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": "auto", + "set_davg_zero": True, + }, + } + contract = DescrptDPA2.neighbor_contract_from_jdata(jdata) + self.assertTrue(contract.is_graph) + self.assertTrue(contract.requires_capacity) + + def test_dpa2_prepare_jdata_does_not_inject_top_level_sel(self) -> None: + jdata = { + "type": "dpa2", + "repinit": { + "rcut": 6.0, + "rcut_smth": 0.5, + "nsel": 20, + "tebd_input_mode": "concat", + "set_davg_zero": True, + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 10, + "set_davg_zero": True, + }, + } + contract = DescrptDPA2.neighbor_contract_from_jdata(jdata) + prepared = DescrptDPA2.prepare_jdata_for_neighbor_contract(jdata, contract) + self.assertNotIn("sel", prepared) + self.assertEqual(prepared["repinit"]["nsel"], 20) + + class TestPrepareNeighborsMinNborDist(unittest.TestCase): def test_skip_capacity_returns_min_nbor_dist_without_top_level_rcut(self) -> None: # DPA2-like nested cutoffs under repinit/repformer only. From af52f6989e15bc18f2e44f37db59a3e3b1d7d86b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:48:10 +0000 Subject: [PATCH 09/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt_expt/infer/deep_eval.py | 4 +--- source/tests/common/dpmodel/test_neighbor_contract.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index d3cfbb29c1..252fb0622e 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -415,9 +415,7 @@ def _init_from_model_json(self, model_json_str: str) -> None: self._dpmodel.get_sel(), requires_capacity=False ) ) - uses_graph = bool( - getattr(self._dpmodel, "uses_graph_lower", lambda: False)() - ) + uses_graph = bool(getattr(self._dpmodel, "uses_graph_lower", lambda: False)()) self._sel = ( [] if self._neighbor_contract.is_graph and uses_graph diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index f32c381781..ee8a8da11e 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -287,7 +287,6 @@ def test_hybrid_graph_merge_keeps_requires_capacity(self) -> None: self.assertTrue(contract.is_graph) self.assertTrue(contract.requires_capacity) - def test_dpa2_auto_nsel_keeps_capacity_discovery(self) -> None: jdata = { "type": "dpa2", From 02e306d3df018aa6e20716795673cb9259386daf Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> Date: Sat, 10 Oct 2026 10:45:52 +0800 Subject: [PATCH 10/12] fix(pt_expt): keep legacy sel/nnei for graph export C++ readers C++ DeepPotPTExpt/DeepSpinPTExpt still read metadata sel/nnei before lower_input_kind; omitting them on graph lowers caused map::at load failures. Always emit capacity or construction placeholders alongside neighbor_contract. Also: DPA2 dense discovery ORs requires_capacity across repinit/ repformer/three_body; graph construction placeholders satisfy repinit.nsel > repformer.nsel; migrate linear-model test mocks from update_sel to prepare_neighbors. --- deepmd/dpmodel/descriptor/dpa2.py | 67 ++++++++++++++++--- deepmd/pt_expt/infer/deep_eval.py | 3 +- deepmd/pt_expt/utils/serialization.py | 29 +++----- .../common/dpmodel/test_neighbor_contract.py | 44 ++++++++++++ .../tests/common/dpmodel/test_zbl_bridging.py | 6 +- .../pt_expt/model/test_get_model_bridging.py | 6 +- .../tests/pt_expt/model/test_linear_model.py | 6 +- 7 files changed, 124 insertions(+), 37 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 9f9cdf79f5..9257eebecf 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -766,19 +766,43 @@ def neighbor_contract_from_jdata(cls, local_jdata: dict) -> NeighborContract: 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): + 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: - return NeighborContract.from_legacy_sel(_block_sel(repinit)) + # 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 local_jdata.get("use_three_body", False): + if use_three_body: return _dense_from_blocks() if not repinit.get("set_davg_zero", True): return _dense_from_blocks() @@ -801,16 +825,39 @@ def prepare_jdata_for_neighbor_contract( local_jdata: dict, contract: NeighborContract, ) -> dict: - """Rewrite block ``nsel`` placeholders; never inject a top-level ``sel``.""" + """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 - for key in ("repinit", "repformer"): - block = dict(out.get(key) or {}) - nsel = block.get("nsel", block.get("sel")) - if nsel is None or is_auto_sel(nsel): - block["nsel"] = int(GRAPH_NATIVE_CONSTRUCTION_SEL) - out[key] = block + + 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: diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 252fb0622e..2d118ac903 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -454,7 +454,8 @@ def _init_from_metadata(self) -> None: self._neighbor_contract = NeighborContract.from_metadata(self.metadata) lower_kind = self.metadata.get("lower_input_kind", "nlist") if self._neighbor_contract.is_graph and lower_kind != "nlist": - # Graph-routed archives omit sel; keep an empty legacy list. + # Graph route ignores dense capacity; keep an empty legacy list even + # when metadata still carries a C++-compat placeholder sel/nnei. self._sel = [] elif "sel" in self.metadata: self._sel = [int(s) for s in self.metadata["sel"]] diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 352f04380b..6aee8d57ce 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1195,29 +1195,20 @@ def _collect_metadata( } ) neighbor_contract = _neighbor_contract_for_model(model) - # Omit sel only for graph-routed lowers. An nlist lower always needs the - # model's dense capacity even if a child described itself as graph. - omit_legacy_sel = neighbor_contract.is_graph and lower_kind != "nlist" + # Always emit legacy sel/nnei. C++ DeepPotPTExpt / DeepSpinPTExpt still + # read them before lower_input_kind (map::at on missing keys). Keep a + # construction placeholder for graph-native models until those readers + # migrate to neighbor_contract; neighbor_contract remains authoritative. + if neighbor_contract.capacity is not None: + legacy_sel = list(neighbor_contract.capacity) + else: + legacy_sel = list(model.get_sel()) meta = { "type_map": model.get_type_map(), "rcut": model.get_rcut(), "neighbor_contract": neighbor_contract.to_dict(), - # Legacy dense fields: omitted for graph-native models so metadata does - # not invent a dummy capacity. Dense models keep sel/nnei for old readers. - **( - {} - if omit_legacy_sel - else { - "sel": list(neighbor_contract.capacity) - if neighbor_contract.capacity is not None - else list(model.get_sel()), - "nnei": ( - sum(neighbor_contract.capacity) - if neighbor_contract.capacity is not None - else sum(model.get_sel()) - ), - } - ), + "sel": legacy_sel, + "nnei": int(sum(legacy_sel)), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), "dim_chg_spin": model.get_dim_chg_spin(), diff --git a/source/tests/common/dpmodel/test_neighbor_contract.py b/source/tests/common/dpmodel/test_neighbor_contract.py index ee8a8da11e..118f641b2a 100644 --- a/source/tests/common/dpmodel/test_neighbor_contract.py +++ b/source/tests/common/dpmodel/test_neighbor_contract.py @@ -330,6 +330,50 @@ def test_dpa2_prepare_jdata_does_not_inject_top_level_sel(self) -> None: self.assertNotIn("sel", prepared) self.assertEqual(prepared["repinit"]["nsel"], 20) + def test_dpa2_dense_mixed_auto_requires_capacity(self) -> None: + # Graph-ineligible (nonzero mean) with explicit repinit + auto repformer + # must still discover capacity for the unresolved block. + jdata = { + "type": "dpa2", + "repinit": { + "rcut": 6.0, + "rcut_smth": 0.5, + "nsel": 20, + "tebd_input_mode": "concat", + "set_davg_zero": False, + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": "auto", + "set_davg_zero": True, + }, + } + contract = DescrptDPA2.neighbor_contract_from_jdata(jdata) + self.assertTrue(contract.is_dense) + self.assertTrue(contract.requires_capacity) + + def test_dpa2_prepare_omitted_nsel_preserves_ordering(self) -> None: + jdata = { + "type": "dpa2", + "repinit": { + "rcut": 6.0, + "rcut_smth": 0.5, + "tebd_input_mode": "concat", + "set_davg_zero": True, + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "set_davg_zero": True, + }, + } + contract = DescrptDPA2.neighbor_contract_from_jdata(jdata) + self.assertTrue(contract.is_graph) + self.assertFalse(contract.requires_capacity) + prepared = DescrptDPA2.prepare_jdata_for_neighbor_contract(jdata, contract) + self.assertGreater(prepared["repinit"]["nsel"], prepared["repformer"]["nsel"]) + class TestPrepareNeighborsMinNborDist(unittest.TestCase): def test_skip_capacity_returns_min_nbor_dist_without_top_level_rcut(self) -> None: diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 7739cdd9f4..951c33946b 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -896,11 +896,13 @@ def test_update_sel_dispatches_and_skips_inner_child(self, monkeypatch) -> None: seen = [] - def _fake_update_sel(train_data, type_map, sub): + def _fake_prepare_neighbors(train_data, type_map, sub): seen.append(copy.deepcopy(sub)) return sub, 0.9 - monkeypatch.setattr(DPModelCommon, "update_sel", staticmethod(_fake_update_sel)) + monkeypatch.setattr( + DPModelCommon, "prepare_neighbors", staticmethod(_fake_prepare_neighbors) + ) cfg = model_args().normalize_value(_canonical_config(), trim_pattern="_*") updated, min_dist = BaseModel.update_sel(None, cfg["type_map"], cfg) assert min_dist == 0.9 diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index e4b9874519..d8de7526e7 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -313,11 +313,13 @@ def test_update_sel_skips_inner_potential_child(monkeypatch) -> None: seen = [] - def _fake_update_sel(train_data, type_map, sub): + def _fake_prepare_neighbors(train_data, type_map, sub): seen.append(copy.deepcopy(sub)) return sub, 0.9 - monkeypatch.setattr(DPModelCommon, "update_sel", staticmethod(_fake_update_sel)) + monkeypatch.setattr( + DPModelCommon, "prepare_neighbors", staticmethod(_fake_prepare_neighbors) + ) cfg = _canonical_config() updated, min_dist = LinearEnergyModel.update_sel(None, cfg["type_map"], cfg) assert min_dist == 0.9 diff --git a/source/tests/pt_expt/model/test_linear_model.py b/source/tests/pt_expt/model/test_linear_model.py index 8c95f3e69e..d7785ab775 100644 --- a/source/tests/pt_expt/model/test_linear_model.py +++ b/source/tests/pt_expt/model/test_linear_model.py @@ -460,8 +460,8 @@ def test_custom_weights(self) -> None: class TestLinearUpdateSel(unittest.TestCase): """Test that update_sel writes updated sub-model configs back.""" - @patch("deepmd.pt_expt.model.dp_linear_model.DPModelCommon.update_sel") - def test_updated_sel_written_back(self, mock_update_sel) -> None: + @patch("deepmd.pt_expt.model.dp_linear_model.DPModelCommon.prepare_neighbors") + def test_updated_sel_written_back(self, mock_prepare_neighbors) -> None: """Verify that update_sel returns configs with updated sel values.""" def side_effect(train_data, type_map, sub_jdata): @@ -469,7 +469,7 @@ def side_effect(train_data, type_map, sub_jdata): updated["descriptor"]["sel"] = 99 return updated, 0.5 - mock_update_sel.side_effect = side_effect + mock_prepare_neighbors.side_effect = side_effect local_jdata = { "type_map": ["O", "H"], From 04a61fd311a1fb2a6ad7a64184b61a931fca6fcd Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <198997215+OutisLi-Bot@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:34:19 +0800 Subject: [PATCH 11/12] fix(dpmodel): annotate get_neighbor_contract return types Satisfy ruff ANN201/ANN202 on NeighborContract publishers so pre-commit.ci passes on #6068. --- deepmd/dpmodel/atomic_model/base_atomic_model.py | 2 +- deepmd/dpmodel/atomic_model/dp_atomic_model.py | 2 +- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 2 +- deepmd/dpmodel/model/make_model.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index b8639e0d3f..c641b1d8d8 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -199,7 +199,7 @@ def uses_graph_lower(self) -> bool: """ return False - def get_neighbor_contract(self): + 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 diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 40133b7747..c9848be3b0 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -257,7 +257,7 @@ def get_sel(self) -> list[int]: """Get the neighbor selection.""" return self.descriptor.get_sel() - def get_neighbor_contract(self): + 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): diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index b409b3da46..15763ae8ca 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -325,7 +325,7 @@ 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): + 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 diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index b5f6f61a63..1c9a3aae5a 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -1220,7 +1220,7 @@ def get_sel(self) -> list[int]: """Returns the number of selected atoms for each type.""" return self.atomic_model.get_sel() - def get_neighbor_contract(self): + def get_neighbor_contract(self) -> "NeighborContract": """Return the neighbor representation contract of this model.""" return self.atomic_model.get_neighbor_contract() From 7883b19bc465e32968e9685ee4999f7b2ca220b2 Mon Sep 17 00:00:00 2001 From: OutisLi_Bot <198997215+OutisLi-Bot@users.noreply.github.com> Date: Tue, 13 Oct 2026 17:25:02 +0800 Subject: [PATCH 12/12] fix(dpmodel): import NeighborContract under TYPE_CHECKING for annotations Ruff F821 flagged the quoted return annotations on get_neighbor_contract after the previous annotation commit. Import the name in TYPE_CHECKING so pre-commit ruff passes without a runtime dependency cycle. --- deepmd/dpmodel/atomic_model/base_atomic_model.py | 3 +++ deepmd/dpmodel/atomic_model/dp_atomic_model.py | 3 +++ deepmd/dpmodel/atomic_model/linear_atomic_model.py | 6 ++++++ deepmd/dpmodel/model/make_model.py | 3 +++ 4 files changed, 15 insertions(+) diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index c641b1d8d8..d86268e768 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/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, ) diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index c9848be3b0..da89c5f36e 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/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, ) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 15763ae8ca..6b9dd2304f 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/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 diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 1c9a3aae5a..d9a575dfc2 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -14,6 +14,9 @@ from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, ) + from deepmd.dpmodel.utils.neighbor_contract import ( + NeighborContract, + ) import array_api_compat import numpy as np