diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index bf41735f89..ee406c11ae 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import dataclasses import functools import math from collections.abc import ( @@ -270,7 +269,11 @@ def forward_common_atomic( extended atom typs, shape: nf x nall for a type < 0 indicating the atomic is virtual. nlist - neighbor list, shape: nf x nloc x nsel + neighbor list, shape: nf x nloc x nsel. CONTRACT: model-level + ``pair_exclude_types`` is already folded in (excluded entries are + ``-1``) — exclusion is a nlist-BUILD transform (decision #18/A4) + applied by the NeighborList builders (Python) or + ``applyPairExclusionNlist`` (C++ ingestion), never re-applied here. mapping extended to local index mapping, shape: nf x nall fparam @@ -294,10 +297,14 @@ def forward_common_atomic( xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist) _, nloc, _ = nlist.shape atype = xp_take_first_n(extended_atype, 1, nloc) - if self.pair_excl is not None: - pair_mask = self.pair_excl.build_type_exclude_mask(nlist, extended_atype) - # exclude neighbors in the nlist - nlist = xp.where(pair_mask == 1, nlist, -1) + # NOTE: model-level ``pair_exclude_types`` is NOT applied here. It is a + # nlist-BUILD transform (decision #18/A4, same as the graph route): + # already folded into the nlist by the NeighborList builders (Python) + # or ``applyPairExclusionNlist`` (C++ ingestion); this method consumes + # a pre-excluded nlist. Fail-safe (eager only): guard against a caller + # that skipped the build seam, which would silently INCLUDE excluded + # pairs (fail-open). + self._assert_nlist_pair_excluded(nlist, extended_atype) ext_atom_mask = self.make_atom_mask(extended_atype) ret_dict = self.forward_atomic( @@ -326,11 +333,11 @@ def forward_common_atomic_graph( The node axis is flat ``(N,)`` (``N = sum(graph.n_node)``); masking and out-stat operate per node. Reuses :meth:`_finalize_atomic_ret`, so virtual-atom masking, ``atom_excl`` and ``apply_out_stat`` match the dense - path. Model-level ``pair_exclude_types`` is graph-native: when - ``self.pair_excl is not None``, an edge-keep mask is ANDed into - ``graph.edge_mask`` before the descriptor forward, so excluded type-pairs - contribute zero to the segment_sum. Descriptor-level ``exclude_types`` is - gated by ``uses_graph_lower()==False``. + path. Model-level ``pair_exclude_types`` is applied at graph BUILD time + (decision #18) — folded into ``graph.edge_mask`` by the NeighborGraph + builder (Python) or ``applyPairExclusion`` (C++), NOT here; this method + consumes a pre-excluded graph. Descriptor-level ``exclude_types`` is + handled inside the descriptor's ``call_graph`` (graph-native). Parameters ---------- @@ -356,14 +363,13 @@ def forward_common_atomic_graph( atype = xp.asarray(atype, device=array_api_compat.device(graph.edge_vec)) atom_mask = self.make_atom_mask(atype) # (N,) bool atype_clamped = xp.where(atom_mask, atype, xp.zeros_like(atype)) - if self.pair_excl is not None: - keep = self.pair_excl.build_edge_exclude_mask( - graph.edge_index, atype_clamped - ) - graph = dataclasses.replace( - graph, - edge_mask=graph.edge_mask * xp.astype(keep, graph.edge_mask.dtype), - ) + # NOTE: model-level ``pair_exclude_types`` is NOT applied here. It is a + # graph-BUILD transform (decision #18) already folded into + # ``graph.edge_mask`` by the NeighborGraph builder (Python) or + # ``applyPairExclusion`` (C++); this method consumes a pre-excluded graph. + # Fail-safe (eager only): guard against a caller that skipped the build + # seam, which would silently INCLUDE excluded pairs (fail-open). + self._assert_graph_pair_excluded(graph, atype_clamped) ret_dict = self.forward_atomic_graph( graph, atype_clamped, @@ -373,6 +379,92 @@ def forward_common_atomic_graph( ) return self._finalize_atomic_ret(ret_dict, atom_mask, atype) + def _assert_nlist_pair_excluded(self, nlist: Array, extended_atype: Array) -> None: + """Fail-safe: assert the nlist reaching the dense seam is pre-excluded. + + Decision #18/A4 makes the build seam the SOLE owner of model-level + ``pair_exclude_types``; this method no longer re-applies it. A caller + that skips the build seam would therefore silently INCLUDE excluded + pairs (fail-open) -- the dangerous direction for an exclusion feature. + This guard turns that into a loud error. + + Eager (numpy) only: it is a data-dependent check, so it is skipped under + ``torch.export`` / jax ``jit`` (where it cannot be traced) and in + compiled production. The C++ / exported-``.pt2`` ingestion paths are + covered by their own ingestion-site regression tests instead. + + Parameters + ---------- + nlist + neighbor list, shape: nf x nloc x nsel; ``-1`` marks empty slots. + extended_atype + extended atom types, shape: nf x nall. + + Raises + ------ + AssertionError + if a real (``>= 0``) neighbour carries an excluded type pair, i.e. + the build seam was skipped. Only when ``self.pair_excl`` is set and + ``nlist`` is a numpy array. + """ + if self.pair_excl is None or not array_api_compat.is_numpy_array(nlist): + return + xp = array_api_compat.array_namespace(nlist) + # keep == 0 marks an excluded type pair; a pre-excluded nlist has already + # set those neighbours to -1, so any *real* (>= 0) neighbour with keep==0 + # is a leak (the build seam was skipped). + keep = self.pair_excl.build_type_exclude_mask(nlist, extended_atype) + leaked = xp.astype(nlist != -1, xp.bool) & (keep == 0) + if bool(xp.any(leaked)): + n_leak = int(xp.sum(xp.astype(leaked, xp.int64))) + raise AssertionError( + f"forward_common_atomic received a nlist that is NOT " + f"pair-excluded: {n_leak} excluded-type neighbour(s) still " + "present. Model-level pair_exclude_types is a nlist-BUILD " + "transform (decision #18/A4) -- apply it at neighbor-list build " + "(Python builders / C++ applyPairExclusionNlist); this seam does " + "not re-apply it." + ) + + def _assert_graph_pair_excluded(self, graph: "NeighborGraph", atype: Array) -> None: + """Fail-safe graph analogue of :meth:`_assert_nlist_pair_excluded`. + + A pre-excluded graph has ``edge_mask == False`` on every excluded edge, + so any *active* edge whose type pair is excluded is a leak. Eager + (numpy) only, for the same reasons. + + Parameters + ---------- + graph + neighbor graph reaching the graph seam; only ``edge_mask`` and + ``edge_index`` are inspected. + atype + flat local node types, shape: N (clamped ``>= 0``). + + Raises + ------ + AssertionError + if an active edge (``edge_mask`` True) carries an excluded type + pair. Only when ``self.pair_excl`` is set and ``graph.edge_mask`` + is a numpy array. + """ + if self.pair_excl is None or not array_api_compat.is_numpy_array( + graph.edge_mask + ): + return + xp = array_api_compat.array_namespace(graph.edge_mask) + keep = self.pair_excl.build_edge_exclude_mask(graph.edge_index, atype) + leaked = xp.astype(graph.edge_mask, xp.bool) & (keep == 0) + if bool(xp.any(leaked)): + n_leak = int(xp.sum(xp.astype(leaked, xp.int64))) + raise AssertionError( + f"forward_common_atomic_graph received a graph that is NOT " + f"pair-excluded: {n_leak} excluded-type edge(s) still active. " + "Model-level pair_exclude_types is a graph-BUILD transform " + "(decision #18) -- apply it at graph build (Python builders / " + "C++ applyPairExclusion); this seam does not re-apply it." + ) + def _finalize_atomic_ret( self, ret_dict: dict, atom_mask: Array, atype: Array ) -> dict: @@ -673,6 +765,9 @@ def model_forward( self.get_sel(), mixed_types=self.mixed_types(), box=box, + # exclusion is a nlist-BUILD transform (decision #18/A4); + # forward_common_atomic consumes a pre-excluded nlist. + pair_excl=self.pair_excl, ) atomic_ret = self.forward_common_atomic( extended_coord, diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 3008daac84..3a47a7319c 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -356,6 +356,9 @@ def __init__( self.tebd_compress = False self.geo_compress = False self.compress = False + # When set, force the legacy dense lower even if the config would + # otherwise be graph-lower eligible (see ``disable_graph_lower``). + self._graph_lower_disabled = False def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -432,9 +435,11 @@ def uses_graph_lower(self) -> bool: The graph-native lower (``call_graph``) covers the factorizable path AND transformer attention (``attn_layer >= 0``, NeighborGraph PR-D) - with concat OR strip type-embedding. Remaining ineligible configs - (``exclude_types``, and compressed descriptors) fall back to the legacy - dense path, so those models keep working unchanged. + with concat OR strip type-embedding. ``exclude_types`` is fully + supported via + :func:`~deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`. + Compressed descriptors are the remaining ineligible config and fall + back to the legacy dense path, so those models keep working unchanged. Eligibility does NOT imply numerical interchangeability with the dense route for every config: with ``smooth_type_embedding=True`` @@ -442,16 +447,30 @@ def uses_graph_lower(self) -> bool: differs from the dense lower by up to ~1e-4 (see the Notes of :meth:`call_graph`). """ + if self._graph_lower_disabled: + return False # compressed descriptors have no graph kernel (geo/tebd tabulation is # dense-only); keep them on the legacy dense path. if self.compress: return False - # exclude_types stays dense (graph exclusion is owned elsewhere); strip is - # now graph-eligible (per-edge factorized embedding, no neighbor coupling). - return ( - self.se_atten.tebd_input_mode in ("concat", "strip") - and not self.se_atten.exclude_types - ) + # strip is graph-eligible (per-edge factorized embedding, no neighbor + # coupling); exclude_types is graph-native via ``apply_pair_exclusion`` + # (owned at this seam). Only compression / the disable flag force dense. + return self.se_atten.tebd_input_mode in ("concat", "strip") + + def disable_graph_lower(self) -> None: + """Force the legacy dense lower for this descriptor. + + This is an explicit opt-out knob used by contexts where the + graph-native lower is unsupported or undesirable (e.g. spin models, + whose carry-all routing diverges on sel-binding spin systems and + whose ``.pt2``/``.pte`` export trips a torch-inductor scatter/ + atomic_add CPU codegen assertion). After calling this, + :meth:`uses_graph_lower` returns ``False`` regardless of the + descriptor configuration. The flag is not serialized; it is + re-derived structurally at spin-model construction/deserialization. + """ + self._graph_lower_disabled = True def share_params( self, base_class: "DescrptDPA1", shared_level: int, resume: bool = False @@ -581,9 +600,9 @@ def call( nall = xp.reshape(coord_ext, (nlist.shape[0], -1)).shape[1] // 3 # graph-eligible configs route through the graph-native adapter (decision # #14: graph = single math source, dense call = thin adapter). Ineligible - # configs (exclude_types, compressed descriptors) and the ghost case with - # no mapping fall back to the legacy dense body. The graph needs `mapping` - # to fold ghosts to local owners; without it only nall == nloc is valid. + # configs (compressed descriptors) and the ghost case with no mapping + # fall back to the legacy dense body. The graph needs `mapping` to fold + # ghosts to local owners; without it only nall == nloc is valid. if self.uses_graph_lower() and (mapping is not None or nall == nloc): return self._call_graph_adapter(coord_ext, atype_ext, nlist, mapping) else: @@ -632,26 +651,16 @@ def _call_graph_adapter( The smooth switch function. shape: nf x nloc x nnei x 1 """ from deepmd.dpmodel.utils.neighbor_graph import ( - from_dense_quartet, + graph_from_dense_quartet, ) xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) - dev = array_api_compat.device(coord_ext) nf, nloc, nnei = nlist.shape - nall = xp.reshape(coord_ext, (nf, -1)).shape[1] // 3 - coord_ext_3 = xp.reshape(coord_ext, (nf, nall, 3)) - if mapping is None: - # default identity mapping (ext == loc, e.g. no-PBC nall == nloc) - mapping_g = xp.broadcast_to( - xp.arange(nall, dtype=xp.int64, device=dev)[None, :], (nf, nall) - ) - else: - mapping_g = xp.reshape(mapping, (nf, nall)) - graph = from_dense_quartet( - coord_ext_3, nlist, mapping_g, layout=None, compact=False + # shape-static graph + flat local center types from the dense quartet + # (shared with the input-stat graph path, see graph_from_dense_quartet). + graph, atype_local = graph_from_dense_quartet( + coord_ext, atype_ext, nlist, mapping ) - # local atom types, flat (nf * nloc,) - atype_local = xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (nf * nloc,)) grrg_flat, rot_mat_flat = self.call_graph( graph, atype_local, @@ -665,9 +674,10 @@ def _call_graph_adapter( grrg = xp.reshape(grrg_flat, (nf, nloc, *grrg_flat.shape[1:])) rot_mat = xp.reshape(rot_mat_flat, (nf, nloc, *rot_mat_flat.shape[1:])) # reconstruct the dense-shaped sw the dense way (env_mat switch masked - # where nlist == -1; the graph path forbids exclude_types, so nlist_mask - # == nlist != -1, matching DescrptBlockSeAtten.call). A dense-layout - # artifact tied to neighbor slots, which the graph does not carry. + # where nlist == -1 OR the neighbor pair is type-excluded, matching + # DescrptBlockSeAtten.call which erases excluded nlist entries to -1 + # before computing sw). A dense-layout artifact tied to neighbor slots, + # which the graph does not carry. _, _, sw = self.se_atten.env_mat.call( coord_ext, atype_ext, @@ -677,6 +687,12 @@ def _call_graph_adapter( ) nlist_mask = (nlist != -1)[:, :, :, None] sw = xp.where(nlist_mask, sw, xp.zeros_like(sw)) + if self.se_atten.exclude_types: + # additionally mask excluded type-pairs (mirrors the block's nlist + # erasure: excluded entries become -1 there, so sw is 0 for them). + exc_mask = self.se_atten.emask.build_type_exclude_mask(nlist, atype_ext) + exc_mask = xp.astype(exc_mask[:, :, :, None], sw.dtype) + sw = sw * exc_mask sw = xp.reshape(sw, (nf, nloc, nnei, 1)) return grrg, rot_mat, None, None, sw @@ -687,7 +703,7 @@ def _call_dense( nlist: Array, ) -> Array: """Legacy dense descriptor body (the ineligible ``call`` path: - compressed descriptors, exclude_types, or the no-mapping ghost case). + compressed descriptors or the no-mapping ghost case). Parameters ---------- @@ -1383,7 +1399,11 @@ def compute_input_stats( The path to the stat file. """ - env_mat_stat = EnvMatStatSe(self) + # dpa1's forward computes its env matrix through the NeighborGraph + # (from_dense_quartet -> edge_env_mat); run the input stat through the + # SAME path so stat and forward share one env-matrix implementation. + # Bit-identical to the dense EnvMat (see test_env_mat_stat_graph.py). + env_mat_stat = EnvMatStatSe(self, use_graph=True) if path is not None: path = path / env_mat_stat.get_hash() if path is None or not path.is_dir(): @@ -1756,9 +1776,10 @@ def call_graph( ----- Known limitations: - ``tebd_input_mode`` in {"concat", "strip"}; compressed descriptors stay dense; - - ``exclude_types`` is not yet supported and raises (lands in a later PR). + - ``exclude_types`` is applied graph-natively via ``apply_pair_exclusion``. """ from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, edge_env_mat, segment_sum, ) @@ -1767,11 +1788,6 @@ def call_graph( raise NotImplementedError( f"graph path does not support tebd_input_mode={self.tebd_input_mode!r}" ) - if self.exclude_types: - raise NotImplementedError( - "graph path does not yet apply exclude_types (NeighborGraph PR-A); " - "type exclusion lands in a later PR" - ) if type_embedding is None: raise ValueError("type_embedding is required for the graph path") xp = array_api_compat.array_namespace(graph.edge_vec) @@ -1779,9 +1795,15 @@ def call_graph( # N == sum(graph.n_node) by contract (atype is (N,)); use the static shape # value so the kernel stays jit/export-traceable (no concretize of n_node). n_total = atype.shape[0] + atype = xp.asarray(atype, device=dev) + # descriptor-level pair exclusion: same canonical transform as the + # model-level ``pair_exclude_types`` (decision #18). Masked edges + # contribute zero to every segment_sum below; the dense path's + # nlist-erasure + env-mat zeroing is reproduced exactly. + # apply_pair_exclusion is a no-op when self.emask has no exclusions. + graph = apply_pair_exclusion(graph, atype, self.emask) src = graph.edge_index[0, :] dst = graph.edge_index[1, :] - atype = xp.asarray(atype, device=dev) center_type = xp.take(atype, dst, axis=0) # (E,) nei_type = xp.take(atype, src, axis=0) # (E,) # per-edge env-mat 4-vector, normalized by the center (dst) atom type. diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 3933ec2533..84175549d1 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -11,6 +11,9 @@ from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) import array_api_compat import numpy as np @@ -85,6 +88,7 @@ def model_call_from_call_lower( coord_corr_for_virial: Array | None = None, charge_spin: Array | None = None, neighbor_list: NeighborList | None = None, + pair_excl: "PairExcludeMask | None" = None, ) -> dict[str, Array]: """Return model prediction from lower interface. @@ -109,6 +113,11 @@ def model_call_from_call_lower( historical behavior. An alternative strategy (e.g. an O(N) cell list) may be injected to speed up neighbor-list construction; it returns the same extended representation, so model outputs are unchanged. + pair_excl + Model-level pair-type exclusion mask. Exclusion is a nlist-BUILD + transform (decision #18/A4): it is folded into the nlist here, at the + build seam, and ``call_lower`` consumes a pre-excluded nlist without + re-applying it. Returns ------- @@ -122,7 +131,7 @@ def model_call_from_call_lower( del coord, box, fparam, aparam builder = neighbor_list if neighbor_list is not None else DefaultNeighborList() extended_coord, extended_atype, nlist, mapping = builder.build( - cc, atype, bb, rcut, sel + cc, atype, bb, rcut, sel, pair_excl=pair_excl ) extended_coord = extended_coord.reshape(nframes, -1, 3) if coord_corr_for_virial is not None: @@ -321,9 +330,8 @@ def call_common( The graph routes (``"dense"``/``"ase"``, and the pt_expt default-flip) require a ``mixed_types`` descriptor with a graph - lower (dpa1/se_atten with concat type embedding and no - ``exclude_types``; attention layers included). At non-binding - ``sel`` the graph matches the dense path exactly for the + lower (dpa1/se_atten with concat type embedding; attention layers included). + At non-binding ``sel`` the graph matches the dense path exactly for the non-smooth branch; at binding ``sel`` the carry-all graph keeps neighbors the dense path truncates, and for ``smooth_type_embedding=True`` the graph drops the dense @@ -386,6 +394,8 @@ def call_common( coord_corr_for_virial=coord_corr_for_virial, charge_spin=cs, neighbor_list=neighbor_list, + # exclusion is a nlist-BUILD transform (decision #18/A4) + pair_excl=getattr(self.atomic_model, "pair_excl", None), ) model_predict = self._output_type_cast(model_predict, input_prec) return model_predict @@ -469,10 +479,20 @@ def _call_common_graph( "neighbor_graph_method requires a mixed_types descriptor with a " "graph lower (e.g. dpa1 attn_layer=0)" ) + # Model-level ``pair_exclude_types`` is a graph-BUILD transform + # (decision #18): apply it here, at the seam where the NeighborGraph + # is constructed, so the graph lower / exported ``.pt2`` consumes an + # already-excluded ``edge_mask`` and never re-applies it. Mirrors the + # pt_expt eager path and the C++ ``applyPairExclusion`` at build. + pair_excl = getattr(self.atomic_model, "pair_excl", None) if method == "dense": - ng = build_neighbor_graph(cc, atype, bb, self.get_rcut()) + ng = build_neighbor_graph( + cc, atype, bb, self.get_rcut(), pair_excl=pair_excl + ) elif method == "ase": - ng = build_neighbor_graph_ase(cc, atype, bb, self.get_rcut()) + ng = build_neighbor_graph_ase( + cc, atype, bb, self.get_rcut(), pair_excl=pair_excl + ) else: raise ValueError( f"unknown neighbor_graph_method {method!r}; the dpmodel/jax backend " diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 373294bece..0d9da07355 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -63,6 +63,16 @@ def __init__( super().__init__() self.backbone_model = backbone_model self.spin = spin + # Spin graph-lower unsupported: carry-all routing diverges on + # sel-binding spin systems and spin export trips inductor scatter + # codegen. Re-derived structurally here so it survives both + # construction and serialize/deserialize round trips (the flag is + # not part of the serialized schema). + dp_atomic_model = self.backbone_model.get_dp_atomic_model() + if dp_atomic_model is not None: + descriptor = getattr(dp_atomic_model, "descriptor", None) + if descriptor is not None and hasattr(descriptor, "disable_graph_lower"): + descriptor.disable_graph_lower() self.ntypes_real = self.spin.ntypes_real self.virtual_scale_mask = self.spin.get_virtual_scale_mask() self.spin_mask = self.spin.get_spin_mask() @@ -170,6 +180,26 @@ def process_spin_input_lower( mapping_updated = None # extend the nlist nlist_updated = self.extend_nlist(extended_atype, nlist) + # This extension CREATES the virtual-atom nlist entries, so it is the + # BUILD site of the spin-extended nlist: fold the backbone's model-level + # pair exclusion in here (decision #18/A4 — the lower consumes a + # pre-excluded nlist and never re-applies it). No-op when the backbone + # has no pair_exclude_types. + pair_excl = getattr( + self.backbone_model.atomic_model + if hasattr(self.backbone_model, "atomic_model") + else self.backbone_model, + "pair_excl", + None, + ) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist_updated = apply_pair_exclusion_nlist( + nlist_updated, extended_atype_updated, pair_excl + ) return ( extended_coord_updated, extended_atype_updated, @@ -628,6 +658,12 @@ def call_common( charge_spin=charge_spin, do_atomic_virial=do_atomic_virial, coord_corr_for_virial=coord_corr_for_virial, + # Spin graph support is not yet implemented; the carry-all graph + # route diverges on sel-binding spin systems (virtual atoms double + # the density). Belt-and-braces: the backbone descriptor already + # has graph-lower disabled in ``__init__``, but force the legacy + # dense-nlist path here too until spin-graph support lands. + neighbor_graph_method="legacy", ) model_output_type = self.backbone_model.model_output_type() if "mask" in model_output_type: diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 3593af5c16..3e439f173e 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -48,6 +48,7 @@ make_multilayer_network, ) from .nlist import ( + apply_pair_exclusion_nlist, build_multiple_neighbor_list, build_neighbor_list, extend_coord_with_ghosts, @@ -94,6 +95,7 @@ "PairExcludeMask", "SameNlocBatchSampler", "aggregate", + "apply_pair_exclusion_nlist", "build_multiple_neighbor_list", "build_neighbor_graph", "build_neighbor_graph_ase", diff --git a/deepmd/dpmodel/utils/default_neighbor_list.py b/deepmd/dpmodel/utils/default_neighbor_list.py index 3628664c5a..d730b43669 100644 --- a/deepmd/dpmodel/utils/default_neighbor_list.py +++ b/deepmd/dpmodel/utils/default_neighbor_list.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Default all-pairs neighbor-list builder (historical deepmd behavior).""" +from typing import ( + TYPE_CHECKING, +) + import array_api_compat from deepmd.dpmodel.array_api import ( @@ -21,6 +25,11 @@ normalize_coord, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + class DefaultNeighborList(NeighborList): """All-pairs builder: replicate the cell into periodic images and rank by @@ -37,7 +46,36 @@ def build( rcut: float, sel: list[int], return_mode: str = "extended", + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Array, Array, Array, Array] | EdgeNeighborList: + """Build extended coordinates and a candidate neighbor list. + + Parameters + ---------- + coord : Array + Local coordinates, shape ``(nf, nloc, 3)`` or ``(nf, nloc*3)``. + atype : Array + Local atom types, shape ``(nf, nloc)``. + box : Array or None + Simulation cell, shape ``(nf, 3, 3)`` or ``(nf, 9)``; ``None`` + for non-periodic systems. + rcut : float + Cutoff radius. + sel : list[int] + Number of selected neighbors per type. + return_mode : str + Must be ``"extended"`` (the only mode this builder supports). + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list immediately after the geometric search by + :func:`~deepmd.dpmodel.utils.nlist.build_neighbor_list`. + + Returns + ------- + tuple[Array, Array, Array, Array] + ``(extended_coord, extended_atype, nlist, mapping)`` as documented + in :meth:`~deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. + """ if return_mode != "extended": raise NotImplementedError( "DefaultNeighborList only supports the extended-coordinate contract." @@ -54,7 +92,8 @@ def build( extended_coord, extended_atype, mapping = extend_coord_with_ghosts( coord_normalized, atype, box, rcut ) - # types are distinguished in the lower interface, so keep them merged here + # types are distinguished in the lower interface, so keep them merged here; + # pair_excl is forwarded so exclusion is applied at build time. nlist = build_neighbor_list( extended_coord, extended_atype, @@ -62,6 +101,7 @@ def build( rcut, sel, distinguish_types=False, + pair_excl=pair_excl, ) extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) return extended_coord, extended_atype, nlist, mapping diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 8d53602b18..543b918b42 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -25,6 +25,10 @@ from deepmd.dpmodel.utils.exclude_mask import ( PairExcludeMask, ) +from deepmd.dpmodel.utils.neighbor_graph import ( + edge_env_mat, + graph_from_dense_quartet, +) from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, ) @@ -143,12 +147,85 @@ class EnvMatStatSe(EnvMatStat): The descriptor of the model. """ - def __init__(self, descriptor: Union["Descriptor", "DescriptorBlock"]) -> None: + def __init__( + self, + descriptor: Union["Descriptor", "DescriptorBlock"], + use_graph: bool = False, + ) -> None: super().__init__() self.descriptor = descriptor self.last_dim = ( self.descriptor.ndescrpt // self.descriptor.nnei ) # se_r=1, se_a=4 + # ``use_graph`` computes the env matrix through the NeighborGraph path + # (``from_dense_quartet`` -> ``edge_env_mat``) instead of the dense + # ``EnvMat``, so the input stat runs the SAME machinery the dpa1 graph + # forward uses. It is BIT-IDENTICAL to the dense path (same neighbor + # set + padding, ``edge_env_mat`` mirrors ``EnvMat.call``, row-major + # ``(frame, center, slot)`` edges reshape 1:1 to ``(nf, nloc, nsel)``); + # only se_a-type (``last_dim == 4``) descriptors may opt in. + self.use_graph = use_graph + + def _graph_env_mat( + self, + extended_coord: Array, + extended_atype: Array, + mapping: Array, + nlist: Array, + ) -> Array: + """Env matrix via the NeighborGraph, shaped ``(nf, nloc, nsel, last_dim)``. + + Bit-identical to the dense ``EnvMat.call`` with zero mean / unit std: + ``from_dense_quartet(compact=False)`` reuses the same neighbor set and + padding (row-major ``(frame, center, slot)`` edges), ``edge_env_mat`` + mirrors ``EnvMat.call``, and padding / model-excluded edges (already + ``-1`` in the pre-excluded ``nlist``) carry ``edge_mask=False`` and are + zeroed -- so the ``(E, 4)`` output reshapes 1:1 back to the dense + ``(nf, nloc, nsel, 4)`` env-matrix tensor. + + Parameters + ---------- + extended_coord + extended coordinates, shape: nf x (nall x 3). + extended_atype + extended atom types, shape: nf x nall. + mapping + extended-to-local index mapping, shape: nf x nall. + nlist + pre-excluded neighbor list, shape: nf x nloc x nsel. + + Returns + ------- + env_mat + the environment matrix, shape: nf x nloc x nsel x last_dim. + """ + xp = array_api_compat.array_namespace(extended_coord, nlist) + dev = array_api_compat.device(extended_coord) + nframes, nloc, nsel = nlist.shape + ntypes = self.descriptor.get_ntypes() + graph, atype_local = graph_from_dense_quartet( + extended_coord, extended_atype, nlist, mapping + ) + # local center type per edge (dst is the local center index) + center_type = xp.take(atype_local, graph.edge_index[1, :], axis=0) + zero2 = xp.zeros((ntypes, 4), dtype=graph.edge_vec.dtype, device=dev) + one2 = xp.ones((ntypes, 4), dtype=graph.edge_vec.dtype, device=dev) + em = edge_env_mat( + graph.edge_vec, + center_type, + zero2, + one2, + self.descriptor.get_rcut(), + self.descriptor.get_rcut_smth(), + protection=self.descriptor.get_env_protection(), + edge_mask=graph.edge_mask, + return_sw=False, + ) # (E, 4) + # zero padding / model-excluded edges (edge_mask=False) so they count + # as 0 -- exactly like empty slots in the dense path. + em = em * xp.astype(graph.edge_mask[:, None], em.dtype) + # row-major (frame, center, slot) -> dense (nf, nloc, nsel, last_dim) + return xp.reshape(em, (nframes, nloc, nsel, self.last_dim)) def iter( self, data: list[dict[str, np.ndarray | list[tuple[int, int]]]] @@ -203,6 +280,11 @@ def iter( system["box"], ) nframes, nloc = atype.shape[:2] + pair_excl = None + if "pair_exclude_types" in system: + pair_excl = PairExcludeMask( + self.descriptor.get_ntypes(), system["pair_exclude_types"] + ) ( extended_coord, extended_atype, @@ -215,20 +297,34 @@ def iter( self.descriptor.get_sel(), mixed_types=self.descriptor.mixed_types(), box=box, + # Model-level pair exclusion is a nlist-BUILD transform + # (decision #18/A4): fold it in here so the input stat matches + # the model forward, which feeds the descriptor a pre-excluded + # nlist. Excluded pairs then behave exactly like empty slots + # (env_mat 0, still counted) -- identical to descriptor-level + # exclude_types, replacing the previous accumulation-deselect. + pair_excl=pair_excl, ) - env_mat_caller = EnvMat( - self.descriptor.get_rcut(), - self.descriptor.get_rcut_smth(), - protection=self.descriptor.get_env_protection(), - ) - env_mat, _, _ = env_mat_caller.call( - extended_coord, - extended_atype, - nlist, - zero_mean, - one_stddev, - radial_only, - ) + if self.use_graph: + # NeighborGraph env matrix (bit-identical to the dense EnvMat + # below): the SAME machinery the dpa1 graph forward uses. + env_mat = self._graph_env_mat( + extended_coord, extended_atype, mapping, nlist + ) + else: + env_mat_caller = EnvMat( + self.descriptor.get_rcut(), + self.descriptor.get_rcut_smth(), + protection=self.descriptor.get_env_protection(), + ) + env_mat, _, _ = env_mat_caller.call( + extended_coord, + extended_atype, + nlist, + zero_mean, + one_stddev, + radial_only, + ) # apply excluded_types exclude_mask = self.descriptor.emask.build_type_exclude_mask( nlist, extended_atype @@ -258,21 +354,11 @@ def iter( (-1, 1), ), ) - if "pair_exclude_types" in system: - pair_exclude_mask = PairExcludeMask( - self.descriptor.get_ntypes(), system["pair_exclude_types"] - ) - pair_exclude_mask.type_mask = xp.asarray( - pair_exclude_mask.type_mask, - device=array_api_compat.device(atype), - ) - # shape: (1, nloc, nnei) - exclude_mask = xp.reshape( - pair_exclude_mask.build_type_exclude_mask(nlist, extended_atype), - (1, nframes * nloc, -1), - ) - # shape: (ntypes, nloc, nnei) - type_idx = xp.logical_and(type_idx[..., None], exclude_mask) + # NOTE: model-level ``pair_exclude_types`` is NOT re-applied here. + # It is folded into the neighbor list at BUILD time above + # (decision #18/A4), so excluded pairs already have env_mat == 0 + # and are counted like empty slots -- the same treatment the model + # forward gives them. for type_i in range(self.descriptor.get_ntypes()): dd = env_mat[type_idx[type_i, ...]] dd = xp.reshape( diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 24fb090309..f29738c05c 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -15,6 +15,7 @@ from .builder import ( build_neighbor_graph, from_dense_quartet, + graph_from_dense_quartet, ) from .derivatives import ( edge_force_virial, @@ -28,6 +29,7 @@ from .graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, frame_id_from_n_node, node_validity_mask, pad_and_guard_edges, @@ -45,6 +47,7 @@ __all__ = [ "GraphLayout", "NeighborGraph", + "apply_pair_exclusion", "build_neighbor_graph", "build_neighbor_graph_ase", "center_edge_pairs", @@ -52,6 +55,7 @@ "edge_force_virial", "frame_id_from_n_node", "from_dense_quartet", + "graph_from_dense_quartet", "neighbor_graph_from_ijs", "node_validity_mask", "pad_and_guard_edges", diff --git a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py index 3b00ee6fac..fa163634a4 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py @@ -24,6 +24,9 @@ from .from_ijs import ( neighbor_graph_from_ijs, ) +from .graph import ( + apply_pair_exclusion, +) if TYPE_CHECKING: from deepmd.dpmodel.array_api import ( @@ -33,6 +36,7 @@ from .graph import ( GraphLayout, NeighborGraph, + PairExcludeMask, ) @@ -42,6 +46,9 @@ def build_neighbor_graph_ase( box: Array | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using ASE's O(N) cell-list search. @@ -66,6 +73,14 @@ def build_neighbor_graph_ase( cutoff radius. layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. Returns ------- @@ -136,6 +151,13 @@ def _to_cpu_numpy(x: Any) -> np.ndarray: i_all, j_all = i_all[keep], j_all[keep] S_all, nframe_all = S_all[keep], nframe_all[keep] - return neighbor_graph_from_ijs( + graph = neighbor_graph_from_ijs( i_all, j_all, S_all, coord, box, nframe_all, nloc, layout=layout ) + if pair_excl is not None: + import array_api_compat + + xp = array_api_compat.array_namespace(coord) + atype_flat = xp.reshape(xp.asarray(atype), (-1,)) + graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) + return graph diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index 71ca699e1b..54f70e0bc8 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -39,9 +39,14 @@ import array_api_compat +from deepmd.dpmodel.array_api import ( + xp_take_first_n, +) + from .graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, pad_and_guard_edges, ) @@ -50,6 +55,10 @@ Array, ) + from .graph import ( + PairExcludeMask, + ) + def from_dense_quartet( extended_coord: Array, @@ -197,12 +206,67 @@ def from_dense_quartet( ) +def graph_from_dense_quartet( + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None, +) -> tuple[NeighborGraph, Array]: + """Shape-static NeighborGraph + flat local center atype from a dense quartet. + + Convenience wrapper over :func:`from_dense_quartet` (``compact=False``) that + also fills the ``mapping is None`` identity case (no-PBC, ``nall == nloc``) + and derives the ghost-free flat local atom types ``(nf * nloc,)``. Shared by + the dpa1 dense->graph forward adapter (``DescrptDPA1._call_graph_adapter``) + and the input-stat graph path (``EnvMatStatSe._graph_env_mat``) so both build + the graph identically. + + Parameters + ---------- + coord_ext + Extended coordinates, ``(nf, nall x 3)`` or ``(nf, nall, 3)``. + atype_ext + Extended atom types, ``(nf, nall)``. + nlist + Dense neighbor list into the extended atoms, ``(nf, nloc, nsel)``; ``-1`` + is padding. + mapping + Extended -> local-owner index, ``(nf, nall)``; ``None`` means the + identity mapping (``nall == nloc``). + + Returns + ------- + graph + Shape-static :class:`NeighborGraph` over the LOCAL atoms. + atype_local + Flat local atom types, ``(nf * nloc,)``. + """ + xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist) + dev = array_api_compat.device(coord_ext) + nf, nloc, _ = nlist.shape + nall = atype_ext.shape[1] + coord_ext_3 = xp.reshape(coord_ext, (nf, nall, 3)) + if mapping is None: + # default identity mapping (ext == loc, e.g. no-PBC nall == nloc) + mapping_g = xp.broadcast_to( + xp.arange(nall, dtype=xp.int64, device=dev)[None, :], (nf, nall) + ) + else: + mapping_g = xp.reshape(mapping, (nf, nall)) + graph = from_dense_quartet(coord_ext_3, nlist, mapping_g, compact=False) + atype_local = xp.reshape(xp_take_first_n(atype_ext, 1, nloc), (nf * nloc,)) + return graph, atype_local + + def build_neighbor_graph( coord: Array, atype: Array, box: Array | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph DIRECTLY from coordinates (``dense`` search). @@ -221,6 +285,11 @@ def build_neighbor_graph( ``method`` key. Edges map every neighbor to its LOCAL owner (``src = mapping[neighbor]``), so the graph is ghost-free. + When ``pair_excl`` is given, :func:`apply_pair_exclusion` is called as a + post-process after the geometric search (the default path). A builder MAY + natively fuse the exclusion into its search in a future PR; the contract is + set-equality of valid-edge sets with the default post-process path. + Parameters ---------- coord @@ -236,6 +305,14 @@ def build_neighbor_graph( at non-binding ``sel``). layout edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. """ from deepmd.dpmodel.utils.nlist import ( extend_coord_with_ghosts, @@ -298,9 +375,13 @@ def build_neighbor_graph( edge_index, edge_vec, layout.edge_capacity, layout.min_edges ) n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) - return NeighborGraph( + graph = NeighborGraph( n_node=n_node, edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, ) + if pair_excl is not None: + atype_flat = xp.reshape(atype, (-1,)) + graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact) + return graph diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 0ce10efdf6..d8fcfb780e 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -25,6 +25,9 @@ from deepmd.dpmodel.array_api import ( Array, ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) @dataclass @@ -167,6 +170,130 @@ def frame_id_from_n_node(n_node: Array, n_total: int | None = None) -> Array: return xp.minimum(frame_id, xp.astype(last_frame, xp.int64)) +def apply_pair_exclusion( + graph: NeighborGraph, + atype: Array, + pair_excl: PairExcludeMask | None, + *, + compact: bool = False, +) -> NeighborGraph: + """Canonical pair-type exclusion transform (decision #18). + + ANDs the per-edge type keep-mask into ``graph.edge_mask`` so excluded + type pairs contribute exactly zero to every downstream ``segment_sum``. + The search stays purely geometric; this transform is applied ONCE at the + atomic-model seam (model-level ``pair_exclude_types``) and, for + descriptor-level ``exclude_types``, inside the descriptor's graph + forward. Identity (returns ``graph`` itself) when ``pair_excl`` is + ``None`` or empty. + + Parameters + ---------- + graph + The neighbor graph; only ``edge_mask`` (and, if ``compact=True``, + ``edge_index``, ``edge_vec``, ``angle_index``, ``angle_mask``) are + replaced. + atype + (N,) flat node types, clamped >= 0 (virtual atoms already handled + by the caller / the builders). + pair_excl + The ``PairExcludeMask`` holding the excluded (ti, tj) set. + compact + If ``False`` (default), only zero-out masked edges via ``edge_mask`` + (shape-static; the ONLY mode allowed in compiled / AOTI paths). + If ``True``, additionally drop masked edges so the returned graph + has no padding on the edge axis (data-dependent shape; eager / + dynamic-nedge only). When the graph carries angle fields, angles are + remapped onto the compacted edge axis and any angle whose constituent + edges were excluded is dropped. ``angle_index`` and ``angle_mask`` must + be a consistent pair (both set or both ``None``; matching ``A``); + anything else raises ``ValueError``. + + Returns + ------- + NeighborGraph + A ``dataclasses.replace`` copy (or the original ``graph`` on early + exit) with the exclusion applied. + + Notes + ----- + The C++ inference-path mirror is ``applyPairExclusion`` in + ``source/api_cc/include/commonPT.h``. It uses the same argument order + (edge_index, edge_mask, atype, ...) and the same variable names + (``type_ij``, ``keep``): it computes + ``type_ij = atype[dst]*(ntypes+1) + atype[src]`` and ANDs the flat + ``(ntypes+1)^2`` table lookup into ``edge_mask`` (mask-only mode; no + compact variant on the compiled path). + """ + import dataclasses + + if pair_excl is None or len(pair_excl.get_exclude_types()) == 0: + return graph + xp = array_api_compat.array_namespace(graph.edge_mask) + keep = pair_excl.build_edge_exclude_mask(graph.edge_index, atype) + out = dataclasses.replace( + graph, + edge_mask=xp.logical_and(graph.edge_mask, xp.astype(keep, xp.bool)), + ) + if compact: + # Angle fields are a coupled pair (produced together by the angle + # builder): both present or both None. Fail fast on any inconsistent + # state — a partial or shape-mismatched pair is a caller bug that would + # otherwise remap silently wrong. + has_ai = out.angle_index is not None + has_am = out.angle_mask is not None + if has_ai != has_am: + raise ValueError( + "apply_pair_exclusion(compact=True): angle_index and angle_mask " + "must both be set or both be None; got " + f"angle_index={'set' if has_ai else 'None'}, " + f"angle_mask={'set' if has_am else 'None'}." + ) + if has_ai: + if out.angle_index.ndim != 2 or out.angle_index.shape[0] != 2: + raise ValueError( + "apply_pair_exclusion(compact=True): angle_index must have " + f"shape (2, A); got {tuple(out.angle_index.shape)}." + ) + if out.angle_index.shape[1] != out.angle_mask.shape[0]: + raise ValueError( + "apply_pair_exclusion(compact=True): angle_index (2, A) and " + f"angle_mask (A,) disagree on A: {out.angle_index.shape[1]} " + f"vs {out.angle_mask.shape[0]}." + ) + (keep_idx,) = xp.nonzero(out.edge_mask) + fields = { + "edge_index": out.edge_index[:, keep_idx], + "edge_vec": xp.take(out.edge_vec, keep_idx, axis=0), + "edge_mask": xp.take(out.edge_mask, keep_idx, axis=0), + } + if has_ai: + # Angles reference PRE-compaction edge positions; remap them to the + # compacted axis and drop any angle whose constituent edges were + # excluded. ``new_pos`` maps old edge position -> new position via an + # exclusive prefix sum over the survivors (-1 for dropped edges). + surv = xp.astype(out.edge_mask, out.edge_index.dtype) # (E,) 0/1 + rank = xp.cumulative_sum(surv, axis=0) - surv # survivors before me + new_pos = xp.where(out.edge_mask, rank, xp.full_like(rank, -1)) + a_new = xp.take(new_pos, out.angle_index[0, :], axis=0) + b_new = xp.take(new_pos, out.angle_index[1, :], axis=0) + both_survive = xp.logical_and(a_new >= 0, b_new >= 0) + angle_keep = xp.logical_and( + xp.astype(out.angle_mask, xp.bool), both_survive + ) + (angle_keep_idx,) = xp.nonzero(angle_keep) + fields["angle_index"] = xp.stack( + [ + xp.take(a_new, angle_keep_idx, axis=0), + xp.take(b_new, angle_keep_idx, axis=0), + ], + axis=0, + ) + fields["angle_mask"] = xp.take(out.angle_mask, angle_keep_idx, axis=0) + out = dataclasses.replace(out, **fields) + return out + + def node_validity_mask(n_node: Array, n_total: int) -> Array: """Derive the (n_total,) real-vs-padding node mask from per-frame counts. diff --git a/deepmd/dpmodel/utils/neighbor_list.py b/deepmd/dpmodel/utils/neighbor_list.py index e63f5099dd..c9b6f5923b 100644 --- a/deepmd/dpmodel/utils/neighbor_list.py +++ b/deepmd/dpmodel/utils/neighbor_list.py @@ -14,6 +14,7 @@ dataclass, ) from typing import ( + TYPE_CHECKING, Literal, ) @@ -21,6 +22,11 @@ Array, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + @dataclass class EdgeNeighborList: @@ -69,6 +75,7 @@ def build( rcut: float, sel: list[int], return_mode: Literal["extended", "edges"] = "extended", + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Array, Array, Array, Array] | EdgeNeighborList: """Build the extended system and a candidate neighbor list. @@ -88,6 +95,11 @@ def build( ``"extended"`` returns the historical extended-coordinate quartet. ``"edges"`` returns :class:`EdgeNeighborList`, where ``edge_vec`` is the only geometric displacement consumed by the model. + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) by calling + :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. + Subclasses are expected to apply this filter before returning. Returns ------- diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 59e68a64a0..2697e57388 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + TYPE_CHECKING, Any, ) @@ -17,6 +18,11 @@ to_face_distance, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + def _is_ndtensorflow_namespace(xp: Any) -> bool: return getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow" @@ -46,6 +52,7 @@ def extend_input_and_build_neighbor_list( sel: list[int], mixed_types: bool = False, box: Array | None = None, + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(coord, atype) nframes, nloc = atype.shape[:2] @@ -66,11 +73,61 @@ def extend_input_and_build_neighbor_list( rcut, sel, distinguish_types=(not mixed_types), + pair_excl=pair_excl, ) extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) return extended_coord, extended_atype, mapping, nlist +def apply_pair_exclusion_nlist( + nlist: Array, + atype_ext: Array, + pair_excl: "PairExcludeMask | None", +) -> Array: + """Apply model-level pair-type exclusion to a dense neighbor list. + + Replaces excluded neighbor entries with ``-1`` so that downstream + descriptors see them as empty slots. Identity (returns ``nlist`` + unchanged) when *pair_excl* is ``None`` or its exclude-types list is + empty. + + This is the nlist-representation counterpart of + :func:`deepmd.dpmodel.utils.neighbor_graph.apply_pair_exclusion`. + + Parameters + ---------- + nlist : Array + Dense neighbor list of shape ``(nf, nloc, nnei)``. Entries equal + to ``-1`` indicate empty / padding slots. + atype_ext : Array + Extended atom types of shape ``(nf, nall)``. + pair_excl : PairExcludeMask or None + Exclusion mask object, or ``None`` / empty to skip. + + Returns + ------- + Array + Neighbor list of the same shape with excluded entries set to ``-1``. + Erasing ``-1`` entries a second time is a no-op (idempotent). + + Notes + ----- + Exclusion is a nlist-BUILD transform (decision #18/A4), same as the graph + route: it is applied ONCE where the nlist is built — the Python builders + (``build_neighbor_list(pair_excl=...)`` / the NeighborList strategies) or + the C++ ingestion seam (``applyPairExclusionNlist`` in + ``source/api_cc/include/commonPT.h``, same table lookup and variable + names). The lower interfaces (``call_lower`` / ``forward_common_atomic`` + and the exported artifacts) consume a pre-excluded nlist and never + re-apply it. + """ + if pair_excl is None or len(pair_excl.exclude_types) == 0: + return nlist + xp = array_api_compat.array_namespace(nlist, atype_ext) + pair_mask = pair_excl.build_type_exclude_mask(nlist, atype_ext) + return xp.where(pair_mask == 1, nlist, xp.full_like(nlist, -1)) + + ## translated from torch implementation by chatgpt def build_neighbor_list( coord: Array, @@ -79,6 +136,7 @@ def build_neighbor_list( rcut: float, sel: int | list[int], distinguish_types: bool = True, + pair_excl: "PairExcludeMask | None" = None, ) -> Array: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -100,6 +158,12 @@ def build_neighbor_list( types. distinguish_types : bool distinguish different types. + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) immediately after the + geometric search. This is a convenience shortcut for calling + :func:`apply_pair_exclusion_nlist` separately. ``None`` (default) + leaves the list unchanged. Returns ------- @@ -195,9 +259,8 @@ def build_neighbor_list( ) if distinguish_types: - return nlist_distinguish_types(nlist, atype, sel) - else: - return nlist + nlist = nlist_distinguish_types(nlist, atype, sel) + return apply_pair_exclusion_nlist(nlist, atype, pair_excl) def nlist_distinguish_types( diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index dba6d6946b..3d569b6a71 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -11,9 +11,17 @@ from collections.abc import ( Callable, ) +from typing import ( + TYPE_CHECKING, +) import tensorflow as tf +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + from deepmd.dpmodel.output_def import ( ModelOutputDef, ) @@ -52,6 +60,7 @@ def model_call_from_call_lower( fparam: tf.Tensor, aparam: tf.Tensor, do_atomic_virial: bool = False, + pair_excl: "PairExcludeMask | None" = None, ) -> dict[str, tf.Tensor]: """Return model prediction from lower interface.""" atype_shape = tf.shape(atype) @@ -78,6 +87,23 @@ def model_call_from_call_lower( # need to be distinguished here distinguish_types=False, ) + if pair_excl is not None and len(pair_excl.get_exclude_types()) > 0: + # Reuse the canonical dpmodel nlist-BUILD transform (decision #18/A4) + # via the vendored ``ndtensorflow`` array-API namespace -- the same way + # the TF2 backend (``deepmd/tf2``) runs dpmodel array-API code on + # TensorFlow. Unlike the neighbor-list *build* (see the docstring of + # ``jax2tf/nlist.py``), the exclusion has no data-dependent Python + # control flow: its only branch is on the static ``exclude_types`` + # config, so it traces cleanly under SavedModel export and does not + # need a hand-written TF twin. + from deepmd._vendors import ndtensorflow as ndtf + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist( + ndtf.asarray(nlist), ndtf.asarray(extended_atype), pair_excl + ).unwrap() extended_coord = tf.reshape(extended_coord, [nframes, -1, 3]) model_predict_lower = call_lower( extended_coord, diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index dd496dedf0..928798baeb 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -176,6 +176,12 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + # exclusion is a nlist-BUILD transform (decision #18/A4); + # the traced lower consumes a pre-excluded nlist. Guard + # atomic_model too: test doubles (DummyModel) lack it. + pair_excl=getattr( + getattr(model, "atomic_model", None), "pair_excl", None + ), ) return call diff --git a/deepmd/jax/jax_md/__init__.py b/deepmd/jax/jax_md/__init__.py index 45080dbf8c..13f23d212b 100644 --- a/deepmd/jax/jax_md/__init__.py +++ b/deepmd/jax/jax_md/__init__.py @@ -323,6 +323,18 @@ def _eval_with_jax_md_neighbor( displacement_fn, displacement_kwargs, ) + # Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision + # #18/A4): ``call_lower`` consumes a pre-excluded nlist and no longer + # re-applies it. The JAX-MD neighbor list is built without exclusion, so + # fold it in at this ingestion seam -- otherwise excluded pairs would be + # silently included (fail-open). + pair_excl = getattr(getattr(model, "atomic_model", None), "pair_excl", None) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) return model.call_lower( extended_coord, extended_atype, diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 0162850eaa..e8795e6a28 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -93,6 +93,25 @@ def __init__( self._has_default_fparam = has_default_fparam self.default_fparam = default_fparam self.numb_dos = numb_dos + # Model-level pair_exclude_types, rebuilt from the training config so + # the outer wrapper can fold it into the nlist at BUILD time (decision + # #18/A4) — the serialized StableHLO lower consumes a pre-excluded + # nlist and never re-applies it. + import json + + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + pet = [] + if model_def_script: + try: + pet = json.loads(model_def_script).get("pair_exclude_types", []) + except (ValueError, AttributeError): + pet = [] + self._pair_excl = ( + PairExcludeMask(len(type_map), [tuple(p) for p in pet]) if pet else None + ) def __call__( self, @@ -176,6 +195,9 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + # exclusion is a nlist-BUILD transform (decision #18/A4); the + # serialized StableHLO lower consumes a pre-excluded nlist. + pair_excl=self._pair_excl, ) def model_output_def(self) -> ModelOutputDef: diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index c77bc944b5..655bf5c581 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -17,9 +17,15 @@ Path, ) from typing import ( + TYPE_CHECKING, Any, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + import numpy as np import optax import orbax.checkpoint as ocp @@ -773,6 +779,7 @@ def _prepare_batch( box=jax_data["box"] if jax_data["find_box"] else None, fparam=jax_data.get("fparam", None), aparam=jax_data.get("aparam", None), + pair_excl=getattr(model.atomic_model, "pair_excl", None), ) return jax_data, extended_coord, extended_atype, nlist, mapping, fp, ap @@ -1007,6 +1014,7 @@ def prepare_input( box: np.ndarray | None = None, fparam: np.ndarray | None = None, aparam: np.ndarray | None = None, + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[ np.ndarray, np.ndarray, @@ -1038,6 +1046,9 @@ def prepare_input( # types will be distinguished in the lower interface, # so it doesn't need to be distinguished here distinguish_types=False, + # model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4); the lower consumes a pre-excluded nlist. + pair_excl=pair_excl, ) extended_coord = extended_coord.reshape(nframes, -1, 3) return extended_coord, extended_atype, nlist, mapping, fp, ap diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index 08fc308f72..df5f3dc119 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -51,6 +51,10 @@ Iterator, ) + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + @contextlib.contextmanager def _suppress_native_stderr() -> Iterator[None]: @@ -155,12 +159,30 @@ def build( rcut: float, sel: list[int], return_mode: str = "extended", + pair_excl: PairExcludeMask | None = None, ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: """Build the extended system and neighbor list. See :meth:`deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. The returned ``nlist`` is distance-sorted and truncated to ``sum(sel)``. + + Parameters + ---------- + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) by + :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. + ``NvNeighborList`` is CUDA-only; the ``pair_excl`` parameter is + accepted for API parity with the other strategies but cannot be + validated on a CPU-only machine. + ``return_mode='edges'`` does not support ``pair_excl``; a + :class:`NotImplementedError` is raised in that combination. """ + if return_mode == "edges" and pair_excl is not None: + raise NotImplementedError( + "pair_excl is not supported with return_mode='edges'; " + "use apply_pair_exclusion (graph variant) on the returned EdgeNeighborList." + ) device = coord.device nf, nloc = atype.shape[:2] target_neighbors = int(sum(sel)) @@ -207,6 +229,12 @@ def build( nlist = _truncate_to_sel_compiled( extended_coord, nlist, target_neighbors, float(rcut) ) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) return extended_coord, extended_atype, nlist, mapping diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index eeb97dcd2c..2e4f747ccb 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -499,8 +499,7 @@ def freeze( Lower-level export form: ``"nlist"`` (default, dense neighbor-list lower) or ``"graph"`` (NeighborGraph edge-list lower). ``"graph"`` is only valid for graph-eligible models (``mixed_types`` and ``uses_graph_lower``: - dpa1/se_atten with concat type embedding and no ``exclude_types``, - attention layers included) and selects the C++ graph inference path; + dpa1/se_atten with concat type embedding) and selects the C++ graph inference path; the per-atom virial is enabled for it (near-free in the graph path: one extra scatter off the shared single backward). NOTE: for ``smooth_type_embedding=True`` the carry-all graph attention @@ -570,7 +569,7 @@ def freeze( m.eval() # The graph lower is opt-in and only valid for graph-eligible models - # (dpa1 with concat tebd and no type exclusion; attention layers included + # (dpa1 with concat tebd, incl. attention layers and exclude_types # -- the carry-all pair enumeration exports via unbacked SymInts). Fail # fast with a clear message rather than emitting a broken .pt2. Enable the # per-atom virial for the graph form -- it is near-free there (one extra @@ -585,8 +584,7 @@ def freeze( raise ValueError( "lower_kind='graph' requires a graph-eligible model " "(mixed_types and a descriptor exposing uses_graph_lower()==True, " - "currently dpa1 with tebd_input_mode='concat' and no " - "exclude_types). Use lower_kind='nlist' for this model." + "currently dpa1 with tebd_input_mode='concat'). Use lower_kind='nlist' for this model." ) do_atomic_virial = True diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 58b62aaf56..ba209eba66 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -65,6 +65,9 @@ if TYPE_CHECKING: import ase.neighborlist + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) from deepmd.dpmodel.utils.neighbor_graph import ( NeighborGraph, ) @@ -954,6 +957,10 @@ def _build_nlist_native( sel = self._sel mixed_types = self._mixed_types + # Model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): fold it in here; the exported dense lower consumes a + # pre-excluded nlist and never re-applies it. + pair_excl = self._model_pair_excl() if self._nlist_builder is not None: # O(N) cell-list strategy (e.g. vesin): builds the same extended # representation. Match the native builder's type handling @@ -963,7 +970,7 @@ def _build_nlist_native( # type-distinguished nlist a non-mixed-type descriptor expects. The # main eval path is unaffected (its ``format_nlist`` re-formats). extended_coord, extended_atype, nlist, mapping = self._nlist_builder.build( - coords, atom_types, cells, rcut, sel + coords, atom_types, cells, rcut, sel, pair_excl=pair_excl ) if not mixed_types: nlist = nlist_distinguish_types(nlist, extended_atype, sel) @@ -988,6 +995,7 @@ def _build_nlist_native( rcut, sel, distinguish_types=not mixed_types, + pair_excl=pair_excl, ) extended_coord = extended_coord.reshape(nframes, -1, 3) return extended_coord, extended_atype, nlist, mapping @@ -1047,10 +1055,21 @@ def _build_nlist_ase( ext_atypes.append(ea) nlists.append(nl) mappings.append(mp) + extended_atype = np.stack(ext_atypes, axis=0) + nlist = np.stack(nlists, axis=0) + # Model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): fold it in here, like the native builder path. + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist( + nlist, extended_atype, self._model_pair_excl() + ) return ( np.stack(ext_coords, axis=0), - np.stack(ext_atypes, axis=0), - np.stack(nlists, axis=0), + extended_atype, + nlist, np.stack(mappings, axis=0), ) @@ -1795,19 +1814,26 @@ def _build_eval_graph( selection is a pure performance choice and results are unchanged. """ method = self._neighbor_graph_method + # Model-level ``pair_exclude_types`` is a graph-BUILD transform + # (decision #18): apply it here so the exported ``.pt2`` lower consumes a + # pre-excluded ``edge_mask`` and never re-applies it (mirrors the C++ + # ``applyPairExclusion`` and the eager dpmodel/pt_expt build path). + pair_excl = self._model_pair_excl() if method == "dense": from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph, ) - return build_neighbor_graph(coord_input, atom_types, box_input, self._rcut) + return build_neighbor_graph( + coord_input, atom_types, box_input, self._rcut, pair_excl=pair_excl + ) if method == "ase": from deepmd.dpmodel.utils.neighbor_graph import ( build_neighbor_graph_ase, ) return build_neighbor_graph_ase( - coord_input, atom_types, box_input, self._rcut + coord_input, atom_types, box_input, self._rcut, pair_excl=pair_excl ) if method in ("vesin", "nv"): cc = torch.as_tensor(coord_input, dtype=torch.float64, device=device) @@ -1824,17 +1850,51 @@ def _build_eval_graph( build_neighbor_graph_vesin, ) - return build_neighbor_graph_vesin(cc, aa, bb, self._rcut) + return build_neighbor_graph_vesin( + cc, aa, bb, self._rcut, pair_excl=pair_excl + ) from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, ) - return build_neighbor_graph_nv(cc, aa, bb, self._rcut) + return build_neighbor_graph_nv(cc, aa, bb, self._rcut, pair_excl=pair_excl) raise ValueError( f"unknown neighbor_graph_method {method!r}; " "use 'dense', 'ase', 'vesin', or 'nv'" ) + def _model_pair_excl(self) -> "PairExcludeMask | None": + """Model-level ``pair_exclude_types`` as a ``PairExcludeMask`` (or None). + + Applied at graph BUILD time (decision #18), NOT inside the exported + ``.pt2`` lower. Reads the excluded pairs from the loaded dpmodel (if any) + or the ``pair_exclude_types`` field in ``metadata.json``, and returns a + FRESH numpy-backed mask. + + A numpy ``type_mask`` converts cleanly onto whichever namespace/device the + builder's ``atype`` uses (dense/ase pass numpy; vesin/nv pass torch). The + dpmodel's own ``pair_excl`` is NOT reused: as a pt_expt module attribute + its ``type_mask`` is a torch (possibly CUDA) buffer, which cannot convert + to a numpy ``atype`` on the dense/ase build path. + + Returns + ------- + PairExcludeMask | None + The exclusion mask, or ``None`` when the model excludes no pairs. + """ + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + if self._dpmodel is not None: + pe = getattr(self._dpmodel.atomic_model, "pair_excl", None) + pet = pe.get_exclude_types() if pe is not None else [] + else: + pet = self.metadata.get("pair_exclude_types", []) + if not pet: + return None + return PairExcludeMask(len(self._type_map), [tuple(p) for p in pet]) + def _get_output_shape( self, odef: OutputVariableDef, nframes: int, natoms: int ) -> list[int]: diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py index ae2e83eada..a080600709 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -468,22 +468,29 @@ def _call_common_graph( "graph lower (e.g. dpa1 attn_layer=0)" ) rcut = self.get_rcut() + # Model-level pair_exclude_types is a graph-BUILD transform + # (decision #18): apply it here, at the single owning site, so the + # exported lower (forward_common_atomic_graph, which no longer + # re-applies it) consumes a pre-excluded edge_mask. + pair_excl = getattr(self.atomic_model, "pair_excl", None) if method == "dense": - ng = build_neighbor_graph(cc, atype, bb, rcut) + ng = build_neighbor_graph(cc, atype, bb, rcut, pair_excl=pair_excl) elif method == "ase": - ng = build_neighbor_graph_ase(cc, atype, bb, rcut) + ng = build_neighbor_graph_ase(cc, atype, bb, rcut, pair_excl=pair_excl) elif method == "vesin": from deepmd.pt_expt.utils.vesin_graph_builder import ( build_neighbor_graph_vesin, ) - ng = build_neighbor_graph_vesin(cc, atype, bb, rcut) + ng = build_neighbor_graph_vesin( + cc, atype, bb, rcut, pair_excl=pair_excl + ) elif method == "nv": from deepmd.pt_expt.utils.nv_graph_builder import ( build_neighbor_graph_nv, ) - ng = build_neighbor_graph_nv(cc, atype, bb, rcut) + ng = build_neighbor_graph_nv(cc, atype, bb, rcut, pair_excl=pair_excl) else: raise ValueError( f"unknown neighbor_graph_method {method!r}; " diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index d2868a4082..5c07d438f7 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -587,8 +587,8 @@ def _model_uses_graph_lower(model: torch.nn.Module) -> bool: :meth:`~deepmd.pt_expt.model.make_model.make_model..CM._resolve_graph_method` for ``neighbor_graph_method is None`` (the training default): a model is graph-eligible iff it is ``mixed_types`` AND its single descriptor reports - ``uses_graph_lower() == True`` (dpa1/se_atten with concat type embedding - and no ``exclude_types``; attention layers included). + ``uses_graph_lower() == True`` (dpa1/se_atten with concat type embedding; + attention layers included). When True the compiled lower must be the GRAPH ``forward_common_lower_graph`` so the compiled path matches eager training (which already default-flips to @@ -941,6 +941,9 @@ def forward( rcut, sel, distinguish_types=False, + # model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4); the compiled dense lower consumes a pre-excluded nlist. + pair_excl=getattr(self.original_model.atomic_model, "pair_excl", None), ) ext_coord = ext_coord.reshape(nframes, -1, 3) @@ -1154,8 +1157,12 @@ def _forward_graph( ) # Carry-all graph (dynamic E, no edge_capacity) — identical to the eager - # uncompiled ``_call_common_graph`` builder so the two paths match. - ng = build_neighbor_graph(coord_3d, atype, box_flat, rcut) + # uncompiled ``_call_common_graph`` builder so the two paths match. Model- + # level pair_exclude is a graph-BUILD transform (decision #18): fold it + # into edge_mask here so the compiled lower consumes a pre-excluded graph + # (the lower no longer re-applies it), matching the eager path exactly. + pair_excl = getattr(_model.atomic_model, "pair_excl", None) + ng = build_neighbor_graph(coord_3d, atype, box_flat, rcut, pair_excl=pair_excl) atype_flat = atype.reshape(nframes * nloc) # Lazy compile of the GRAPH lower (cached per structure key). diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py index 06a30c7dd4..fe588c94e7 100644 --- a/deepmd/pt_expt/utils/nv_graph_builder.py +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -22,14 +22,21 @@ ) from typing import ( + TYPE_CHECKING, Any, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + import torch from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, neighbor_graph_from_ijs, ) from deepmd.pt.utils.nv_nlist import ( @@ -204,6 +211,9 @@ def build_neighbor_graph_nv( box: Any | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using nvalchemiops' GPU cell list. @@ -219,6 +229,14 @@ def build_neighbor_graph_nv( cutoff radius. layout edge-axis length policy; ``None`` => dynamic with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. Returns ------- @@ -230,6 +248,12 @@ def build_neighbor_graph_nv( ------ ImportError if ``nvalchemi-toolkit-ops`` (CUDA) is not installed. + + Notes + ----- + The ``pair_excl`` path of this builder has no local oracle set-equality test + because nvalchemiops requires CUDA; the set-equality contract must be + validated on a GPU box (same pattern as :class:`~deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph_ase`). """ if not is_nv_available(): raise ImportError( @@ -275,6 +299,10 @@ def build_neighbor_graph_nv( center_local, src_local = center_local[keep], src_local[keep] shift, frame_idx = shift[keep], frame_idx[keep] - return neighbor_graph_from_ijs( + graph = neighbor_graph_from_ijs( center_local, src_local, shift, coord, box_out, frame_idx, nloc, layout=layout ) + if pair_excl is not None: + at_flat = torch.as_tensor(atype, device=device).reshape(-1) + graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) + return graph diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 9317e87f4d..c6d26dcde8 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -738,6 +738,27 @@ def _probe_has_message_passing(obj: object) -> bool | None: # "graph" → NeighborGraph (atype, n_node, edge_index, edge_vec, edge_mask) # The C++ loader branches on this to build the matching inputs. meta["lower_input_kind"] = "graph" if lower_kind == "graph" else "nlist" + + # Model-level pair-type exclusion (``pair_exclude_types``): a list of + # ``[ti, tj]`` type pairs whose interaction is dropped. Exclusion is a + # BUILD-time transform on BOTH routes (decision #18/A4): the exported + # lower (graph edge_mask / dense nlist) consumes a pre-excluded input and + # never re-applies it, so every feeder — Python builders, DeepEval, C++ + # ``applyPairExclusion`` / ``applyPairExclusionNlist`` — MUST fold the + # exclusion in at build. This metadata field is what lets external + # feeders (C++ ``DeepPotPTExpt::init``, metadata-only DeepEval) rebuild + # the mask. Descriptor-level ``exclude_types`` needs NO metadata: it is + # fully inside the compiled artifact. + pair_exclude_types: list[list[int]] = [] + for obj in ( + getattr(model, "atomic_model", None), + model, + ): + pet = getattr(obj, "pair_exclude_types", None) + if pet: + pair_exclude_types = [[int(ti), int(tj)] for (ti, tj) in pet] + break + meta["pair_exclude_types"] = pair_exclude_types return meta diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py index 3f86fc7b75..414b356f68 100644 --- a/deepmd/pt_expt/utils/vesin_graph_builder.py +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -19,15 +19,22 @@ ) from typing import ( + TYPE_CHECKING, Any, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + import array_api_compat import torch from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, NeighborGraph, + apply_pair_exclusion, neighbor_graph_from_ijs, ) from deepmd.pt_expt.utils.vesin_neighbor_list import ( @@ -89,11 +96,40 @@ def build_neighbor_graph_vesin( box: Any | None, rcut: float, layout: GraphLayout | None = None, + *, + pair_excl: PairExcludeMask | None = None, + compact: bool = False, ) -> NeighborGraph: """Build a CARRY-ALL NeighborGraph using vesin.torch's O(N) cell list. Mirrors :func:`deepmd.dpmodel.utils.neighbor_graph.build_neighbor_graph_ase` but runs on the input tensor's device via ``vesin.torch``. + + Parameters + ---------- + coord + (nf, nloc, 3) or (nf, nloc*3) local coordinates (torch tensor). + atype + (nf, nloc) local atom types; ``type < 0`` marks a virtual atom. + box + (nf, 3, 3) simulation cell, or ``None`` for non-periodic. + rcut + cutoff radius. + layout + edge-axis length policy; ``None`` => dynamic with ``min_edges`` guards. + pair_excl + Optional :class:`~deepmd.dpmodel.utils.neighbor_graph.graph.PairExcludeMask` + for model-level ``pair_exclude_types``. When given, + :func:`apply_pair_exclusion` is applied after the geometric search. ``None`` + (default) leaves all geometrically valid edges present. + compact + Passed to :func:`apply_pair_exclusion`; see that function for details. + Ignored when ``pair_excl`` is ``None``. + + Returns + ------- + graph + The carry-all :class:`NeighborGraph` over the LOCAL atoms. """ if not is_vesin_torch_available(): raise ImportError( @@ -162,6 +198,10 @@ def build_neighbor_graph_vesin( # (grad-carrying). Unlike the nv builder, vesin's cell list handles # out-of-cell (unwrapped) positions natively, so no normalize_coord is # needed and S is consistent with the original coords as searched. - return neighbor_graph_from_ijs( + graph = neighbor_graph_from_ijs( i_all, j_all, S_all, coord, box, nf_all, nloc, layout=layout ) + if pair_excl is not None: + at_flat = torch.as_tensor(atype, device=dev).reshape(-1) + graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact) + return graph diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 39a0dff883..0c7e5a22ec 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -19,6 +19,7 @@ """ from typing import ( + TYPE_CHECKING, Any, ) @@ -28,6 +29,9 @@ EdgeNeighborList, NeighborList, ) + +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_ij_shifts, merge_frame_edge_schemas, @@ -60,6 +64,7 @@ def build( rcut: float, sel: list[int], return_mode: str = "extended", + pair_excl: "PairExcludeMask | None" = None, ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: """Build the extended system + candidate neighbor list with vesin. @@ -67,7 +72,24 @@ def build( returned ``nlist`` is distance-sorted and truncated to ``sum(sel)`` (matching the default builder); the lower interface still re-formats / type-splits it. + + Parameters + ---------- + pair_excl : PairExcludeMask or None, optional + When provided, excluded type pairs are erased from the returned + neighbor list (entries set to ``-1``) by + :func:`~deepmd.dpmodel.utils.nlist.apply_pair_exclusion_nlist`. + This is the OWNING application site (decision #18/A4): the exported + lower (``forward_common_atomic``) no longer re-applies model-level + ``pair_exclude_types`` -- it is applied once here at nlist BUILD + time. ``return_mode='edges'`` does not support ``pair_excl``; a + :class:`NotImplementedError` is raised in that combination. """ + if return_mode == "edges" and pair_excl is not None: + raise NotImplementedError( + "pair_excl is not supported with return_mode='edges'; " + "use apply_pair_exclusion (graph variant) on the returned EdgeNeighborList." + ) is_numpy = not isinstance(coord, torch.Tensor) # vesin runs on the device of the inputs: numpy (the dpmodel backend) is # bridged through CPU torch; torch tensors stay on their own device. Pin @@ -146,6 +168,13 @@ def build( nlist = torch.stack(nlists, dim=0) mapping = torch.stack(mappings, dim=0) + if pair_excl is not None: + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + ) + + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) + if is_numpy: return ( extended_coord.detach().cpu().numpy(), diff --git a/deepmd/tf2/make_model.py b/deepmd/tf2/make_model.py index 458e0b5a06..e08890247e 100644 --- a/deepmd/tf2/make_model.py +++ b/deepmd/tf2/make_model.py @@ -3,6 +3,7 @@ Callable, ) from typing import ( + TYPE_CHECKING, Any, ) @@ -18,6 +19,7 @@ NeighborList, ) from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, nlist_distinguish_types, ) from deepmd.tf2.common import ( @@ -38,6 +40,11 @@ normalize_coord, ) +if TYPE_CHECKING: + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + def _unwrap_tuple(values: tuple[Array, ...]) -> tuple[tf.Tensor, ...]: return tuple(to_tf_tensor(value) for value in values) @@ -68,6 +75,7 @@ def model_call_from_call_lower( charge_spin: Array | None = None, neighbor_list: NeighborList | None = None, pass_lower_kwargs: bool = False, + pair_excl: "PairExcludeMask | None" = None, ) -> dict[str, Array]: """Return model prediction from lower interface. @@ -124,6 +132,13 @@ def model_call_from_call_lower( charge_spin=charge_spin, neighbor_list=neighbor_list, ) + if pair_excl is not None: + # Model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): fold it into the freshly built nlist here (already an + # ndtensorflow array, so no wrap/unwrap) via the canonical dpmodel + # helper, before the lower consumes it. Identity when nothing is + # excluded. + nlist = apply_pair_exclusion_nlist(nlist, extended_atype, pair_excl) lower_kwargs: dict[str, Any] = {"fparam": fp, "aparam": ap} if pass_lower_kwargs: if nlist_is_formatted: diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py index 5de81a2bdf..2bc7e0d42a 100644 --- a/deepmd/tf2/utils/serialization.py +++ b/deepmd/tf2/utils/serialization.py @@ -390,6 +390,12 @@ def call( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + # exclusion is a nlist-BUILD transform (decision #18/A4); + # the traced lower consumes a pre-excluded nlist. Guard + # atomic_model too: test doubles (DummyModel) lack it. + pair_excl=getattr( + getattr(model, "atomic_model", None), "pair_excl", None + ), ) ) diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 177c652bed..56bde5dbdc 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -157,7 +157,8 @@ In other backends, type embedding is within this descriptor with the {ref}`tebd_ TensorFlow and other backends have different implementations for {ref}`smooth_type_embedding `. The results are inconsistent when `smooth_type_embedding` is `true`. -In the pt_expt backend, graph-eligible descriptors (mixed types, `tebd_input_mode` `"concat"`, no descriptor-level `exclude_types` or compression) are evaluated by default through the carry-all neighbor-graph path instead of the legacy dense neighbor list. +In the pt_expt backend, graph-eligible descriptors (mixed types, `tebd_input_mode` `"concat"`, not compressed) are evaluated by default through the carry-all neighbor-graph path instead of the legacy dense neighbor list. +Descriptor-level `exclude_types` is supported on the graph path (it is folded into the neighbor graph), so it no longer disqualifies a descriptor from graph evaluation. The graph path considers all neighbors within the cutoff, so its result does not depend on {ref}`sel `. When `smooth_type_embedding` is `true` and {ref}`attn_layer ` is larger than 0 (the defaults), the dense path keeps `sel`-padding phantom terms in the attention softmax denominator while the graph path drops them, so checkpoints trained under the dense semantics shift by up to about 1e-4 in energy when evaluated on the graph path. Passing `neighbor_graph_method="legacy"` to the model forward (or the corresponding evaluation option) restores the dense-path numbers exactly. diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 6c30fae991..acda36dad3 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -361,6 +361,16 @@ class DeepPotPTExpt : public DeepPotBackend { // continue to work; GNN archives must be regenerated to opt into // the fail-fast guard against the silent-corruption bug. bool has_message_passing_ = false; + // Device-resident (ntypes+1)^2 model-level pair-type keep table, uploaded + // ONCE in ``init`` from the ``pair_exclude_types`` metadata field (see + // ``deepmd::buildPairExcludeTable``). An UNDEFINED tensor => no model-level + // exclusion (identity). The device is fixed at ``init`` (``gpu_id`` / + // ``gpu_enabled``), so the seam helpers ``index_select`` it directly with no + // per-step CPU clone / H2D upload. Exclusion is a BUILD-time transform + // (decision #18/A4): the C++ ingestion seam is the single application site + // (``applyPairExclusion`` graph / ``applyPairExclusionNlist`` dense); the + // exported .pt2 lowers consume pre-excluded inputs and never re-apply it. + torch::Tensor pair_exclude_table_; std::unique_ptr with_comm_tempfile_; std::unique_ptr with_comm_loader; diff --git a/source/api_cc/include/commonPT.h b/source/api_cc/include/commonPT.h index 02c25aa047..fd88d933ac 100644 --- a/source/api_cc/include/commonPT.h +++ b/source/api_cc/include/commonPT.h @@ -7,8 +7,10 @@ #include #include #include +#include #include #include +#include #include #include "common.h" @@ -462,6 +464,145 @@ inline GraphTensorPack buildGraphTensors( return pack; } +/** + * @brief Build the flat ``(ntypes+1)^2`` pair-type keep table. + * + * Inference-path mirror of the Python ``PairExcludeMask`` constructor + * (``deepmd/dpmodel/utils/exclude_mask.py``). The table is row-major over + * ``[tj][ti]`` (flat index ``tj * (ntypes+1) + ti``); an entry is ``0`` when + * the ordered pair ``(ti, tj)`` is excluded and ``1`` otherwise. Both ``(ti, + * tj)`` and ``(tj, ti)`` are inserted into the exclude set, so the table is + * symmetric. Type ``ntypes`` is the reserved virtual-atom row/column. + * + * Returns an empty vector when ``exclude_types`` is empty, so callers can treat + * an empty table as "no exclusion" (identity) just like the Python + * ``pair_excl is None`` early-exit. + * + * @param ntypes Number of real atom types. + * @param exclude_types List of excluded ``(ti, tj)`` type pairs. + */ +inline std::vector buildPairExcludeTable( + const int ntypes, const std::vector>& exclude_types) { + if (exclude_types.empty()) { + return {}; + } + const int n1 = ntypes + 1; + std::set> excl; + for (const auto& tt : exclude_types) { + excl.insert({tt.first, tt.second}); + excl.insert({tt.second, tt.first}); + } + // type_mask[tj][ti] == 0 iff (ti, tj) is excluded (mirrors the Python + // list comprehension in PairExcludeMask.__init__, reshape(-1)). + std::vector type_mask(static_cast(n1) * n1, 1); + for (int tj = 0; tj < n1; ++tj) { + for (int ti = 0; ti < n1; ++ti) { + if (excl.count({ti, tj})) { + type_mask[static_cast(tj) * n1 + ti] = 0; + } + } + } + return type_mask; +} + +/** + * @brief Graph pair-type exclusion: AND the per-edge keep-mask into + * ``edge_mask``. + * + * Inference-path twin of Python ``apply_pair_exclusion`` (mask-only mode) in + * ``deepmd/dpmodel/utils/neighbor_graph/graph.py`` + + * ``PairExcludeMask.build_edge_exclude_mask``. Kept side-by-side reviewable: + * same argument order (edge_index, edge_mask, atype, ...) and same variable + * names (``type_ij``, ``keep``). + * + * OWNERSHIP: exclusion is a BUILD-time transform (decision #18/A4). The + * exported ``.pt2`` graph lower consumes a pre-excluded ``edge_mask`` and + * never re-applies it, so the C++ ingestion seam is the sole owner: this helper + * is called once on each dispatched graph run path (single-rank and the + * non-message-passing multi-rank extended-region path), never twice on the same + * tensors. + * + * @param edge_index (2, E) int64 ``[src, dst]``; src = neighbor, dst = center. + * @param edge_mask (E,) bool real-vs-padding mask to be ANDed in place. + * @param atype (N,) int64 flat node types (clamped >= 0). + * @param type_mask_table Device-resident flat ``(ntypes+1)^2`` int32 keep table + * uploaded once in ``init`` (``buildPairExcludeTable`` + one H2D copy). An + * UNDEFINED tensor => identity (returns ``edge_mask``), mirroring the old + * empty-vector early-exit. Already on the model device (== ``edge_mask`` + * device), so ``index_select`` needs no per-call transfer. + * @param ntypes Number of real atom types. + * @return New ``edge_mask`` with excluded edges cleared. + */ +inline torch::Tensor applyPairExclusion(const torch::Tensor& edge_index, + const torch::Tensor& edge_mask, + const torch::Tensor& atype, + const torch::Tensor& type_mask_table, + const int ntypes) { + if (!type_mask_table.defined()) { + return edge_mask; + } + const auto src = edge_index.index({0}); // (E,) neighbour + const auto dst = edge_index.index({1}); // (E,) center + const auto src_t = atype.index_select(0, src); + const auto dst_t = atype.index_select(0, dst); + // type_ij = atype[dst] * (ntypes + 1) + atype[src] (matches Python) + const auto type_ij = dst_t * (ntypes + 1) + src_t; + const auto keep = type_mask_table.index_select(0, type_ij).to(torch::kBool); + return torch::logical_and(edge_mask, keep); +} + +/** + * @brief Dense-nlist pair-type exclusion: erase excluded neighbours to ``-1``. + * + * Inference-path twin of Python ``apply_pair_exclusion_nlist`` in + * ``deepmd/dpmodel/utils/nlist.py`` + + * ``PairExcludeMask.build_type_exclude_mask``. Same argument order (nlist, + * atype_ext, ...) and same variable names (``type_ij``, ``keep``). + * + * OWNERSHIP: exclusion is a BUILD-time transform (decision #18/A4). The + * exported ``.pt2`` dense lower consumes a pre-excluded nlist and never + * re-applies it, so the C++ ingestion seam is the sole owner: this helper is + * called once on each dispatched dense run path (single-rank ``run_model`` and + * multi-rank ``run_model_with_comm``), never twice on the same tensors. + * + * @param nlist (nf, nloc, nnei) int64 neighbour list; ``-1`` == empty slot. + * @param atype_ext (nf, nall) int64 extended atom types. + * @param type_mask_table Device-resident flat ``(ntypes+1)^2`` int32 keep table + * uploaded once in ``init``. An UNDEFINED tensor => identity (returns + * ``nlist``). Already on the model device (== ``nlist`` device), so + * ``index_select`` needs no per-call transfer. + * @param ntypes Number of real atom types. + * @return New neighbour list with excluded entries set to ``-1``. + */ +inline torch::Tensor applyPairExclusionNlist( + const torch::Tensor& nlist, + const torch::Tensor& atype_ext, + const torch::Tensor& type_mask_table, + const int ntypes) { + if (!type_mask_table.defined()) { + return nlist; + } + const std::int64_t nf = nlist.size(0); + const std::int64_t nloc = nlist.size(1); + const std::int64_t nnei = nlist.size(2); + const std::int64_t nall = atype_ext.size(1); + // center types: first nloc extended atoms. type_i = atype * (ntypes + 1). + const auto type_i = atype_ext.slice(1, 0, nloc) * (ntypes + 1); // (nf, nloc) + // append virtual atom of type ntypes; map -1 neighbours to it. + const auto ae = torch::cat( + {atype_ext, torch::full({nf, 1}, ntypes, atype_ext.options())}, 1); + const auto nlist_for_type = + torch::where(nlist == -1, torch::full_like(nlist, nall), nlist); + const auto type_j = torch::gather( + ae.unsqueeze(1).expand({nf, nloc, nall + 1}), 2, nlist_for_type); + // type_ij = type_i * (ntypes + 1) + type_j (matches Python: type_i already + // scaled above; here just add the neighbour type). + const auto type_ij = type_i.unsqueeze(2) + type_j; // (nf, nloc, nnei) + const auto keep = type_mask_table.index_select(0, type_ij.reshape({-1})) + .reshape({nf, nloc, nnei}); + return torch::where(keep == 1, nlist, torch::full_like(nlist, -1)); +} + /** * @brief Remap NeighborGraph (graph-schema) public outputs onto the dense * internal-key layout the rest of ``compute`` consumes. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index 9088762fcb..a60bd681e4 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -201,6 +201,39 @@ void DeepPotPTExpt::init(const std::string& model, // scripts and carry the explicit value. has_message_passing_ = metadata.obj_val.count("has_message_passing") && metadata["has_message_passing"].as_bool(); + + // Model-level pair-type exclusion table. ``pair_exclude_types`` is a list + // of [ti, tj] pairs; rebuild the flat (ntypes+1)^2 keep table exactly like + // the Python ``PairExcludeMask`` ctor. Exclusion is a BUILD-time transform + // (decision #18/A4): the C++ ingestion seam is the single application site + // (applyPairExclusion graph / applyPairExclusionNlist dense); the exported + // lowers consume pre-excluded inputs and never re-apply it. + // + // Upload the table to the model device ONCE here (device is fixed by + // ``gpu_id`` / ``gpu_enabled``), so the per-step seam helpers only + // ``index_select`` it -- no per-``compute()`` CPU clone + H2D copy. Leave + // ``pair_exclude_table_`` undefined (=> identity) when there is no exclusion. + { + std::vector> pair_exclude_types; + if (metadata.obj_val.count("pair_exclude_types")) { + for (const auto& v : metadata["pair_exclude_types"].as_array()) { + pair_exclude_types.emplace_back(v[0].as_int(), v[1].as_int()); + } + } + std::vector tbl = + deepmd::buildPairExcludeTable(ntypes, pair_exclude_types); + if (!tbl.empty()) { + torch::Device device(torch::kCUDA, gpu_id); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + } + pair_exclude_table_ = + torch::from_blob(tbl.data(), {static_cast(tbl.size())}, + torch::TensorOptions().dtype(torch::kInt32)) + .clone() + .to(device); + } + } if (has_comm_artifact_) { try { // Extract the nested ``extra/forward_lower_with_comm.pt2`` into a @@ -823,9 +856,20 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor, charge_spin_tensor, comm_tensors); } } else { + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported dense lower consumes a pre-excluded nlist and + // never re-applies it. The multi-rank (with-comm) dense route shares the + // same dense nlist as the single-rank path below, so it applies the SAME + // seam -- otherwise a message-passing .pt2 with pair_exclude_types would + // silently include excluded pairs on the with-comm path (multi-rank != + // single-rank). The cross-rank ghost exchange happens inside + // run_model_with_comm and does not change the nlist's meaning, so + // pre-excluding it is correct per rank. + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = run_model_with_comm( - coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, - fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); + coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, fparam_tensor, + aparam_tensor, charge_spin_tensor, comm_tensors); } } else { if (lower_input_is_edge_) { @@ -876,14 +920,27 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, torch::full({1}, n_node_count, int_option).to(device); at::Tensor node_atype = atype_Tensor.slice(1, 0, n_node_count).reshape({n_node_count}); + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported graph lower consumes a pre-excluded edge_mask + // and never re-applies it; this is the single application site on the + // C++ graph route. + const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( + edge_tensors.edge_index, edge_tensors.edge_mask, node_atype, + pair_exclude_table_, ntypes); flat_outputs = run_model_graph(node_atype, n_node_tensor, edge_tensors.edge_index, - edge_tensors.edge_vec, edge_tensors.edge_mask, - fparam_tensor, aparam_tensor, charge_spin_tensor); + edge_tensors.edge_vec, graph_edge_mask, fparam_tensor, + aparam_tensor, charge_spin_tensor); } else { - flat_outputs = run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, - mapping_tensor, fparam_tensor, aparam_tensor, - charge_spin_tensor); + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported dense lower consumes a pre-excluded nlist and + // never re-applies it. Single-rank dense application site; the + // multi-rank (with-comm) dense sibling above applies the same seam. + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + firstneigh_tensor, atype_Tensor, pair_exclude_table_, ntypes); + flat_outputs = + run_model(coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, + fparam_tensor, aparam_tensor, charge_spin_tensor); } } @@ -1239,13 +1296,26 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, edge_tensors.edge_index_ext, edge_tensors.edge_mask, fparam_tensor, aparam_tensor, charge_spin_tensor); } else if (lower_input_is_graph_) { + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported graph lower consumes a pre-excluded edge_mask and + // never re-applies it; this is the single application site on the C++ + // graph route. + const at::Tensor graph_edge_mask = deepmd::applyPairExclusion( + graph_tensors.edge_index, graph_tensors.edge_mask, graph_tensors.atype, + pair_exclude_table_, ntypes); flat_outputs = run_model_graph( graph_tensors.atype, graph_tensors.n_node, graph_tensors.edge_index, - graph_tensors.edge_vec, graph_tensors.edge_mask, fparam_tensor, - aparam_tensor, charge_spin_tensor); + graph_tensors.edge_vec, graph_edge_mask, fparam_tensor, aparam_tensor, + charge_spin_tensor); } else { + // Model-level pair exclusion is a BUILD-time transform (decision + // #18/A4): the exported dense lower consumes a pre-excluded nlist and + // never re-applies it; this is the single application site on the C++ + // dense route. + const at::Tensor excl_nlist = deepmd::applyPairExclusionNlist( + nlist_tensor, atype_Tensor, pair_exclude_table_, ntypes); flat_outputs = - run_model(coord_Tensor, atype_Tensor, nlist_tensor, mapping_tensor, + run_model(coord_Tensor, atype_Tensor, excl_nlist, mapping_tensor, fparam_tensor, aparam_tensor, charge_spin_tensor); } diff --git a/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc new file mode 100644 index 0000000000..0b2934ecf9 --- /dev/null +++ b/source/api_cc/tests/test_deeppot_dpa1_pairexcl_ptexpt.cc @@ -0,0 +1,179 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test the C++ model-level pair-exclusion ingestion seam of the pt_expt +// backend (Task A3/A4). Two DPA1(attn_layer=0) models with identical weights +// and model-level pair_exclude_types=[[0,1]] are exported, one through each C++ +// ingestion route: +// - deeppot_dpa1_pairexcl_graph.pt2 -> applyPairExclusion (graph route) +// - deeppot_dpa1_pairexcl_nlist.pt2 -> applyPairExclusionNlist (dense route) +// A no-exclusion baseline (deeppot_dpa1_pairexcl_none.pt2, empty exclude table) +// exercises the identity/pre-change branch of both helpers. +// +// Exclusion is a BUILD-time transform (decision #18/A4): the exported .pt2 +// lowers consume pre-excluded inputs (graph edge_mask / dense nlist) and never +// re-apply it, so the C++ helpers here are the SINGLE application site — these +// tests are what proves they are load-bearing. The reference values +// (.expected sidecars) come from the Python DeepEval of the SAME .pt2, so a +// 1e-10 match validates the whole chain (pair_exclude_types metadata +// round-trip + init table build + build-time apply + compiled math). A +// separate assertion (excluded energy != baseline energy) proves the exclusion +// is genuinely active and not silently dropped. +#include + +#include +#include + +#include "DeepPot.h" +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "test_utils.h" + +namespace { +constexpr const char* kGraphModel = + "../../tests/infer/deeppot_dpa1_pairexcl_graph.pt2"; +constexpr const char* kNlistModel = + "../../tests/infer/deeppot_dpa1_pairexcl_nlist.pt2"; +constexpr const char* kNoneModel = + "../../tests/infer/deeppot_dpa1_pairexcl_none.pt2"; +constexpr const char* kGraphRef = + "../../tests/infer/deeppot_dpa1_pairexcl_graph.expected"; +constexpr const char* kNlistRef = + "../../tests/infer/deeppot_dpa1_pairexcl_nlist.expected"; +} // namespace + +template +class TestInferDpa1PairExclPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + + // Excluded models (one per ingestion route) + no-exclusion baseline. + static deepmd::DeepPot dp_graph; + static deepmd::DeepPot dp_nlist; + static deepmd::DeepPot dp_none; + + static void SetUpTestSuite() { +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT + dp_graph.init(kGraphModel); + dp_nlist.init(kNlistModel); + dp_none.init(kNoneModel); +#endif + } + + static void TearDownTestSuite() { + dp_graph = deepmd::DeepPot(); + dp_nlist = deepmd::DeepPot(); + dp_none = deepmd::DeepPot(); + } + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + } + + // Load per-atom reference from a .expected sidecar and reduce to totals. + void load_ref(const char* path, + double& tot_e, + std::vector& per_f, + std::vector& tot_v, + int& natoms) { + deepmd_test::ExpectedRef ref; + ref.load(path); + const auto per_e = ref.get("pbc", "expected_e"); + per_f = ref.get("pbc", "expected_f"); + const auto per_v = ref.get("pbc", "expected_v"); + natoms = per_e.size(); + tot_e = 0.; + for (int ii = 0; ii < natoms; ++ii) { + tot_e += per_e[ii]; + } + tot_v.assign(9, 0.); + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 9; ++dd) { + tot_v[dd] += per_v[ii * 9 + dd]; + } + } + } + + // Run one model through the standalone build-nlist path and check it against + // its Python DeepEval reference at EPSILON. + void check_against_ref(deepmd::DeepPot& dp, const char* ref_path) { + double tot_e; + std::vector per_f, tot_v; + int natoms; + load_ref(ref_path, tot_e, per_f, tot_v, natoms); + + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - per_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - tot_v[ii]), EPSILON); + } + } +}; + +template +deepmd::DeepPot TestInferDpa1PairExclPtExpt::dp_graph; +template +deepmd::DeepPot TestInferDpa1PairExclPtExpt::dp_nlist; +template +deepmd::DeepPot TestInferDpa1PairExclPtExpt::dp_none; + +TYPED_TEST_SUITE(TestInferDpa1PairExclPtExpt, ValueTypes); + +// Graph route: applyPairExclusion at the ingestion seam + compiled exclusion. +TYPED_TEST(TestInferDpa1PairExclPtExpt, graph_route_matches_python_ref) { + this->check_against_ref(this->dp_graph, kGraphRef); +} + +// Dense route: applyPairExclusionNlist at the ingestion seam + compiled +// exclusion. +TYPED_TEST(TestInferDpa1PairExclPtExpt, nlist_route_matches_python_ref) { + this->check_against_ref(this->dp_nlist, kNlistRef); +} + +// The two ingestion routes carry the SAME weights and exclusion, so at +// non-binding sel they must agree bit-for-bit (fp64 ~1e-10). +TYPED_TEST(TestInferDpa1PairExclPtExpt, graph_equals_nlist_route) { + using VALUETYPE = TypeParam; + double e_g, e_n; + std::vector f_g, v_g, f_n, v_n; + this->dp_graph.compute(e_g, f_g, v_g, this->coord, this->atype, this->box); + this->dp_nlist.compute(e_n, f_n, v_n, this->coord, this->atype, this->box); + EXPECT_LT(fabs(e_g - e_n), EPSILON); + ASSERT_EQ(f_g.size(), f_n.size()); + for (size_t ii = 0; ii < f_g.size(); ++ii) { + EXPECT_LT(fabs(f_g[ii] - f_n[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(v_g[ii] - v_n[ii]), EPSILON); + } +} + +// The no-exclusion baseline exercises the EMPTY-table (identity) branch of the +// C++ helpers; it must run cleanly and produce an energy that DIFFERS from the +// excluded models (proving pair_exclude_types is genuinely active, not +// dropped). +TYPED_TEST(TestInferDpa1PairExclPtExpt, exclusion_is_active_vs_baseline) { + using VALUETYPE = TypeParam; + double e_none, e_g, e_n; + std::vector f, v; + this->dp_none.compute(e_none, f, v, this->coord, this->atype, this->box); + this->dp_graph.compute(e_g, f, v, this->coord, this->atype, this->box); + this->dp_nlist.compute(e_n, f, v, this->coord, this->atype, this->box); + + EXPECT_TRUE(std::isfinite(e_none)); + // Excluding all O-H pairs changes the energy well above the fp64 tolerance. + EXPECT_GT(fabs(e_g - e_none), 1e-6); + EXPECT_GT(fabs(e_n - e_none), 1e-6); +} diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 2247dcb46c..2b1c7fddb2 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -96,7 +96,10 @@ else: env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4.py & PID9=$! + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa1_pairexcl.py & + PID10=$! wait $PID9 + wait $PID10 env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_spin.py & PID7=$! diff --git a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py index ab898e15f6..b54dbd8cce 100644 --- a/source/lmp/tests/test_lammps_dpa1_graph_pt2.py +++ b/source/lmp/tests/test_lammps_dpa1_graph_pt2.py @@ -44,6 +44,21 @@ pb_file = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa1_graph.pt2" ) +# Graph-lower dpa1 with model-level pair_exclude_types=[[0,1]] and its +# same-weights no-exclusion baseline (source/tests/infer/gen_dpa1_pairexcl.py). +# Used to prove exclusion survives the extended-region multi-rank graph path. +pb_file_pairexcl = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa1_pairexcl_graph.pt2" +) +pb_file_pairexcl_none = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa1_pairexcl_none.pt2" +) ref_file = ( Path(__file__).parent.parent.parent / "tests" @@ -193,6 +208,7 @@ def _run_mpi_subprocess( data_path: Path | None = None, processors: str | None = None, runner_args: list[str] | None = None, + pb: Path | None = None, ) -> dict: """Invoke the (backend-agnostic) DPA3 MPI runner under ``mpirun -n `` against the dpa1 graph .pt2 and return @@ -200,10 +216,13 @@ def _run_mpi_subprocess( ``nprocs == 1`` forces ``--processors 1 1 1`` so the C++ side sees ``nprocs == 1`` and routes to the single-rank graph path — a - same-archive reference for the multi-rank comparison. + same-archive reference for the multi-rank comparison. ``pb`` overrides + the model archive (defaults to the no-exclusion ``deeppot_dpa1_graph.pt2``). """ if data_path is None: data_path = data_file + if pb is None: + pb = pb_file with tempfile.NamedTemporaryFile(mode="r", suffix=".out", delete=False) as f: out_path = f.name try: @@ -214,7 +233,7 @@ def _run_mpi_subprocess( sys.executable, str(mpi_runner), str(data_path.resolve()), - str(pb_file.resolve()), + str(pb.resolve()), out_path, ] if processors is not None: @@ -317,3 +336,47 @@ def test_pair_deepmd_mpi_dpa1_graph_empty_subdomain() -> None: out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=0 ) assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + + +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +@pytest.mark.skipif( + not pb_file_pairexcl.exists(), + reason="gen_dpa1_pairexcl.py .pt2 fixtures not generated", +) +def test_pair_deepmd_mpi_dpa1_pairexcl_graph_matches_single_rank() -> None: + """Model-level ``pair_exclude_types`` must survive the extended-region + multi-rank graph path (cell 5). + + Exclusion is a BUILD-time transform applied at the C++ ingestion seam + (``applyPairExclusion`` on ``edge_mask``); the SAME seam serves the + single-rank (``n_node == nloc``) and multi-rank (``n_node == nall_real``, + extended region) graph routes, so multi-rank must equal single-rank on the + excluded model. A regression that skipped the seam on the extended-region + path -- or fed it the wrong (local vs extended) atypes -- would diverge here. + + Two checks: + 1. MP (``-n 2``) ≡ SP (``-n 1``) on the excluded archive. + 2. The excluded run differs from the SAME-weights no-exclusion baseline + (``deeppot_dpa1_pairexcl_none.pt2``), so a silently-dropped exclusion + on BOTH ranks cannot pass check 1 trivially. + """ + out_mpi = _run_mpi_subprocess(nprocs=2, pb=pb_file_pairexcl) + out_ref = _run_mpi_subprocess(nprocs=1, pb=pb_file_pairexcl) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=0 + ) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + + # Exclusion must be ACTIVE on the multi-rank path: the excluded energy must + # differ from the same-weights no-exclusion baseline (O-H pairs dropped). + out_none = _run_mpi_subprocess(nprocs=2, pb=pb_file_pairexcl_none) + assert abs(out_mpi["pe"] - out_none["pe"]) > 1e-6, ( + "pair_exclude_types had no effect on the multi-rank graph path " + f"(|E_excl - E_none| = {abs(out_mpi['pe'] - out_none['pe']):.2e})" + ) diff --git a/source/lmp/tests/test_lammps_dpa3_pt2.py b/source/lmp/tests/test_lammps_dpa3_pt2.py index a55d80561c..b359ba5028 100644 --- a/source/lmp/tests/test_lammps_dpa3_pt2.py +++ b/source/lmp/tests/test_lammps_dpa3_pt2.py @@ -35,6 +35,16 @@ pb_file_mpi = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa3_mpi.pt2" ) +# Same as deeppot_dpa3_mpi.pt2 but with model-level pair_exclude_types=[[0,1]] +# (identical weights). Used to check that model-level exclusion survives the +# dense multi-rank (run_model_with_comm) path; deeppot_dpa3_mpi.pt2 is its +# no-exclusion baseline. Produced by source/tests/infer/gen_dpa3.py. +pb_file_pairexcl_mpi = ( + Path(__file__).parent.parent.parent + / "tests" + / "infer" + / "deeppot_dpa3_pairexcl_mpi.pt2" +) ref_file = ( Path(__file__).parent.parent.parent / "tests" / "infer" / "deeppot_dpa3.expected" ) @@ -540,6 +550,50 @@ def test_pair_deepmd_mpi_dpa3() -> None: ) +@pytest.mark.skipif( + shutil.which("mpirun") is None, reason="MPI is not installed on this system" +) +@pytest.mark.skipif( + importlib.util.find_spec("mpi4py") is None, reason="mpi4py is not installed" +) +@pytest.mark.skipif( + not pb_file_pairexcl_mpi.exists(), + reason="gen_dpa3.py pair-exclude .pt2 fixture not generated", +) +def test_pair_deepmd_mpi_dpa3_pairexcl_matches_single_rank() -> None: + """Model-level ``pair_exclude_types`` must survive the dense multi-rank + (with-comm) path (cell 3). + + DPA3 is message-passing, so ``use_loc_mapping=False`` routes multi-rank + through ``run_model_with_comm`` (the dense with-comm lower). Decision + #18/A4 removed exclusion from the exported lower, so it must be applied at + the C++ ingestion seam (``applyPairExclusionNlist``) on THIS path too -- + otherwise a message-passing model with ``pair_exclude_types`` silently + includes excluded pairs multi-rank (multi-rank != single-rank). + + Two checks: + 1. MP (``-n 2``) ≡ SP (``-n 1``) on the excluded archive. + 2. The excluded run differs from the SAME-weights no-exclusion baseline + (``deeppot_dpa3_mpi.pt2``), so a silently-dropped exclusion on BOTH + ranks cannot pass check 1 trivially. + """ + out_mpi = _run_mpi_subprocess(nprocs=2, pb_path=pb_file_pairexcl_mpi) + out_ref = _run_mpi_subprocess(nprocs=1, pb_path=pb_file_pairexcl_mpi) + np.testing.assert_allclose(out_mpi["forces"], out_ref["forces"], atol=1e-8, rtol=0) + np.testing.assert_allclose( + out_mpi["virials"], out_ref["virials"], atol=1e-8, rtol=0 + ) + assert out_mpi["pe"] == pytest.approx(out_ref["pe"], rel=1e-8, abs=1e-10) + + # Exclusion must be ACTIVE on the with-comm path: energy must differ from + # the same-weights no-exclusion baseline (O-H pairs dropped). + out_none = _run_mpi_subprocess(nprocs=2, pb_path=pb_file_mpi) + assert abs(out_mpi["pe"] - out_none["pe"]) > 1e-6, ( + "pair_exclude_types had no effect on the dense multi-rank path " + f"(|E_excl - E_none| = {abs(out_mpi['pe'] - out_none['pe']):.2e})" + ) + + @pytest.mark.skipif( shutil.which("mpirun") is None, reason="MPI is not installed on this system" ) diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion.py b/source/tests/common/dpmodel/test_apply_pair_exclusion.py new file mode 100644 index 0000000000..b73f4c8d38 --- /dev/null +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import pytest + +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + apply_pair_exclusion, +) + + +def _toy_graph(): + # 4 nodes, types [0, 1, 0, 1]; 5 edges incl. one already-masked pad edge. + # edge_index rows: [src(neighbor), dst(center)] + edge_index = np.array([[1, 2, 3, 0, 0], [0, 0, 1, 3, 0]], dtype=np.int64) + edge_vec = np.ones((5, 3), dtype=np.float64) + edge_mask = np.array([1, 1, 1, 1, 0], dtype=np.int32) # last = padding + n_node = np.array([4], dtype=np.int64) + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) + + +def test_none_and_empty_are_identity() -> None: + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + assert apply_pair_exclusion(g, atype, None) is g + assert apply_pair_exclusion(g, atype, PairExcludeMask(2, [])) is g + + +def test_excluded_pairs_are_masked_and_padding_stays_masked() -> None: + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)])) + # edges (dst_t, src_t): e0 (0,1) excl, e1 (0,0) keep, e2 (1,1) keep, + # e3 (1,0) excl (symmetric), e4 padding stays 0. + np.testing.assert_array_equal(out.edge_mask, [0, 1, 1, 0, 0]) + # non-mask fields untouched, input not mutated + np.testing.assert_array_equal(g.edge_mask, [1, 1, 1, 1, 0]) + assert out.edge_index is g.edge_index + assert out.edge_vec is g.edge_vec + + +def test_no_exclusion_empty_list_is_identity() -> None: + """Cover PairExcludeMask with non-None but empty exclude list.""" + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + result = apply_pair_exclusion(g, atype, PairExcludeMask(2, [])) + assert result is g + + +def test_no_excluded_edges_in_graph() -> None: + """Exclusion list non-empty but no edge matches — all real edges stay.""" + g = _toy_graph() + atype = np.array([0, 0, 0, 0], dtype=np.int64) # all same type + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)])) + # (0,1) never appears — all edges kept (except pre-existing padding) + np.testing.assert_array_equal(out.edge_mask, [1, 1, 1, 1, 0]) + + +def test_torch_namespace_smoke() -> None: + torch = pytest.importorskip("torch") + g = _toy_graph() + gt = NeighborGraph( + n_node=torch.from_numpy(g.n_node), + edge_index=torch.from_numpy(g.edge_index), + edge_vec=torch.from_numpy(g.edge_vec), + edge_mask=torch.from_numpy(g.edge_mask), + ) + atype = torch.tensor([0, 1, 0, 1], dtype=torch.int64) + out = apply_pair_exclusion(gt, atype, PairExcludeMask(2, [(0, 1)])) + np.testing.assert_array_equal(out.edge_mask.numpy(), [0, 1, 1, 0, 0]) + + +# --------------------------------------------------------------------------- +# compact=True tests +# --------------------------------------------------------------------------- + + +def test_compact_drops_masked_edges() -> None: + """compact=True must keep exactly the valid edges (after exclusion).""" + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + out_mask = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)])) + out_compact = apply_pair_exclusion( + g, atype, PairExcludeMask(2, [(0, 1)]), compact=True + ) + # Expected kept edges: indices 1 and 2 (mask-only has [0,1,1,0,0]) + assert out_compact.edge_index.shape[1] == 2 + assert out_compact.edge_vec.shape[0] == 2 + assert out_compact.edge_mask.shape[0] == 2 + # all remaining edge_mask entries must be 1 + np.testing.assert_array_equal(out_compact.edge_mask, [1, 1]) + # edge_index content matches kept edges from mask path + np.testing.assert_array_equal( + out_compact.edge_index, out_mask.edge_index[:, [1, 2]] + ) + + +def test_compact_drops_preexisting_padding_too() -> None: + """Pre-existing padding (edge 4) must be dropped even with no exclusions.""" + g = _toy_graph() + atype = np.array([0, 0, 0, 0], dtype=np.int64) # no type exclusions + # compact=True with empty exclusion list -> graph has no exclusion keep_idx change + # The brief says compact on identity returns graph unchanged, + # but with a non-empty excl list that matches nothing, out has same edge_mask as g. + # Let's use a real exclusion that changes something so compact is non-trivial: + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + # (0,1) never appears in this graph (all types are 0) → all real edges kept + # but pre-existing padding (edge 4) should be dropped + assert out.edge_index.shape[1] == 4 # only 4 real edges, padding gone + np.testing.assert_array_equal(out.edge_mask, [1, 1, 1, 1]) + + +def test_compact_torch_smoke() -> None: + torch = pytest.importorskip("torch") + g = _toy_graph() + gt = NeighborGraph( + n_node=torch.from_numpy(g.n_node), + edge_index=torch.from_numpy(g.edge_index), + edge_vec=torch.from_numpy(g.edge_vec), + edge_mask=torch.from_numpy(g.edge_mask), + ) + atype = torch.tensor([0, 1, 0, 1], dtype=torch.int64) + out = apply_pair_exclusion(gt, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + np.testing.assert_array_equal(out.edge_mask.numpy(), [1, 1]) + assert out.edge_index.shape[1] == 2 + + +def test_compact_invariance_vs_mask_only() -> None: + """Descriptor-level invariance: segment_sum over mask-only == compact. + + Masked edges contribute zero to the sum; dropping them should give identical + results. + """ + g = _toy_graph() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + excl = PairExcludeMask(2, [(0, 1)]) + + out_mask = apply_pair_exclusion(g, atype, excl, compact=False) + out_compact = apply_pair_exclusion(g, atype, excl, compact=True) + + # Build fake per-edge values (like edge_env_mat output) + vals_mask = np.arange(5, dtype=np.float64).reshape(5, 1) + 1.0 + vals_mask_valid = vals_mask * out_mask.edge_mask[:, None] + + # Map compact edge indices to the original edge values + # Kept edges are 1 and 2 (mask [0,1,1,0,0]) + vals_compact = vals_mask[out_mask.edge_mask.astype(bool)] + + # segment_sum over dst (center) node axis: 4 nodes + N = 4 + dst_mask = out_mask.edge_index[1] # (5,) + dst_compact = out_compact.edge_index[1] # (2,) + + # manual segment_sum for mask path + result_mask = np.zeros((N, 1), dtype=np.float64) + for ei, v in zip(dst_mask, vals_mask_valid, strict=True): + result_mask[ei] += v + + result_compact = np.zeros((N, 1), dtype=np.float64) + for ei, v in zip(dst_compact, vals_compact, strict=True): + result_compact[ei] += v + + np.testing.assert_allclose(result_mask, result_compact) + + +# --------------------------------------------------------------------------- +# compact=True with angle fields — remap onto the compacted edge axis +# --------------------------------------------------------------------------- + + +def _toy_graph_with_angles(): + """Same base graph as _toy_graph but with angle_index/angle_mask populated.""" + g = _toy_graph() + import dataclasses + + # Two toy angles (pairs of edges sharing a center), into edge positions [0,5) + # angle0 = (edge0, edge1); angle1 = (edge1, edge2) + angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) + angle_mask = np.array([1, 1], dtype=np.int32) + return dataclasses.replace(g, angle_index=angle_index, angle_mask=angle_mask) + + +def test_compact_drops_angles_touching_excluded_edges() -> None: + """compact=True remaps angle_index and drops angles whose edges were excluded.""" + g = _toy_graph_with_angles() + atype = np.array([0, 1, 0, 1], dtype=np.int64) + # exclusion (0,1) masks edges 0 and 3 (see the mask-only test): survivors are + # old edges [1, 2] -> new positions [0, 1]; padding edge 4 is dropped too. + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + # edges compacted to the two survivors + assert out.edge_index.shape[1] == 2 + # angle0 = (edge0, edge1): edge0 excluded -> angle0 DROPPED. + # angle1 = (edge1, edge2): both survive -> kept, remapped (edge1->0, edge2->1). + np.testing.assert_array_equal(out.angle_index, [[0], [1]]) + np.testing.assert_array_equal(out.angle_mask, [1]) + + +def test_compact_remaps_angles_when_no_angle_dropped() -> None: + """When exclusion drops no referenced edge, angles are remapped, none lost.""" + g = _toy_graph_with_angles() + atype = np.array([0, 0, 0, 0], dtype=np.int64) # (0,1) matches nothing + # all 4 real edges kept (new positions == old for [0,1,2,3]); padding dropped. + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + assert out.edge_index.shape[1] == 4 + # both angles survive; indices unchanged (survivor ranks equal old positions) + np.testing.assert_array_equal(out.angle_index, [[0, 1], [1, 2]]) + np.testing.assert_array_equal(out.angle_mask, [1, 1]) + + +def test_compact_raises_when_only_angle_index_present() -> None: + """compact=True fails fast when angle_index is set but angle_mask is None.""" + import dataclasses + + g = _toy_graph() + angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) + g2 = dataclasses.replace(g, angle_index=angle_index) # angle_mask stays None + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(ValueError, match="both be set or both be None"): + apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + + +def test_compact_raises_when_only_angle_mask_present() -> None: + """compact=True fails fast when angle_mask is set but angle_index is None.""" + import dataclasses + + g = _toy_graph() + angle_mask = np.array([1], dtype=np.int32) + g_with_mask = dataclasses.replace(g, angle_mask=angle_mask) + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(ValueError, match="both be set or both be None"): + apply_pair_exclusion( + g_with_mask, atype, PairExcludeMask(2, [(0, 1)]), compact=True + ) + + +def test_compact_raises_on_angle_dim_mismatch() -> None: + """compact=True fails fast when angle_index (2,A) and angle_mask (A,) disagree.""" + import dataclasses + + g = _toy_graph() + angle_index = np.array([[0, 1], [1, 2]], dtype=np.int64) # A == 2 + angle_mask = np.array([1], dtype=np.int32) # A == 1 (mismatch) + g2 = dataclasses.replace(g, angle_index=angle_index, angle_mask=angle_mask) + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(ValueError, match="disagree on A"): + apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + + +def test_compact_raises_on_bad_angle_index_shape() -> None: + """compact=True fails fast when angle_index is not (2, A).""" + import dataclasses + + g = _toy_graph() + angle_index = np.array([[0, 1, 2]], dtype=np.int64) # (1, 3), not (2, A) + angle_mask = np.array([1, 1, 1], dtype=np.int32) + g2 = dataclasses.replace(g, angle_index=angle_index, angle_mask=angle_mask) + atype = np.array([0, 1, 0, 1], dtype=np.int64) + with pytest.raises(ValueError, match=r"shape \(2, A\)"): + apply_pair_exclusion(g2, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + + +def test_compact_angle_torch_smoke() -> None: + """compact-mode angle remap runs under the torch namespace.""" + torch = pytest.importorskip("torch") + g = _toy_graph_with_angles() + gt = NeighborGraph( + n_node=torch.from_numpy(g.n_node), + edge_index=torch.from_numpy(g.edge_index), + edge_vec=torch.from_numpy(g.edge_vec), + edge_mask=torch.from_numpy(g.edge_mask), + angle_index=torch.from_numpy(g.angle_index), + angle_mask=torch.from_numpy(g.angle_mask), + ) + atype = torch.tensor([0, 1, 0, 1], dtype=torch.int64) + out = apply_pair_exclusion(gt, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + np.testing.assert_array_equal(out.angle_index.numpy(), [[0], [1]]) + np.testing.assert_array_equal(out.angle_mask.numpy(), [1]) + + +def test_compact_works_when_angle_fields_are_none() -> None: + """compact=True must NOT raise when angle_index and angle_mask are both None.""" + g = _toy_graph() # angle_index=None, angle_mask=None by default + atype = np.array([0, 1, 0, 1], dtype=np.int64) + # Should succeed; reuse the existing compact assertion + out = apply_pair_exclusion(g, atype, PairExcludeMask(2, [(0, 1)]), compact=True) + assert out.edge_index.shape[1] == 2 diff --git a/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py new file mode 100644 index 0000000000..0d4f2c2304 --- /dev/null +++ b/source/tests/common/dpmodel/test_apply_pair_exclusion_nlist.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for :func:`apply_pair_exclusion_nlist` and the +``pair_excl`` parameter of :func:`build_neighbor_list` / +``DefaultNeighborList.build``. +""" + +import numpy as np +import pytest + +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + build_neighbor_list, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _make_nlist_atype(): + """2-frame, 3-local, 4-neighbor test fixture. + + nlist[f, i, k] = index of the k-th neighbor of atom i in frame f. + -1 = empty slot. + + Frame 0 - atoms 0,1,2 (types 0,1,0); ghost atom 3 (type 1). + Frame 1 - same layout, different neighbor assignment. + """ + # shape (nf=2, nloc=3, nnei=4) + nlist = np.array( + [ + # frame 0 + [[1, 2, 3, -1], [0, 3, -1, -1], [0, 1, -1, -1]], + # frame 1 + [[2, 3, -1, -1], [0, -1, -1, -1], [1, 2, 3, -1]], + ], + dtype=np.int64, + ) + # extended atype: local atoms + 1 ghost; shape (nf=2, nall=4) + # types: atom0=0, atom1=1, atom2=0, ghost3=1 + atype_ext = np.array( + [[0, 1, 0, 1], [0, 1, 0, 1]], + dtype=np.int64, + ) + return nlist, atype_ext + + +# --------------------------------------------------------------------------- +# apply_pair_exclusion_nlist — unit tests +# --------------------------------------------------------------------------- + + +def test_none_is_identity() -> None: + nlist, atype_ext = _make_nlist_atype() + result = apply_pair_exclusion_nlist(nlist, atype_ext, None) + assert result is nlist + + +def test_empty_exclude_list_is_identity() -> None: + nlist, atype_ext = _make_nlist_atype() + excl = PairExcludeMask(2, []) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + assert result is nlist + + +def test_excluded_pairs_become_minus_one() -> None: + """Neighbors of excluded type should become -1.""" + nlist, atype_ext = _make_nlist_atype() + # Exclude (type0, type1) pairs (symmetric: also type1,type0). + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + + # Frame 0, atom 0 (type 0): neighbors 1(type1), 2(type0), 3(type1), -1 + # (0,1) and (0,3) are excluded; (0,2) kept; -1 stays -1 + np.testing.assert_array_equal(result[0, 0], [-1, 2, -1, -1]) + # Frame 0, atom 1 (type 1): neighbors 0(type0), 3(type1), -1, -1 + # (1,0) excluded; (1,3) kept (type1,type1 not excluded) + np.testing.assert_array_equal(result[0, 1], [-1, 3, -1, -1]) + # Frame 0, atom 2 (type 0): neighbors 0(type0), 1(type1), -1, -1 + # (2,0) kept (0,0); (2,1) excluded (0,1) + np.testing.assert_array_equal(result[0, 2], [0, -1, -1, -1]) + + +def test_already_empty_slots_preserved() -> None: + """Entries that are already -1 must remain -1 after exclusion.""" + nlist, atype_ext = _make_nlist_atype() + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + # -1 positions in the original nlist must still be -1 + original_minus_one = nlist == -1 + np.testing.assert_array_equal(result[original_minus_one], -1) + + +def test_ghost_atom_type_respected() -> None: + """Ghost atoms (index >= nloc) are indexed by atype_ext; their types count.""" + nlist, atype_ext = _make_nlist_atype() + # atype_ext[*,3] == 1; atom 0 is type 0. (0,1) is excluded. + # atom 0 in frame 0 lists neighbor 3 (ghost, type 1) -> should be excluded. + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + assert result[0, 0, 2] == -1 # was 3 (type 1), now excluded + + +def test_idempotent_double_application() -> None: + """Applying exclusion twice must give the same result as applying once.""" + nlist, atype_ext = _make_nlist_atype() + excl = PairExcludeMask(2, [(0, 1)]) + once = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + twice = apply_pair_exclusion_nlist(once, atype_ext, excl) + np.testing.assert_array_equal(once, twice) + + +def test_no_matching_pairs_leaves_nlist_unchanged() -> None: + """Exclusion list non-empty but no edge matches — nlist unchanged.""" + nlist, atype_ext = _make_nlist_atype() + # All atoms are type 0 or 1; exclude (2,3) which doesn't appear. + excl = PairExcludeMask(4, [(2, 3)]) + result = apply_pair_exclusion_nlist(nlist, atype_ext, excl) + np.testing.assert_array_equal(result, nlist) + + +# --------------------------------------------------------------------------- +# apply_pair_exclusion_nlist — torch namespace smoke test +# --------------------------------------------------------------------------- + + +def test_torch_namespace_smoke() -> None: + torch = pytest.importorskip("torch") + nlist, atype_ext = _make_nlist_atype() + nlist_t = torch.from_numpy(nlist) + atype_t = torch.from_numpy(atype_ext) + excl = PairExcludeMask(2, [(0, 1)]) + result = apply_pair_exclusion_nlist(nlist_t, atype_t, excl) + np.testing.assert_array_equal( + result.numpy(), apply_pair_exclusion_nlist(nlist, atype_ext, excl) + ) + + +# --------------------------------------------------------------------------- +# build_neighbor_list pair_excl parameter +# --------------------------------------------------------------------------- + + +def _simple_extended_system(): + """2-atom local system with 1 ghost to give a nontrivial nlist. + + Atoms: local 0 (type 0) at (0,0,0), local 1 (type 1) at (1,0,0). + Ghost 2 (type 1) at (-1,0,0) (periodic image of atom 1 across pbc). + + rcut=1.5: atom 0 sees neighbors 1, 2; atom 1 sees neighbor 0. + """ + # nf=1, nall=3 + coord_ext = np.array( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]], + dtype=np.float64, + ) + atype_ext = np.array([[0, 1, 1]], dtype=np.int64) + return coord_ext, atype_ext + + +def test_build_neighbor_list_pair_excl_equals_post_application() -> None: + """build_neighbor_list(pair_excl=excl) must equal apply_pair_exclusion_nlist + applied to the result without pair_excl (oracle equivalence). + """ + coord_ext, atype_ext = _simple_extended_system() + nloc = 2 + rcut = 1.5 + sel = [4] + excl = PairExcludeMask(2, [(0, 1)]) + + nlist_plain = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False + ) + nlist_excl_builtin = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False, pair_excl=excl + ) + nlist_excl_manual = apply_pair_exclusion_nlist(nlist_plain, atype_ext, excl) + + np.testing.assert_array_equal(nlist_excl_builtin, nlist_excl_manual) + + +def test_build_neighbor_list_none_excl_unchanged() -> None: + """pair_excl=None must not change the nlist.""" + coord_ext, atype_ext = _simple_extended_system() + nloc = 2 + rcut = 1.5 + sel = [4] + + nlist_plain = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False + ) + nlist_with_none = build_neighbor_list( + coord_ext, atype_ext, nloc, rcut, sel, distinguish_types=False, pair_excl=None + ) + np.testing.assert_array_equal(nlist_plain, nlist_with_none) + + +# --------------------------------------------------------------------------- +# DefaultNeighborList.build pair_excl parameter +# --------------------------------------------------------------------------- + + +def _local_system(): + """Return a simple local (non-extended) 2-atom system for DefaultNeighborList.""" + # shape (nf=1, nloc=2, 3) + coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]], dtype=np.float64) + atype = np.array([[0, 1]], dtype=np.int64) + return coord, atype + + +def test_default_neighbor_list_pair_excl_equals_seam() -> None: + """DefaultNeighborList(pair_excl=excl) nlist equals build-then-apply.""" + from deepmd.dpmodel.utils.default_neighbor_list import ( + DefaultNeighborList, + ) + + coord, atype = _local_system() + rcut = 1.5 + sel = [4] + excl = PairExcludeMask(2, [(0, 1)]) + + builder = DefaultNeighborList() + # Build without exclusion + ext_coord, ext_atype, nlist_plain, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel + ) + nlist_manual = apply_pair_exclusion_nlist(nlist_plain, ext_atype, excl) + + # Build with exclusion at builder level + _, ext_atype2, nlist_builtin, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel, pair_excl=excl + ) + np.testing.assert_array_equal(ext_atype, ext_atype2) + np.testing.assert_array_equal(nlist_builtin, nlist_manual) + + +# --------------------------------------------------------------------------- +# VesinNeighborList.build pair_excl (torch only) +# --------------------------------------------------------------------------- + + +def test_vesin_neighbor_list_pair_excl_equals_seam() -> None: + """VesinNeighborList(pair_excl=excl) nlist equals build-then-apply.""" + torch = pytest.importorskip("torch") + from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + is_vesin_torch_available, + ) + + if not is_vesin_torch_available(): + pytest.skip("vesin.torch not installed") + + coord = torch.tensor([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]], dtype=torch.float64) + atype = torch.tensor([[0, 1]], dtype=torch.int64) + rcut = 1.5 + sel = [4] + excl = PairExcludeMask(2, [(0, 1)]) + + builder = VesinNeighborList() + ext_coord, ext_atype, nlist_plain, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel + ) + nlist_manual = apply_pair_exclusion_nlist(nlist_plain, ext_atype, excl) + + _, ext_atype2, nlist_builtin, _ = builder.build( + coord, atype, box=None, rcut=rcut, sel=sel, pair_excl=excl + ) + np.testing.assert_array_equal(nlist_builtin.numpy(), nlist_manual.numpy()) + + +# --------------------------------------------------------------------------- +# NvNeighborList.build: edges + pair_excl raises +# --------------------------------------------------------------------------- + + +def test_nv_nlist_edges_pair_excl_raises(): + """NvNeighborList.build raises NotImplementedError for edges+pair_excl. + + The guard fires before any CUDA search, so this test runs on CPU. + NvNeighborList requires CUDA to produce results, but the early-exit + raise is device-independent. + """ + import torch + + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + from deepmd.pt.utils.nv_nlist import ( + NvNeighborList, + ) + + coord = torch.zeros((1, 4, 3), dtype=torch.float64) + atype = torch.zeros((1, 4), dtype=torch.int64) + pe = PairExcludeMask(2, [(0, 1)]) + nl = NvNeighborList() + with pytest.raises(NotImplementedError, match="return_mode='edges'"): + nl.build(coord, atype, None, 2.0, [4], return_mode="edges", pair_excl=pe) diff --git a/source/tests/common/dpmodel/test_dp_atomic_model.py b/source/tests/common/dpmodel/test_dp_atomic_model.py index ae5a9c610d..721a4f5895 100644 --- a/source/tests/common/dpmodel/test_dp_atomic_model.py +++ b/source/tests/common/dpmodel/test_dp_atomic_model.py @@ -13,6 +13,9 @@ from deepmd.dpmodel.fitting import ( InvarFitting, ) +from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, +) from .case_single_frame_with_nlist import ( TestCaseSingleFrameWithNlist, @@ -76,8 +79,15 @@ def test_self_consistency( md0.reinit_pair_exclude(pair_excl) md1 = DPAtomicModel.deserialize(md0.serialize()) - ret0 = md0.forward_common_atomic(self.coord_ext, self.atype_ext, self.nlist) - ret1 = md1.forward_common_atomic(self.coord_ext, self.atype_ext, self.nlist) + # forward_common_atomic consumes a pre-excluded nlist (decision + # #18/A4). md0 and md1 share the same pair_excl (deserialized), so + # fold it into the nlist once; the serialize/deserialize consistency + # check is unaffected by which (identical) nlist both consume. + nlist = self.nlist + if md0.pair_excl is not None: + nlist = apply_pair_exclusion_nlist(nlist, self.atype_ext, md0.pair_excl) + ret0 = md0.forward_common_atomic(self.coord_ext, self.atype_ext, nlist) + ret1 = md1.forward_common_atomic(self.coord_ext, self.atype_ext, nlist) np.testing.assert_allclose(ret0["energy"], ret1["energy"]) @@ -111,10 +121,16 @@ def test_excl_consistency(self) -> None: md1.descriptor.reinit_exclude(pair_excl) md1.fitting_net.reinit_exclude(atom_excl) - # check energy consistency - args = [self.coord_ext, self.atype_ext, self.nlist] - ret0 = md0.forward_common_atomic(*args) - ret1 = md1.forward_common_atomic(*args) + # check energy consistency. Model-level pair exclusion is a + # nlist-BUILD transform (decision #18/A4): forward_common_atomic + # consumes a pre-excluded nlist, so fold md0's exclusion into its + # nlist here (the descriptor-level md1 keeps the raw nlist; its + # emask applies inside the descriptor). + nlist0 = apply_pair_exclusion_nlist( + self.nlist, self.atype_ext, md0.pair_excl + ) + ret0 = md0.forward_common_atomic(self.coord_ext, self.atype_ext, nlist0) + ret1 = md1.forward_common_atomic(self.coord_ext, self.atype_ext, self.nlist) np.testing.assert_allclose( ret0["energy"], ret1["energy"], diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py index fb01caec55..11db9a7785 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_block.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_block.py @@ -93,27 +93,75 @@ def test_block_graph_equals_dense_any_sel(self, sel, type_one_side) -> None: # attn_layer > 0 is supported since NeighborGraph PR-D; parity is covered # by test_dpa1_graph_attention_parity.py (the fail-fast test was removed). - def test_exclude_types_raises(self) -> None: - """The graph block kernel fail-fasts for exclude_types (not yet applied).""" - # the graph path does not yet apply type exclusion; it must fail-fast - # rather than silently diverge from the dense path (which masks edges). + @pytest.mark.parametrize("type_one_side", [False, True]) # tebd concat branch + @pytest.mark.parametrize("attn_layer", [0, 2]) # no-attn and multi-layer attention + def test_exclude_types_graph_parity(self, attn_layer, type_one_side) -> None: + """The graph block kernel supports exclude_types via apply_pair_exclusion. + + Excluded edges are masked (edge_mask zeroed) before segment_sum, so + the graph block output matches the dense block output at rtol=atol=1e-12 + for a non-binding sel. Parametrized over attn_layer (0 and 2) and + type_one_side (False and True). smooth_type_embedding=False is used + for attn_layer>0 to avoid the known by-design divergence where the dense + smooth path keeps sel-padding terms in the softmax denominator. + """ + from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, + ) + + rng = np.random.default_rng(99) + nloc = 4 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype_arr = np.array([[0, 1, 0, 1]], dtype=np.int64) dd = DescrptDPA1( rcut=4.0, rcut_smth=0.5, - sel=[20], + sel=[30], # non-binding ntypes=2, - attn_layer=0, + attn_layer=attn_layer, + axis_neuron=2, + neuron=[6, 12], exclude_types=[(0, 1)], + type_one_side=type_one_side, + smooth_type_embedding=False, # avoid dense smooth divergence at attn>0 ) - ng = from_dense_quartet( - self.coord, - -np.ones((1, self.nloc, 1), dtype=np.int64), # any graph; guard fires first - np.arange(self.nloc, dtype=np.int64)[None], + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, + atype_arr, + dd.get_rcut(), + dd.get_sel(), + mixed_types=dd.mixed_types(), + box=None, + ) + ng = from_dense_quartet(ext_coord, nlist, mapping, compact=False) + atype_local = atype_arr.reshape(-1) + tebd = dd.type_embedding.call() + nf, nall = ext_atype.shape + atype_embd_ext = np.reshape( + np.take(tebd, ext_atype.reshape(-1), axis=0), + (nf, nall, dd.tebd_dim), + ) + # dense block call (apply_pair_exclusion inside) + grrg_dense, *_ = dd.se_atten.call( + nlist, + ext_coord, + ext_atype, + atype_embd_ext=atype_embd_ext, + mapping=None, + type_embedding=tebd, + ) + # graph block call (apply_pair_exclusion via edge_mask zeroing) + grrg_blk, _rot_mat = dd.se_atten.call_graph( + ng, atype_local, type_embedding=tebd + ) + assert grrg_blk.shape[0] == atype_local.shape[0] # flat N axis + assert not np.any(np.isnan(grrg_blk)) + np.testing.assert_allclose( + grrg_blk.reshape(grrg_dense.shape), + grrg_dense, + rtol=1e-12, + atol=1e-12, ) - with pytest.raises(NotImplementedError): - dd.se_atten.call_graph( - ng, self.atype.reshape(-1), type_embedding=dd.type_embedding.call() - ) class TestDpa1BlockCallGraphStrip: diff --git a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py index 6632d82894..b889644abc 100644 --- a/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py +++ b/source/tests/common/dpmodel/test_dpa1_call_graph_descriptor.py @@ -96,21 +96,38 @@ def test_descriptor_graph_equals_dense_full_tuple(self, sel) -> None: # sw np.testing.assert_allclose(out[4], ref[4], rtol=1e-12, atol=1e-12) + # NOTE: after the strip PR (#5747) and this PR both landed, neither strip + # nor exclude_types falls back to dense — both are graph-eligible. The only + # non-disabled ineligible config is compression, whose gate is covered by + # TestDpa1StripRouting.test_uses_graph_lower_strip_gate, and the no-mapping + # ghost fallback by test_eligible_no_mapping_with_ghosts_falls_back below. + @pytest.mark.parametrize( - "kwargs", - [ - {"exclude_types": [(0, 1)]}, # type exclusion: graph unsupported -> dense - ], + "exclude_types", + [[], [(0, 1)]], # empty exclusions AND non-trivial exclusion ) - def test_ineligible_config_falls_back_to_dense(self, kwargs) -> None: - """attn_layer=0 configs the graph can't handle (exclude_types) must report - uses_graph_lower()=False and run the dense body without raising (strip is - now graph-eligible; its routing + parity is covered by TestDpa1StripRouting). + def test_exclude_types_graph_eligible_and_parity(self, exclude_types) -> None: + """exclude_types (Task 3): descriptor is graph-eligible (uses_graph_lower() + True) regardless of the exclusion list. Graph output must match the dense + reference at rtol=atol=1e-12 for a non-binding sel. """ + from deepmd.dpmodel.utils.neighbor_graph import ( + from_dense_quartet, + ) + dd = DescrptDPA1( - rcut=4.0, rcut_smth=0.5, sel=[30], ntypes=2, attn_layer=0, **kwargs + rcut=4.0, + rcut_smth=0.5, + sel=[30], # non-binding sel + ntypes=2, + attn_layer=0, + axis_neuron=2, + neuron=[6, 12], + exclude_types=exclude_types, ) - assert dd.uses_graph_lower() is False + # gate: with any exclude list the descriptor must now be graph-eligible + assert dd.uses_graph_lower() is True + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( self.coord, self.atype, @@ -119,8 +136,28 @@ def test_ineligible_config_falls_back_to_dense(self, kwargs) -> None: mixed_types=dd.mixed_types(), box=None, ) - out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) # must not raise + # dense reference (calls block directly) + ref = self._dense_reference(dd, ext_coord, ext_atype, nlist) + # graph-routed public call + out = dd.call(ext_coord, ext_atype, nlist, mapping=mapping) assert len(out) == 5 + np.testing.assert_allclose(out[0], ref[0], rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(out[1], ref[1], rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(out[4], ref[4], rtol=1e-12, atol=1e-12) + + if exclude_types: + # verify excluded pairs contribute sw == 0 in the dense reference + # (atype=[0,1,0,1] -> pairs (0,1) and (1,0) should be masked) + # sw shape: (nf, nloc, nnei, 1); just check the graph output is also 0 + # for excluded-pair edges by checking call_graph sw channel + graph = from_dense_quartet(ext_coord, nlist, mapping, compact=False) + atype_local = self.atype.reshape(-1) + grrg_g, rot_mat_g = dd.call_graph( + graph, atype_local, type_embedding=dd.type_embedding.call() + ) + # no nan/inf in output with exclusions applied + assert not np.any(np.isnan(grrg_g)) + assert not np.any(np.isinf(grrg_g)) def test_eligible_no_mapping_with_ghosts_falls_back(self) -> None: """An eligible (concat) attn_layer=0 descriptor called with mapping=None @@ -231,15 +268,17 @@ def _make(self, type_one_side: bool, smooth: bool, attn_layer: int) -> DescrptDP ) def test_uses_graph_lower_strip_gate(self) -> None: - """The gate admits non-compressed strip; excludes compressed and exclude_types.""" + """The gate admits non-compressed strip (and exclude_types, now + graph-native via the pair-exclude seam); only compression forces dense. + """ dd = self._make(type_one_side=False, smooth=True, attn_layer=2) assert dd.uses_graph_lower() is True # strip is now graph-eligible # negative contract: compression keeps the descriptor on the dense path dd.compress = True assert dd.uses_graph_lower() is False dd.compress = False - # negative contract: exclude_types still forces dense (separate feature, - # owned by the pair-exclude PR; strip does NOT change its eligibility) + # exclude_types is now graph-native (owned by the pair-exclude seam), so + # strip + exclude_types stays graph-eligible. dd_excl = DescrptDPA1( rcut=4.0, rcut_smth=0.5, @@ -250,7 +289,7 @@ def test_uses_graph_lower_strip_gate(self) -> None: exclude_types=[(0, 1)], resnet_dt=False, ) - assert dd_excl.uses_graph_lower() is False + assert dd_excl.uses_graph_lower() is True @pytest.mark.parametrize("type_one_side", [False, True]) # strip table branch @pytest.mark.parametrize( diff --git a/source/tests/common/dpmodel/test_env_mat_stat_exclude.py b/source/tests/common/dpmodel/test_env_mat_stat_exclude.py new file mode 100644 index 0000000000..2f76602c4f --- /dev/null +++ b/source/tests/common/dpmodel/test_env_mat_stat_exclude.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Piece A: model-level pair exclusion enters the descriptor INPUT stat at +nlist BUILD, matching the model forward. + +Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision +#18/A4). In the forward the descriptor is fed a pre-excluded nlist, so an +excluded pair is treated exactly like an empty slot (env_mat 0, still counted). +The input stat now folds the same exclusion into the neighbor list it builds +(``EnvMatStatSe.iter``) instead of deselecting excluded pairs during +accumulation. + +The load-bearing invariant: excluding a type-pair via model-level +``pair_exclude_types`` must give the SAME stat (davg/dstd) as excluding it via +descriptor-level ``exclude_types`` -- both zero-and-count the excluded pairs. +Before this change they diverged (descriptor-level zero-counted, model-level +deselected). A no-exclusion control proves the exclusion is genuinely active. +""" + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) + + +def _sample() -> dict: + rng = np.random.default_rng(0) + nf, nloc = 2, 6 + coord = rng.normal(size=(nf, nloc, 3)) * 2.0 + # both types present so the (0, 1) exclusion actually removes pairs + atype = np.array([[0, 1, 0, 1, 0, 1], [1, 0, 1, 0, 1, 0]], dtype=np.int64) + box = np.tile((np.eye(3) * 12.0).reshape(1, 9), (nf, 1)) + return {"coord": coord, "atype": atype, "box": box} + + +def _stats(exclude_types, pair_exclude_types): + """Compute (davg, dstd) for a se_e2_a descriptor under the given exclusions.""" + ds = DescrptSeA(6.0, 0.5, [10, 10], exclude_types=exclude_types) + sample = _sample() + if pair_exclude_types is not None: + sample["pair_exclude_types"] = pair_exclude_types + ds.compute_input_stats([sample]) + return np.asarray(ds.davg), np.asarray(ds.dstd) + + +def test_model_level_equals_descriptor_level_exclusion() -> None: + """Model-level pair_exclude of (0,1) == descriptor-level exclude_types of (0,1).""" + davg_model, dstd_model = _stats(exclude_types=[], pair_exclude_types=[[0, 1]]) + davg_desc, dstd_desc = _stats(exclude_types=[[0, 1]], pair_exclude_types=None) + # Both zero-and-count the same excluded pairs, so the env-matrix + # distribution -- hence the stat -- is bit-identical. + np.testing.assert_allclose(davg_model, davg_desc, rtol=1e-14, atol=1e-14) + np.testing.assert_allclose(dstd_model, dstd_desc, rtol=1e-14, atol=1e-14) + + +def test_exclusion_is_active() -> None: + """Excluding (0,1) shifts the stat vs no exclusion (proves it is applied).""" + davg_none, dstd_none = _stats(exclude_types=[], pair_exclude_types=None) + davg_excl, dstd_excl = _stats(exclude_types=[], pair_exclude_types=[[0, 1]]) + assert not np.allclose(dstd_none, dstd_excl) diff --git a/source/tests/common/dpmodel/test_env_mat_stat_graph.py b/source/tests/common/dpmodel/test_env_mat_stat_graph.py new file mode 100644 index 0000000000..0f0243aa11 --- /dev/null +++ b/source/tests/common/dpmodel/test_env_mat_stat_graph.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Piece B: the input stat's env matrix computed via the NeighborGraph +(``from_dense_quartet`` -> ``edge_env_mat``) is BIT-IDENTICAL to the dense +``EnvMat`` path. + +This is the same machinery the dpa1 graph forward uses. ``from_dense_quartet`` +reuses the same neighbor set and padding, ``edge_env_mat`` mirrors +``EnvMat.call``, and the row-major ``(frame, center, slot)`` edge order maps 1:1 +back to the dense ``(nf, nloc, nsel, 4)`` env-matrix tensor -- so the stored +``davg``/``dstd`` are unchanged. The parity must hold with no exclusion, with +model-level ``pair_exclude_types`` (folded into the nlist at BUILD, Piece A), +and with descriptor-level ``exclude_types``. +""" + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptDPA1, + DescrptSeA, +) +from deepmd.dpmodel.utils.env_mat_stat import ( + EnvMatStatSe, +) + + +def _sample() -> dict: + rng = np.random.default_rng(0) + nf, nloc = 2, 6 + coord = rng.normal(size=(nf, nloc, 3)) * 2.0 + atype = np.array([[0, 1, 0, 1, 0, 1], [1, 0, 1, 0, 1, 0]], dtype=np.int64) + box = np.tile((np.eye(3) * 12.0).reshape(1, 9), (nf, 1)) + return {"coord": coord, "atype": atype, "box": box} + + +def _stats(exclude_types, pair_exclude_types, use_graph): + ds = DescrptSeA(6.0, 0.5, [10, 10], exclude_types=exclude_types) + sample = _sample() + if pair_exclude_types is not None: + sample["pair_exclude_types"] = pair_exclude_types + st = EnvMatStatSe(ds, use_graph=use_graph) + st.load_or_compute_stats([sample], None) + mean, stddev = st() + return np.asarray(mean), np.asarray(stddev) + + +def _assert_graph_equals_dense(exclude_types, pair_exclude_types) -> None: + m_dense, s_dense = _stats(exclude_types, pair_exclude_types, use_graph=False) + m_graph, s_graph = _stats(exclude_types, pair_exclude_types, use_graph=True) + np.testing.assert_allclose(m_graph, m_dense, rtol=1e-15, atol=1e-15) + np.testing.assert_allclose(s_graph, s_dense, rtol=1e-15, atol=1e-15) + + +def test_graph_equals_dense_no_exclusion() -> None: + _assert_graph_equals_dense(exclude_types=[], pair_exclude_types=None) + + +def test_graph_equals_dense_model_level_exclusion() -> None: + _assert_graph_equals_dense(exclude_types=[], pair_exclude_types=[[0, 1]]) + + +def test_graph_equals_dense_descriptor_level_exclusion() -> None: + _assert_graph_equals_dense(exclude_types=[[0, 1]], pair_exclude_types=None) + + +def _dpa1_block_stats(pair_exclude_types, use_graph): + """Stats for the actual dpa1 (mixed_types) block, the wired descriptor.""" + ds = DescrptDPA1(6.0, 0.5, 20, ntypes=2, attn_layer=0) + block = ds.se_atten + sample = _sample() + if pair_exclude_types is not None: + sample["pair_exclude_types"] = pair_exclude_types + st = EnvMatStatSe(block, use_graph=use_graph) + st.load_or_compute_stats([sample], None) + mean, stddev = st() + return np.asarray(mean), np.asarray(stddev) + + +def test_dpa1_block_graph_equals_dense() -> None: + """The wired dpa1 block (use_graph=True) is bit-identical to the dense path.""" + for pair_exclude_types in (None, [[0, 1]]): + m_dense, s_dense = _dpa1_block_stats(pair_exclude_types, use_graph=False) + m_graph, s_graph = _dpa1_block_stats(pair_exclude_types, use_graph=True) + np.testing.assert_allclose(m_graph, m_dense, rtol=1e-15, atol=1e-15) + np.testing.assert_allclose(s_graph, s_dense, rtol=1e-15, atol=1e-15) diff --git a/source/tests/common/dpmodel/test_graph_atomic_parity.py b/source/tests/common/dpmodel/test_graph_atomic_parity.py index 7de084a25f..34b5c27b04 100644 --- a/source/tests/common/dpmodel/test_graph_atomic_parity.py +++ b/source/tests/common/dpmodel/test_graph_atomic_parity.py @@ -14,7 +14,12 @@ from deepmd.dpmodel.model.ener_model import ( EnergyModel, ) +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, + build_neighbor_graph, from_dense_quartet, ) from deepmd.dpmodel.utils.nlist import ( @@ -121,10 +126,10 @@ def test_graph_matches_dense_over_flags(virtual, type_one_side, nf): assert int(np.asarray(g["mask"])[0, -1]) == 0 # virtual atom masked -def test_pair_exclude_types_falls_back_to_dense(): - """Pair exclude_types is unsupported on the graph -> uses_graph_lower False.""" +def test_descriptor_exclude_types_is_graph_eligible(): + """Descriptor-level exclude_types (Task 3): uses_graph_lower() is True.""" m = _ener_model([30], exclude_types=[(0, 1)]) - assert m.atomic_model.descriptor.uses_graph_lower() is False + assert m.atomic_model.descriptor.uses_graph_lower() is True def test_model_pair_exclude_types_graph_matches_dense(): @@ -162,6 +167,119 @@ def test_model_pair_exclude_types_graph_matches_dense(): ), "pair exclusion must change the graph energy (same weights)" +def test_model_pair_exclude_applied_at_build_not_in_lower(): + """Seam contract + fail-safe (decision #18): model-level pair_exclude is a + graph-BUILD transform; the graph lower does NOT re-apply it. Because the + consume-time backstop was removed, the lower would otherwise be fail-OPEN + (a non-excluded input silently INCLUDES excluded pairs). To keep that from + being silent, the lower guards its input (eager): feeding a NON-excluded + graph to a model that HAS ``pair_exclude_types`` RAISES. Applying exclusion + at BUILD is the correct path and changes the lower's output. + """ + rng = np.random.default_rng(4) + nloc = 6 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0, 1]], dtype=np.int64) + box = np.eye(3).reshape(1, 9) * 20.0 + ds = DescrptDPA1(rcut=4.0, rcut_smth=0.5, sel=[200], ntypes=2, attn_layer=0) + ft = InvarFitting("energy", 2, ds.get_dim_out(), 1, mixed_types=True) + model = EnergyModel(ds, ft, type_map=["a", "b"], pair_exclude_types=[(0, 1)]) + assert model.atomic_model.pair_excl is not None + + # RAW graph: built WITHOUT pair_excl (no exclusion baked into edge_mask). + ng_raw = build_neighbor_graph(coord, atype, box, model.get_rcut()) + kw = { + "atype": atype.reshape(-1), + "n_node": ng_raw.n_node, + "edge_index": ng_raw.edge_index, + "edge_vec": ng_raw.edge_vec, + "edge_mask": ng_raw.edge_mask, + } + # Fail-safe: the lower refuses a non-excluded graph (contract boundary) — + # the excluded (0, 1) pairs are within rcut here, so at least one leaks. + with pytest.raises(AssertionError, match="NOT pair-excluded"): + model.call_lower_graph(**kw) + + # Positive control: applying exclusion at BUILD (excluded edge_mask) passes + # the guard AND changes the output vs the same lower with no exclusion. + ng_excl = build_neighbor_graph( + coord, atype, box, model.get_rcut(), pair_excl=PairExcludeMask(2, [(0, 1)]) + ) + out_built_excl = model.call_lower_graph( + atype=atype.reshape(-1), + n_node=ng_excl.n_node, + edge_index=ng_excl.edge_index, + edge_vec=ng_excl.edge_vec, + edge_mask=ng_excl.edge_mask, + ) + # No-exclusion reference: clearing pair_excl makes the raw graph valid again + # (guard skipped when pair_excl is None), so the lower runs on it. + model.atomic_model.reinit_pair_exclude([]) + assert model.atomic_model.pair_excl is None + out_no_excl_model = model.call_lower_graph(**kw) + assert not np.allclose( + np.asarray(out_built_excl["energy_redu"]), + np.asarray(out_no_excl_model["energy_redu"]), + rtol=1e-9, + atol=1e-9, + ), "build-time pair exclusion must change the graph energy" + + +def test_model_pair_exclude_applied_at_build_not_in_dense_lower(): + """Dense-route seam contract + fail-safe (decision #18/A4, mirror of the + graph test): model-level pair_exclude is a nlist-BUILD transform; the dense + lower (``call_lower``) does NOT re-apply it. With the consume-time backstop + removed, a non-excluded nlist would silently INCLUDE excluded pairs, so the + lower guards its input (eager): feeding a RAW nlist to a model that HAS + ``pair_exclude_types`` RAISES. Folding exclusion in at BUILD is correct and + changes the output. + """ + from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, + extend_input_and_build_neighbor_list, + ) + + rng = np.random.default_rng(4) + nloc = 6 + coord = rng.normal(size=(1, nloc, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0, 1]], dtype=np.int64) + box = np.eye(3).reshape(1, 9) * 20.0 + ds = DescrptDPA1(rcut=4.0, rcut_smth=0.5, sel=[200], ntypes=2, attn_layer=0) + ft = InvarFitting("energy", 2, ds.get_dim_out(), 1, mixed_types=True) + model = EnergyModel(ds, ft, type_map=["a", "b"], pair_exclude_types=[(0, 1)]) + assert model.atomic_model.pair_excl is not None + + # RAW quartet: built WITHOUT pair_excl (no exclusion folded into the nlist). + coord_ext, atype_ext, mapping, nlist = extend_input_and_build_neighbor_list( + coord.reshape(1, -1), + atype, + model.get_rcut(), + model.get_sel(), + mixed_types=True, + box=box, + ) + # Fail-safe: the lower refuses a non-excluded nlist (contract boundary). + with pytest.raises(AssertionError, match="NOT pair-excluded"): + model.call_lower(coord_ext, atype_ext, nlist, mapping) + + # Positive control: folding the exclusion in at BUILD passes the guard AND + # changes the output vs the same lower with no exclusion. + nlist_excl = apply_pair_exclusion_nlist( + nlist, atype_ext, PairExcludeMask(2, [(0, 1)]) + ) + out_built_excl = model.call_lower(coord_ext, atype_ext, nlist_excl, mapping) + # No-exclusion reference: clearing pair_excl makes the raw nlist valid again. + model.atomic_model.reinit_pair_exclude([]) + assert model.atomic_model.pair_excl is None + out_no_excl_model = model.call_lower(coord_ext, atype_ext, nlist, mapping) + assert not np.allclose( + np.asarray(out_built_excl["energy"]), + np.asarray(out_no_excl_model["energy"]), + rtol=1e-9, + atol=1e-9, + ), "build-time pair exclusion must change the dense energy" + + def test_graph_matches_dense_with_fparam(): """Frame parameter is gathered to nodes by frame_id in forward_atomic_graph and fed to the fitting's call_graph; the graph path must match dense at 1e-12 @@ -306,3 +424,32 @@ def test_graph_matches_dense_with_out_bias(): ) # non-vacuous: the bias actually shifted the graph energy assert not np.allclose(np.asarray(g["energy"]), np.asarray(g_zero["energy"])) + + +# ── apply_pair_exclusion idempotence (Task 2) ───────────────────────────────── + + +@pytest.mark.parametrize( + "pair_exclude_types", [[], [(0, 1)]] +) # empty branch AND non-empty branch +def test_apply_pair_exclusion_idempotent(pair_exclude_types): + """Applying apply_pair_exclusion twice gives the same edge_mask as once. + + Covers both the empty pair_excl branch (identity) and non-empty branch. + """ + rng = np.random.default_rng(42) + coord = rng.normal(size=(1, 5, 3)) * 1.5 + atype = np.array([[0, 1, 0, 1, 0]], dtype=np.int64) + ext_coord, ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + coord, atype, 4.0, [200], mixed_types=True, box=None + ) + ng = from_dense_quartet(ext_coord, nlist, mapping) + pair_excl = PairExcludeMask(2, pair_exclude_types) if pair_exclude_types else None + atype_flat = atype.reshape(-1) + once = apply_pair_exclusion(ng, atype_flat, pair_excl) + twice = apply_pair_exclusion(once, atype_flat, pair_excl) + # Masks must be exactly equal (AND-idempotent for 0/1 values) + np.testing.assert_array_equal( + np.asarray(once.edge_mask), + np.asarray(twice.edge_mask), + ) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 9ba25c0ccb..8f91f5af8c 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -15,8 +15,12 @@ import numpy as np +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, + apply_pair_exclusion, build_neighbor_graph, from_dense_quartet, ) @@ -313,5 +317,161 @@ def test_adapter_maps_ghost_to_local_owner(self) -> None: np.testing.assert_allclose(ev[0], np.array([3.0, 0.0, 0.0])) +def valid_edge_set(ng): + """Return the set of (src, dst, rounded edge_vec) for all real edges.""" + ei = ng.edge_index[:, ng.edge_mask] + ev = ng.edge_vec[ng.edge_mask] + return { + (int(ei[0, k]), int(ei[1, k]), tuple(np.round(ev[k], 6))) + for k in range(ei.shape[1]) + } + + +class TestBuildNeighborGraphPairExclOracle(unittest.TestCase): + """Oracle harness: builder(pair_excl=X) == builder() + apply_pair_exclusion(X). + + Covers both ``pair_excl=None`` (identity; no exclusion applied) and a + non-empty exclusion set, for the ``dense`` backend. The oracle asserts + SET-EQUALITY of the valid-edge set, matching the Task 3b contract. + """ + + def setUp(self) -> None: + self.rcut = 4.0 + # 4 atoms, 2 types (0 and 1); atom 2 offset avoids degenerate rcut alignment. + self.coord = np.array( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.3, 0.0], [3.5, 0.0, 0.0]], + dtype=np.float64, + ).reshape(1, 4, 3) + # type sequence: 0,1,0,1 -- pairs (0,1) and (1,0) are heterogeneous + self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64) + self.ntypes = 2 + + def _pair_excl(self, exclude_pairs): + """Build a PairExcludeMask from a list of (ti, tj) tuples.""" + return PairExcludeMask(self.ntypes, exclude_pairs) + + def test_pair_excl_none_identity_dense(self) -> None: + """pair_excl=None: builder output unchanged (identity).""" + ng_ref = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + ng_excl = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, pair_excl=None + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_excl)) + + def test_pair_excl_empty_list_identity_dense(self) -> None: + """pair_excl with empty exclude set: builder output unchanged.""" + pe = self._pair_excl([]) + ng_ref = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + ng_excl = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_excl)) + + def test_oracle_set_equality_dense_nonperiodic(self) -> None: + """Builder with pair_excl==(0,1) == builder() + apply_pair_exclusion.""" + pe = self._pair_excl([(0, 1), (1, 0)]) + # reference: build without exclusion then apply separately + ng_base = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + atype_flat = self.atype.reshape(-1) + ng_post = apply_pair_exclusion(ng_base, atype_flat, pe) + # under test: builder applies exclusion internally + ng_fused = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_post), valid_edge_set(ng_fused)) + # sanity: exclusion actually REMOVED some edges + self.assertLess(int(ng_fused.edge_mask.sum()), int(ng_base.edge_mask.sum())) + + def test_oracle_set_equality_dense_periodic(self) -> None: + """Periodic PBC: builder with pair_excl==(0,0) == builder() + apply.""" + pe = self._pair_excl([(0, 0)]) + box = np.eye(3, dtype=np.float64)[None] * 6.0 + ng_base = build_neighbor_graph(self.coord, self.atype, box, self.rcut) + atype_flat = self.atype.reshape(-1) + ng_post = apply_pair_exclusion(ng_base, atype_flat, pe) + ng_fused = build_neighbor_graph( + self.coord, self.atype, box, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_post), valid_edge_set(ng_fused)) + # type-0 centers: atoms 0,2; type-0 neighbors excluded; fewer edges expected + self.assertLess(int(ng_fused.edge_mask.sum()), int(ng_base.edge_mask.sum())) + + def test_oracle_set_equality_dense_multiframe(self) -> None: + """Multi-frame: set-equality holds per frame.""" + pe = self._pair_excl([(0, 1), (1, 0)]) + coord2 = np.concatenate([self.coord, self.coord + 0.5], axis=0) + atype2 = np.concatenate([self.atype, self.atype], axis=0) + ng_base = build_neighbor_graph(coord2, atype2, None, self.rcut) + atype_flat = atype2.reshape(-1) + ng_post = apply_pair_exclusion(ng_base, atype_flat, pe) + ng_fused = build_neighbor_graph(coord2, atype2, None, self.rcut, pair_excl=pe) + self.assertEqual(valid_edge_set(ng_post), valid_edge_set(ng_fused)) + + +class TestBuildNeighborGraphAseOracle(unittest.TestCase): + """Oracle harness for the ASE builder pair_excl parameter. + + Skipped when ``ase`` is not installed. Asserts set-equality of the + valid-edge set between the ASE builder called with ``pair_excl`` and + the dense reference builder + separate :func:`apply_pair_exclusion`. + """ + + @classmethod + def setUpClass(cls) -> None: + try: + import ase # noqa: F401 + except ImportError as e: + raise unittest.SkipTest("ase not installed") from e + + def setUp(self) -> None: + self.rcut = 4.0 + self.coord = np.array( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.3, 0.0], [3.5, 0.0, 0.0]], + dtype=np.float64, + ).reshape(1, 4, 3) + self.atype = np.array([[0, 1, 0, 1]], dtype=np.int64) + self.ntypes = 2 + + def _pair_excl(self, exclude_pairs): + return PairExcludeMask(self.ntypes, exclude_pairs) + + def test_ase_pair_excl_none_identity(self) -> None: + """pair_excl=None: ASE builder output unchanged.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph_ase, + ) + + ng_ref = build_neighbor_graph_ase(self.coord, self.atype, None, self.rcut) + ng_excl = build_neighbor_graph_ase( + self.coord, self.atype, None, self.rcut, pair_excl=None + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_excl)) + + def test_ase_oracle_set_equality(self) -> None: + """ASE builder with pair_excl == dense ref + apply_pair_exclusion.""" + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph_ase, + ) + + pe = self._pair_excl([(0, 1), (1, 0)]) + # dense reference + separate post-process + ng_dense = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + atype_flat = self.atype.reshape(-1) + ng_ref = apply_pair_exclusion(ng_dense, atype_flat, pe) + # ASE builder with fused post-process + ng_ase = build_neighbor_graph_ase( + self.coord, self.atype, None, self.rcut, pair_excl=pe + ) + self.assertEqual(valid_edge_set(ng_ref), valid_edge_set(ng_ase)) + # exclusion actually removed edges + ng_ase_plain = build_neighbor_graph_ase(self.coord, self.atype, None, self.rcut) + self.assertLess(int(ng_ase.edge_mask.sum()), int(ng_ase_plain.edge_mask.sum())) + + +# NOTE: nvalchemiops builder has no local oracle set-equality test for pair_excl +# because it requires CUDA; validation is deferred to GPU box tests (PR-C/nv-gtest). +# See deepmd.pt_expt.utils.nv_graph_builder.build_neighbor_graph_nv docstring. + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/common/dpmodel/test_spin_model_legacy_routing.py b/source/tests/common/dpmodel/test_spin_model_legacy_routing.py new file mode 100644 index 0000000000..1a309a1ca4 --- /dev/null +++ b/source/tests/common/dpmodel/test_spin_model_legacy_routing.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression test: SpinModel backbone must stay on the legacy dense-nlist path. + +A spin system uses virtual/placeholder types whose pair exclusions double the +effective atom density. Before the fix (commit 6c2b007c9), the dpmodel +``call_common`` auto-flip (decision #17) moved graph-eligible mixed_types +backbones to the carry-all graph route. When a spin model's backbone +(dpa1 attn_layer=0) crossed to the graph route, its ``call_common`` diverged +from the dense lower interface (``call_common_lower``) used by pt_expt eager +inference -- the carry-all graph keeps neighbors the capped dense nlist +discards. + +Fix: ``SpinModel.call_common`` forces ``neighbor_graph_method="legacy"`` on +the backbone call so the spin backbone always runs dense regardless of the +default flip. + +This test pins that contract: + +1. ``SpinModel.call_common`` energy == backbone energy computed with explicit + ``neighbor_graph_method="legacy"`` on the spin-doubled coordinate inputs. +2. The backbone in graph mode (``neighbor_graph_method="ase"``) on the same + doubled inputs gives a DIFFERENT energy, confirming the fixture is + sel-binding (i.e., the virtual-atom density really does trigger divergence). +""" + +import numpy as np +import pytest + +from deepmd.dpmodel.model.model import ( + get_spin_model, +) + + +def _spin_dpa1_config() -> dict: + """Minimal dpa1-backed spin model config (sel-binding on doubled inputs).""" + return { + "type": "standard", + "type_map": ["Fe", "H"], + "spin": { + "use_spin": [True, False], + "virtual_scale": [0.3], + }, + "descriptor": { + "type": "dpa1", + "rcut": 4.0, + "rcut_smth": 0.5, + # Small sel: with 8 real + 8 virtual = 16 atoms in a 6 Å box, + # sel=8 is well below the average neighbor count → sel-binding. + "sel": 8, + "ntypes": 4, # expanded to 4 by get_spin_model (2 real + 2 virtual) + "attn_layer": 0, + "axis_neuron": 2, + "neuron": [4, 8], + "seed": 0, + }, + "fitting_net": { + "type": "ener", + "neuron": [4, 4], + "seed": 0, + }, + } + + +def _make_test_frame(rng: np.random.Generator): + """Return (coord, atype, box) for a small PBC cell with 8 atoms.""" + natoms = 8 # 4 Fe + 4 H + coord = rng.random((1, natoms, 3)) * 4.0 # 1 frame + atype = np.array([[0, 0, 0, 0, 1, 1, 1, 1]]) + box = np.eye(3).reshape(1, 9) * 6.0 + return coord, atype, box + + +def test_spin_model_backbone_routes_legacy() -> None: + """SpinModel.call_common energy must equal backbone with explicit legacy. + + Procedure: + - Build the SpinModel and obtain the doubled coords/atypes via + ``process_spin_input`` (the same transform used internally). + - Call the backbone directly with ``neighbor_graph_method="legacy"`` on + those doubled inputs to get the expected dense-path energy. + - Call ``SpinModel.call_common`` and extract its energy. + - Assert the two energies are exactly equal. + - Additionally assert the backbone in ``neighbor_graph_method="ase"`` mode + on the same doubled inputs gives a DIFFERENT energy (sel-binding guard). + """ + pytest.importorskip("ase") # ase builder needed for divergence check + + rng = np.random.default_rng(42) + coord, atype, box = _make_test_frame(rng) + spin = np.zeros_like(coord) + spin[:, :4, 2] = 1.0 # spin-z on Fe atoms only + + model = get_spin_model(_spin_dpa1_config()) + backbone = model.backbone_model + + # --- Get doubled inputs via the model's own transform --- + coord_doubled, atype_doubled, _corr = model.process_spin_input(coord, atype, spin) + + # The backbone descriptor has graph-lower disabled by the spin-model + # opt-out knob. Temporarily re-enable it only for the sel-binding + # divergence probe below; the legacy/spin comparison uses the dense path. + backbone_descriptor = backbone.get_dp_atomic_model().descriptor + assert backbone_descriptor._graph_lower_disabled is True + + # --- Backbone with explicit legacy routing --- + legacy_ret = backbone.call_common( + coord_doubled, atype_doubled, box, neighbor_graph_method="legacy" + ) + legacy_energy = np.array(legacy_ret["energy"]) + + # --- SpinModel.call_common (must route legacy internally) --- + spin_ret = model.call_common(coord, atype, spin, box) + spin_energy = np.array(spin_ret["energy"]) + + # ``energy`` here is per-atom energy; the backbone returns all 2*nloc + # atoms while the spin model truncates to the first nloc real atoms. + # Compare the total energy (sum over all atoms) which should be equal + # because the virtual-atom energies are zeroed by the exclusion mask. + np.testing.assert_allclose( + float(spin_energy.sum()), + float(legacy_energy.sum()), + rtol=0, + atol=0, + err_msg=( + "SpinModel.call_common total energy does not match backbone(legacy) " + "on doubled inputs; the spin model may be routing through the " + "carry-all graph instead of the legacy dense path." + ), + ) + + # --- Sel-binding guard: backbone graph must DIFFER from backbone legacy --- + # Re-enable graph lower on the backbone descriptor only for this probe. + backbone_descriptor._graph_lower_disabled = False + try: + graph_ret = backbone.call_common( + coord_doubled, atype_doubled, box, neighbor_graph_method="ase" + ) + finally: + backbone_descriptor._graph_lower_disabled = True + graph_energy = np.array(graph_ret["energy"]) + + assert not np.allclose(legacy_energy.sum(), graph_energy.sum(), rtol=1e-6), ( + "Backbone legacy and graph give the same energy on the doubled spin " + "system — sel is not binding with these inputs; the regression fixture " + "is too weak. Reduce sel or increase atom density." + ) + + +def test_spin_backbone_descriptor_graph_lower_disabled() -> None: + """The spin backbone descriptor must have graph-lower disabled. + + The explicit ``disable_graph_lower`` opt-out knob is set structurally at + spin-model construction and must survive a serialize -> deserialize round + trip (the flag is NOT part of the serialized schema, so it must be + re-derived on deserialization). + """ + model = get_spin_model(_spin_dpa1_config()) + descriptor = model.backbone_model.get_dp_atomic_model().descriptor + # dpa1 with attn_layer=0 concat tebd would otherwise be graph-eligible. + assert descriptor._graph_lower_disabled is True + assert descriptor.uses_graph_lower() is False + + # --- survive serialize -> deserialize --- + data = model.serialize() + from deepmd.dpmodel.model.spin_model import ( + SpinModel, + ) + + model2 = SpinModel.deserialize(data) + descriptor2 = model2.backbone_model.get_dp_atomic_model().descriptor + assert descriptor2._graph_lower_disabled is True + assert descriptor2.uses_graph_lower() is False + + +def test_spin_model_call_common_deterministic() -> None: + """SpinModel.call_common is deterministic (no stochastic routing).""" + rng = np.random.default_rng(7) + coord, atype, box = _make_test_frame(rng) + spin = np.zeros_like(coord) + spin[:, :4, 2] = 1.0 + + model = get_spin_model(_spin_dpa1_config()) + + ret1 = model.call_common(coord, atype, spin, box) + ret2 = model.call_common(coord, atype, spin, box) + + np.testing.assert_array_equal( + float(np.array(ret1["energy"]).sum()), + float(np.array(ret2["energy"]).sum()), + err_msg="SpinModel.call_common is non-deterministic", + ) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 9601ac7607..be1808c0a9 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -311,3 +311,59 @@ def setUp(self) -> None: def tearDown(self) -> None: IOTest.tearDown(self) + + +@unittest.skipIf( + not DP_TEST_TF2_ONLY, + "pair_exclude_types is not supported by the TF v1 backend; it is validated " + "in the TF2-only job, where test_deep_eval also exercises the jax2tf " + "'.savedmodel' export path (see the backend table in test_deep_eval).", +) +class TestDeepPotPairExclude(unittest.TestCase, IOTest): + """Model-level ``pair_exclude_types`` is a nlist-BUILD transform (decision + #18/A4). Every backend folds it in where the neighbor list is built (the + jax2tf ``.savedmodel`` export reuses the dpmodel + ``apply_pair_exclusion_nlist`` via the ``ndtensorflow`` namespace), so the + exported models must still eval-agree across backends. + """ + + def setUp(self) -> None: + model_def_script = { + "type_map": ["O", "H"], + "pair_exclude_types": [[0, 1]], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 3, + 6, + ], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "ener", + "neuron": [ + 5, + 5, + ], + "resnet_dt": True, + "precision": "float64", + "atom_ener": [], + "seed": 1, + }, + } + model = get_model(copy.deepcopy(model_def_script)) + self.data = { + "model": model.serialize(), + "backend": "test", + "model_def_script": model_def_script, + } + + def tearDown(self) -> None: + IOTest.tearDown(self) diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index 912b897cd6..9522122d91 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -11,6 +11,9 @@ ) from deepmd.dpmodel.model.ener_model import EnergyModel as EnergyModelDP from deepmd.dpmodel.model.model import get_model as get_model_dp +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.nlist import ( build_neighbor_list, extend_coord_with_ghosts, @@ -484,6 +487,14 @@ def setUp(self) -> None: 6.0, [20, 20], distinguish_types=True, + # model-level pair exclusion is a nlist-BUILD transform (decision + # #18/A4): the dpmodel-family lowers consume a pre-excluded nlist; + # legacy pt/pd re-apply internally, which is an idempotent no-op. + pair_excl=PairExcludeMask( + self.ntypes, [tuple(p) for p in self.data["pair_exclude_types"]] + ) + if self.data["pair_exclude_types"] + else None, ) extended_coord = extended_coord.reshape(nframes, -1, 3) self.nlist = nlist diff --git a/source/tests/infer/gen_common.py b/source/tests/infer/gen_common.py index fea29fc63e..4ea27feb32 100644 --- a/source/tests/infer/gen_common.py +++ b/source/tests/infer/gen_common.py @@ -2,8 +2,50 @@ """Common utilities shared by model generation scripts (gen_*.py).""" import glob +import json import os import sys +import zipfile + + +def derive_pair_exclude_pt2(src_pt2, dst_pt2, exclude_types): + """Derive a model-level ``pair_exclude_types`` variant of a ``.pt2`` archive. + + Parameters + ---------- + src_pt2 : str + Path to the no-exclusion baseline ``.pt2``. Must have the same weights + and ``lower_kind`` as the desired variant (only the exclusion list may + differ); a graph-lower baseline cannot derive an nlist-lower variant. + dst_pt2 : str + Path to write the derived ``.pt2``. + exclude_types : list[list[int]] + The ``pair_exclude_types`` list to inject, e.g. ``[[0, 1]]``. + + Notes + ----- + Model-level ``pair_exclude_types`` is a build-time transform (decision + #18/A4) applied at the C++ ingestion seam (from ``metadata.json``) and the + Python DeepEval build seam (from the serialized ``model.json``); it is not + baked into the exported graph, so the compiled AOTI artifact (including any + nested with-comm ``.pt2``) is byte-identical to the baseline. Only the two + JSON blobs that carry the list are patched -- both, so the C++ and Python + inference paths agree; patching one alone makes them disagree. This avoids + a second inductor compile and is bit-identical to a full recompile. + """ + with zipfile.ZipFile(src_pt2) as zin: + entries = [(info, zin.read(info.filename)) for info in zin.infolist()] + with zipfile.ZipFile(dst_pt2, "w") as zout: + for info, blob in entries: + if info.filename == "model/extra/metadata.json": + meta = json.loads(blob) + meta["pair_exclude_types"] = exclude_types + blob = json.dumps(meta).encode() + elif info.filename == "model/extra/model.json": + mdl = json.loads(blob) + mdl["model"]["pair_exclude_types"] = exclude_types + blob = json.dumps(mdl).encode() + zout.writestr(info, blob) def ensure_inductor_compiler(): diff --git a/source/tests/infer/gen_dpa1_pairexcl.py b/source/tests/infer/gen_dpa1_pairexcl.py new file mode 100644 index 0000000000..452a7b4e86 --- /dev/null +++ b/source/tests/infer/gen_dpa1_pairexcl.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate DPA1 test models carrying model-level ``pair_exclude_types``. + +Produces two graph-eligible DPA1(attn_layer=0) models with identical weights, +one per C++ ingestion route, plus a no-exclusion baseline: + + - deeppot_dpa1_pairexcl_graph.pt2 (lower_kind="graph", pair_exclude=[[0,1]]) + - deeppot_dpa1_pairexcl_nlist.pt2 (lower_kind="nlist", pair_exclude=[[0,1]]) + - deeppot_dpa1_pairexcl_none.pt2 (lower_kind="graph", NO exclusion) + +Only two inductor compiles are needed: the graph baseline (``_none``) and the +nlist route (``_nlist``). ``_graph`` has the same weights and lower_kind as +``_none`` and differs only in the exclusion list, so it is DERIVED from +``_none`` by patching the archive (``gen_common.derive_pair_exclude_pt2``) -- +no third compile. + +The pair models exercise the C++ pair-exclusion ownership (decision #18/A4: +exclusion is a BUILD-time transform on BOTH routes) plus the +``pair_exclude_types`` metadata round-trip in ``DeepPotPTExpt::init``: + +- graph route: folded into ``edge_mask`` at build by ``applyPairExclusion`` + (C++) / the NeighborGraph builder (Python DeepEval). +- dense route: folded into the nlist at build by ``applyPairExclusionNlist`` + (C++) / ``build_neighbor_list(pair_excl=...)`` (Python DeepEval). + +The exported lowers consume pre-excluded inputs and never re-apply the +exclusion. The gtest validates C++ energy/force vs the Python ``DeepEval`` +reference at 1e-10 and, by comparing against the ``_none`` baseline, confirms +the exclusion is actually active. + +Reference sidecar files (.expected) consumed by the C++ gtest are written from +the Python ``DeepEval`` evaluation of each pair model (PBC + NoPBC sections). + +exclude_types = [[0, 1]] drops every O-H (cross-type) interaction while keeping +O-O and H-H, so both energy and forces change measurably but stay non-degenerate +(H-H pairs survive). + +This is a SEPARATE script from ``gen_dpa1.py`` so it can be run independently on +boxes where the unrelated attn_layer=2 model in ``gen_dpa1.py`` trips a torch +inductor codegen bug. +""" + +import copy +import os +import sys + +import numpy as np + +# Ensure the source tree is on the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) + +from gen_common import ( + derive_pair_exclude_pt2, + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + + +def main(): + from deepmd.dpmodel.model.model import ( + get_model, + ) + + ensure_inductor_compiler() + + base_dir = os.path.dirname(__file__) + + # Graph-eligible DPA1 (attn_layer=0); same descriptor as gen_dpa1.py + # Section B so the two fixtures stay comparable. + base_config = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut_smth": 2.0, + "rcut": 6.0, + "neuron": [2, 4, 8], + "axis_neuron": 4, + "attn": 5, + "attn_layer": 0, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "scaling_factor": 1.0, + "normalize": True, + "temperature": 1.0, + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [5, 5, 5], + "resnet_dt": True, + "seed": 1, + }, + } + + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file as pt_expt_deserialize_to_file, + ) + + load_custom_ops() + + from deepmd.infer import ( + DeepPot, + ) + + coord = np.array( + [ + 12.83, 2.56, 2.18, 12.09, 2.87, 2.74, 0.25, 3.32, 1.68, + 3.36, 3.00, 1.81, 3.51, 2.51, 2.60, 4.27, 3.22, 1.56, + ], + dtype=np.float64, + ) # fmt: skip + atype = [0, 1, 1, 0, 1, 1] + box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], dtype=np.float64) + + def build_data(exclude_types): + cfg = copy.deepcopy(base_config) + if exclude_types: + cfg["pair_exclude_types"] = exclude_types + model = get_model(copy.deepcopy(cfg)) + return { + "model": copy.deepcopy(model.serialize()), + "model_def_script": cfg, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + # ---- No-exclusion baseline (graph route) ---- + none_pt2 = os.path.join(base_dir, "deeppot_dpa1_pairexcl_none.pt2") + print(f"Exporting no-exclusion baseline to {none_pt2} ...") # noqa: T201 + pt_expt_deserialize_to_file( + none_pt2, build_data(None), do_atomic_virial=True, lower_kind="graph" + ) + dp_none = DeepPot(none_pt2) + e_none = dp_none.eval(coord, box, atype, atomic=False)[0][0, 0] + print(f" baseline PBC energy: {e_none:.18e}") # noqa: T201 + + # ---- Pair-exclusion models (graph + dense routes) ---- + exclude_types = [[0, 1]] + # graph route: SAME weights + graph lower as the no-exclusion baseline + # above, only the exclusion list differs -> derive by patching the archive, + # NO second inductor compile. See gen_common.derive_pair_exclude_pt2. + graph_pt2 = os.path.join(base_dir, "deeppot_dpa1_pairexcl_graph.pt2") + print( # noqa: T201 + f"Deriving {graph_pt2} from {none_pt2} " + "(pair_exclude_types patch, no recompile) ..." + ) + derive_pair_exclude_pt2(none_pt2, graph_pt2, exclude_types) + # nlist route: a DIFFERENT lower_kind (a different exported graph) that the + # graph baseline cannot supply -> compile it. + nlist_pt2 = os.path.join(base_dir, "deeppot_dpa1_pairexcl_nlist.pt2") + print( # noqa: T201 + f"Exporting to {nlist_pt2} (lower_kind='nlist', " + f"pair_exclude_types={exclude_types}) ..." + ) + pt_expt_deserialize_to_file( + nlist_pt2, + copy.deepcopy(build_data(exclude_types)), + do_atomic_virial=True, + lower_kind="nlist", + ) + # Eval + activeness check + reference sidecar for both routes. + for tag in ("graph", "nlist"): + pt2_path = os.path.join(base_dir, f"deeppot_dpa1_pairexcl_{tag}.pt2") + dp_e = DeepPot(pt2_path) + e_e1, f_e1, v_e1, ae_e1, av_e1 = dp_e.eval(coord, box, atype, atomic=True) + e_enp, f_enp, v_enp, ae_enp, av_enp = dp_e.eval(coord, None, atype, atomic=True) + + # Confirm the exclusion is ACTIVE: energy must differ from the + # no-exclusion baseline (identical weights minus pair_exclude_types). + e_diff = float(abs(e_e1[0, 0] - e_none)) + print(f" {tag}: |E(excl) - E(no-excl)| = {e_diff:.3e}") # noqa: T201 + if e_diff < 1e-6: + raise RuntimeError( + f"BLOCKED: pair_exclude_types had no effect on the {tag} model " + f"(energy delta {e_diff:.2e} < 1e-6); exclusion may be dropped." + ) + f_max = float(np.max(np.abs(f_e1))) + if f_max < 1e-10: + raise RuntimeError( + f"Pair-exclude {tag} forces are degenerate (max={f_max:.2e})." + ) + + ref_path = os.path.join(base_dir, f"deeppot_dpa1_pairexcl_{tag}.expected") + write_expected_ref( + ref_path, + sections={ + "pbc": { + "expected_e": ae_e1[0, :, 0], + "expected_f": f_e1[0], + "expected_v": av_e1[0], + }, + "nopbc": { + "expected_e": ae_enp[0, :, 0], + "expected_f": f_enp[0], + "expected_v": av_enp[0], + }, + }, + source_script="source/tests/infer/gen_dpa1_pairexcl.py", + ) + print(f" Wrote {ref_path}") # noqa: T201 + + print("\nAll pair-exclude models generated.") # noqa: T201 + print("Done!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/infer/gen_dpa3.py b/source/tests/infer/gen_dpa3.py index c0a4434e33..feb064eb37 100644 --- a/source/tests/infer/gen_dpa3.py +++ b/source/tests/infer/gen_dpa3.py @@ -17,6 +17,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) from gen_common import ( + derive_pair_exclude_pt2, ensure_inductor_compiler, load_custom_ops, write_expected_ref, @@ -111,6 +112,19 @@ def main(): pt2_mpi_path, copy.deepcopy(data_mpi), do_atomic_virial=True ) + # Multi-rank variant WITH model-level ``pair_exclude_types`` — derived + # CHEAPLY from ``deeppot_dpa3_mpi.pt2`` by patching the exclusion list into + # the archive, with NO second inductor compile (same weights + lower_kind, + # only the exclusion list differs). ``deeppot_dpa3_mpi.pt2`` is the exact + # no-exclusion baseline. See ``derive_pair_exclude_pt2`` for why this is + # sound, and test_lammps_dpa3_pt2.py::test_pair_deepmd_mpi_dpa3_pairexcl_*. + pt2_pairexcl_mpi_path = os.path.join(base_dir, "deeppot_dpa3_pairexcl_mpi.pt2") + print( # noqa: T201 + f"Deriving {pt2_pairexcl_mpi_path} from {pt2_mpi_path} " + "(pair_exclude_types patch, no recompile) ..." + ) + derive_pair_exclude_pt2(pt2_mpi_path, pt2_pairexcl_mpi_path, [[0, 1]]) + # Float32 multi-rank variant — same architecture as the float64 # MPI fixture but with ``precision: float32``. Used by # source/lmp/tests/test_lammps_dpa3_pt2_fp32.py to validate that diff --git a/source/tests/jax/test_jax_md.py b/source/tests/jax/test_jax_md.py index ce4b7288c6..557da898ac 100644 --- a/source/tests/jax/test_jax_md.py +++ b/source/tests/jax/test_jax_md.py @@ -109,6 +109,36 @@ def test_dense_neighbor_uses_jax_md_displacement_convention(): np.testing.assert_allclose(potential(coord, neighbor=neighbor), 0.04, atol=1e-12) +def test_dense_neighbor_applies_model_pair_exclusion(): + """Model-level ``pair_exclude_types`` must be folded into the JAX-MD nlist + at the ``call_lower`` ingestion seam (decision #18/A4): the lower no longer + re-applies it, so without this the JAX-MD path would silently INCLUDE + excluded pairs (fail-open). Types are (0, 1); excluding that pair drops the + only edge, so the energy is 0 instead of the 0.04 the same geometry gives + without exclusion. + """ + from types import ( + SimpleNamespace, + ) + + from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, + ) + + class ExclEdgeModel(EdgeModel): + atomic_model = SimpleNamespace(pair_excl=PairExcludeMask(2, [(0, 1)])) + + disp = lambda ra, rb: (ra - rb) - 10.0 * jnp.round((ra - rb) / 10.0) # noqa: E731 + coord = jnp.asarray([[0.1, 0.0, 0.0], [9.9, 0.0, 0.0]]) + neighbor = DenseNeighbor(jnp.asarray([[1], [0]], dtype=jnp.int32)) + + excl = energy_fn(ExclEdgeModel(), [0, 1], displacement_fn=disp) + np.testing.assert_allclose(excl(coord, neighbor=neighbor), 0.0, atol=1e-12) + # control: identical geometry WITHOUT exclusion keeps the edge (0.04). + noexcl = energy_fn(EdgeModel(), [0, 1], displacement_fn=disp) + np.testing.assert_allclose(noexcl(coord, neighbor=neighbor), 0.04, atol=1e-12) + + def test_dense_neighbor_rejects_scalar_metric(): potential = energy_fn( EdgeModel(), diff --git a/source/tests/pt_expt/descriptor/test_dpa1.py b/source/tests/pt_expt/descriptor/test_dpa1.py index 2bfde949c1..9ab2cc7b6b 100644 --- a/source/tests/pt_expt/descriptor/test_dpa1.py +++ b/source/tests/pt_expt/descriptor/test_dpa1.py @@ -114,7 +114,10 @@ def test_consistency(self, idt, sm, to, tm, prec, ect) -> None: @pytest.mark.parametrize("idt", [False, True]) # resnet_dt @pytest.mark.parametrize("prec", ["float64", "float32"]) # precision - def test_exportable(self, idt, prec) -> None: + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion + def test_exportable(self, idt, prec, excl_types) -> None: rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape davg = rng.normal(size=(self.nt, nnei, 4)) @@ -130,6 +133,7 @@ def test_exportable(self, idt, prec) -> None: attn_layer=2, precision=prec, resnet_dt=idt, + exclude_types=excl_types, seed=GLOBAL_SEED, ).to(self.device) dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device) @@ -252,15 +256,20 @@ def fn(coord_ext, atype_ext, nlist): atol=atol, ) + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion @pytest.mark.parametrize("tm", ["concat", "strip"]) # tebd_input_mode @pytest.mark.parametrize("prec", ["float64"]) # precision - def test_make_fx_graph(self, prec, tm) -> None: + def test_make_fx_graph(self, prec, tm, excl_types) -> None: """make_fx (export-readiness) of the attn_layer=0 GRAPH forward. For ``attn_layer == 0`` the dense ``forward`` routes through the graph-native path (``from_dense_quartet -> call_graph``). This proves that graph forward + ``autograd.grad`` is fx-traceable (full .pt2 - export is PR-B). + export is PR-B). Parametrized over ``excl_types``: with non-empty + exclusion the ``build_edge_exclude_mask`` (mask-only, shape-static) + path is exercised — a tracing failure here would be a bug. """ rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape @@ -278,6 +287,7 @@ def test_make_fx_graph(self, prec, tm) -> None: attn_layer=0, tebd_input_mode=tm, precision=prec, + exclude_types=excl_types, seed=GLOBAL_SEED, ).to(self.device) dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device) @@ -313,10 +323,13 @@ def fn(coord_ext, atype_ext, nlist, mapping): atol=atol, ) + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion @pytest.mark.parametrize("tm", ["concat", "strip"]) # tebd_input_mode @pytest.mark.parametrize("smooth", [False, True]) # smooth attention branch @pytest.mark.parametrize("prec", ["float64"]) # precision - def test_make_fx_graph_attn(self, prec, smooth, tm) -> None: + def test_make_fx_graph_attn(self, prec, smooth, tm, excl_types) -> None: """make_fx (export-readiness) of the GRAPH forward with attention. MERGE BLOCKER (NeighborGraph PR-D): pt_expt compiled training routes @@ -324,6 +337,9 @@ def test_make_fx_graph_attn(self, prec, smooth, tm) -> None: (``attn_layer > 0``) must be fx-traceable — the shape-static ``center_edge_pairs`` form keeps the pair enumeration ``nonzero``-free. Covers both the smooth and non-smooth attention branches. + Parametrized over ``excl_types``: with non-empty exclusion the + ``build_edge_exclude_mask`` (mask-only, shape-static) path is exercised + concurrently with attention — a tracing failure here would be a bug. """ rng = np.random.default_rng(GLOBAL_SEED) _, _, nnei = self.nlist.shape @@ -342,6 +358,7 @@ def test_make_fx_graph_attn(self, prec, smooth, tm) -> None: attn_dotr=True, smooth_type_embedding=smooth, precision=prec, + exclude_types=excl_types, seed=GLOBAL_SEED, ).to(self.device) dd0.se_atten.mean = torch.tensor(davg, dtype=dtype, device=self.device) diff --git a/source/tests/pt_expt/model/test_dpa1_graph_lower.py b/source/tests/pt_expt/model/test_dpa1_graph_lower.py index 11c49cdca2..d85e6dae27 100644 --- a/source/tests/pt_expt/model/test_dpa1_graph_lower.py +++ b/source/tests/pt_expt/model/test_dpa1_graph_lower.py @@ -15,14 +15,18 @@ reduced (per-frame) virial are frame/local quantities and compare directly. """ +import copy + import numpy as np import pytest import torch from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, from_dense_quartet, ) from deepmd.dpmodel.utils.nlist import ( + apply_pair_exclusion_nlist, build_neighbor_list, extend_coord_with_ghosts, ) @@ -91,7 +95,13 @@ def setup_method(self) -> None: [[0, 0, 0, 1, 1]], dtype=torch.int64, device=self.device ) - def _make_model(self, attn_layer: int = 0, smooth: bool = False) -> EnergyModel: + def _make_model( + self, + attn_layer: int = 0, + smooth: bool = False, + pair_excl_types: list | None = None, + descr_excl_types: list | None = None, + ) -> EnergyModel: ds = DescrptDPA1( self.rcut, self.rcut_smth, @@ -111,6 +121,7 @@ def _make_model(self, attn_layer: int = 0, smooth: bool = False) -> EnergyModel: set_davg_zero=False, type_one_side=True, precision="float64", + exclude_types=descr_excl_types or [], seed=GLOBAL_SEED, ).to(self.device) ft = InvarFitting( @@ -122,7 +133,12 @@ def _make_model(self, attn_layer: int = 0, smooth: bool = False) -> EnergyModel: precision="float64", seed=GLOBAL_SEED, ).to(self.device) - return EnergyModel(ds, ft, type_map=self.type_map).to(self.device) + return EnergyModel( + ds, + ft, + type_map=self.type_map, + pair_exclude_types=pair_excl_types or [], + ).to(self.device) def _prepare_lower_inputs(self, periodic: bool): """Build extended coords, atype, nlist, mapping as torch tensors.""" @@ -172,13 +188,20 @@ def _prepare_lower_inputs(self, periodic: bool): @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention @pytest.mark.parametrize("periodic", [True, False]) # PBC vs non-PBC @pytest.mark.parametrize("do_av", [False, True]) # atom-virial off / on - def test_force_virial_parity_vs_legacy(self, periodic, do_av, attn_layer) -> None: + @pytest.mark.parametrize( + "excl_types", [[], [(0, 1)]] + ) # no exclusion / type-0-1 pair exclusion + def test_force_virial_parity_vs_legacy( + self, periodic, do_av, attn_layer, excl_types + ) -> None: """Graph lower energy/force/virial/atom_virial == legacy dense lower on the SAME neighbor set (regime-1 graph from from_dense_quartet). attn_layer=2 exercises graph attention through model-level autograd (smooth=False: exact carry-all parity regime, NeighborGraph PR-D). + Parametrized over exclude_types: empty list (no exclusion) and + [(0,1)] (model-level pair exclusion applied identically on both routes). """ - model = self._make_model(attn_layer=attn_layer) + model = self._make_model(attn_layer=attn_layer, pair_excl_types=excl_types) model.eval() tol = ( {"rtol": 1e-12, "atol": 1e-12} @@ -189,6 +212,13 @@ def test_force_virial_parity_vs_legacy(self, periodic, do_av, attn_layer) -> Non nf = ext_coord.shape[0] nloc = self.natoms + # Model-level pair_exclude is a nlist-BUILD transform (decision + # #18/A4): BOTH lowers consume pre-excluded inputs, so fold the + # exclusion into the dense nlist here (mirrors the C++ + # ``applyPairExclusionNlist`` build step). + nlist = apply_pair_exclusion_nlist( + nlist, ext_atype, model.atomic_model.pair_excl + ) legacy = model.forward_common_lower( ext_coord.clone().requires_grad_(True), ext_atype, @@ -202,6 +232,10 @@ def test_force_virial_parity_vs_legacy(self, periodic, do_av, attn_layer) -> Non # returned edge_vec is already a torch tensor on env.DEVICE. ng = from_dense_quartet(ext_coord, nlist, mapping) atype_local = ext_atype[:, :nloc].reshape(nf * nloc) + # Model-level pair_exclude is a BUILD-time transform (decision #18): the + # converter does not bake it in and the lower no longer re-applies it, so + # apply it to the graph here (mirrors the C++ ``applyPairExclusion`` step). + ng = apply_pair_exclusion(ng, atype_local, model.atomic_model.pair_excl) graph = model.forward_common_lower_graph( atype_local, ng.n_node, @@ -330,6 +364,129 @@ def test_smooth_attention_divergence_pinned(self) -> None: assert e_diff < 1e-3, f"smooth divergence too large: {e_diff:.3e}" assert f_diff < 1e-3, f"smooth force divergence too large: {f_diff:.3e}" + def test_pair_exclude_types_graph_vs_legacy(self) -> None: + """Model-level pair_exclude_types: graph route and legacy dense agree + bit-tight (fp64, 1e-12), AND the excluded model output differs from the + no-exclude baseline (exclusion is not vacuous). + + Strategy: build the no-exclude model, serialize it, inject + ``pair_exclude_types=[[0,1]]`` into the serialized dict, deserialize + to get an exclude model with IDENTICAL weights, then run both routes. + """ + import copy + + # 1. build the reference (no-exclude) model + model_ref = self._make_model(attn_layer=0) + model_ref.eval() + + # 2. derive the exclude model by patching the serialized dict + data = copy.deepcopy(model_ref.serialize()) + data["pair_exclude_types"] = [[0, 1]] + model_excl = EnergyModel.deserialize(data).to(self.device) + model_excl.eval() + + tol = ( + {"rtol": 1e-12, "atol": 1e-12} + if self.device.type == "cpu" + else {"rtol": 1e-10, "atol": 1e-10} + ) + box = self.cell.reshape(1, 9) + + # 3. graph route (build-time pair exclusion) + graph_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + # 4. legacy dense route (seam backstop in forward_atomic_graph) + legacy_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + # parity: graph == legacy + torch.testing.assert_close( + graph_out["energy_redu"], legacy_out["energy_redu"], **tol + ) + torch.testing.assert_close( + graph_out["energy_derv_r"], legacy_out["energy_derv_r"], **tol + ) + + # 5. reference (no-exclude) via graph route + ref_out = model_ref.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + # exclusion must have an effect + e_diff = (graph_out["energy_redu"] - ref_out["energy_redu"]).abs().max().item() + assert e_diff > 1e-10, ( + f"pair_exclude_types had no effect on energy; diff={e_diff:.3e}" + ) + + @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention + def test_descriptor_exclude_types_graph_vs_legacy(self, attn_layer) -> None: + """Descriptor-level exclude_types: graph route and legacy dense agree + bit-tight (fp64, 1e-12) when exclusion is on the DESCRIPTOR (not the + model pair_exclude_types). Uses identical weights across both routes; + also checks exclusion is non-vacuous vs a no-exclude baseline. + """ + # 1. no-exclude model (graph route = reference) + model_ref = self._make_model(attn_layer=attn_layer) + model_ref.eval() + + # 2. exclude model: inject exclude_types into the serialized dict + data = copy.deepcopy(model_ref.serialize()) + data["descriptor"]["exclude_types"] = [[0, 1]] + model_excl = EnergyModel.deserialize(data).to(self.device) + model_excl.eval() + + tol = ( + {"rtol": 1e-12, "atol": 1e-12} + if self.device.type == "cpu" + else {"rtol": 1e-10, "atol": 1e-10} + ) + box = self.cell.reshape(1, 9) + + # 3. graph route + graph_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + # 4. legacy dense route + legacy_out = model_excl.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="legacy", + ) + torch.testing.assert_close( + graph_out["energy_redu"], legacy_out["energy_redu"], **tol + ) + torch.testing.assert_close( + graph_out["energy_derv_r"], legacy_out["energy_derv_r"], **tol + ) + torch.testing.assert_close( + graph_out["energy_derv_c_redu"], legacy_out["energy_derv_c_redu"], **tol + ) + + # 5. exclusion must be non-vacuous + ref_out = model_ref.call_common( + self.coord.clone().requires_grad_(True), + self.atype, + box, + neighbor_graph_method="dense", + ) + e_diff = (graph_out["energy_redu"] - ref_out["energy_redu"]).abs().max().item() + assert e_diff > 1e-10, ( + f"descriptor exclude_types had no effect on energy; diff={e_diff:.3e}" + ) + @pytest.mark.parametrize("attn_layer", [0, 2]) # factorizable AND attention def test_graph_route_float32(self, attn_layer) -> None: """A float32 model runs the graph route and matches the dense route. diff --git a/source/tests/pt_expt/utils/test_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py index dea25adab3..c216676994 100644 --- a/source/tests/pt_expt/utils/test_vesin_graph_builder.py +++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py @@ -3,7 +3,11 @@ import pytest import torch +from deepmd.dpmodel.utils.exclude_mask import ( + PairExcludeMask, +) from deepmd.dpmodel.utils.neighbor_graph import ( + apply_pair_exclusion, build_neighbor_graph, ) @@ -105,3 +109,66 @@ def test_vesin_excludes_virtual_atoms_like_dense(): ei = np.asarray(ng.edge_index)[:, np.asarray(ng.edge_mask)] at = atype.reshape(-1).numpy() assert np.all(at[ei[0]] >= 0) and np.all(at[ei[1]] >= 0) + + +def _valid_edge_set(ng): + """Return the set of (src, dst, rounded edge_vec) for all real edges.""" + ei = np.asarray(ng.edge_index) + ev = np.asarray(ng.edge_vec) + em = np.asarray(ng.edge_mask) + return { + (int(ei[0, k]), int(ei[1, k]), tuple(np.round(ev[k], 6))) + for k in range(ei.shape[1]) + if em[k] + } + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC and PBC +def test_vesin_pair_excl_none_identity(periodic): + """pair_excl=None: vesin builder output is unchanged (identity).""" + coord, atype, box = _system(periodic) + coord = coord.reshape(1, 4, 3) + box_3d = None if box is None else box.reshape(1, 3, 3) + ng_ref = vesin_builder.build_neighbor_graph_vesin(coord, atype, box_3d, 2.0) + ng_excl = vesin_builder.build_neighbor_graph_vesin( + coord, atype, box_3d, 2.0, pair_excl=None + ) + assert _valid_edge_set(ng_ref) == _valid_edge_set(ng_excl) + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC and PBC +def test_vesin_pair_excl_oracle_set_equality(periodic): + """Vesin builder(pair_excl=X) == dense ref + apply_pair_exclusion(X).""" + coord, atype, box = _system(periodic) + coord = coord.reshape(1, 4, 3) + box_3d = None if box is None else box.reshape(1, 3, 3) + rcut = 2.0 + pe = PairExcludeMask(2, [(0, 1), (1, 0)]) + # dense reference + separate post-process + ng_dense = build_neighbor_graph(coord, atype, box_3d, rcut) + atype_flat = atype.reshape(-1) + ng_ref = apply_pair_exclusion(ng_dense, atype_flat, pe) + # vesin builder with fused post-process + ng_vesin = vesin_builder.build_neighbor_graph_vesin( + coord, atype, box_3d, rcut, pair_excl=pe + ) + assert _valid_edge_set(ng_ref) == _valid_edge_set(ng_vesin) + # exclusion actually removed edges + ng_plain = vesin_builder.build_neighbor_graph_vesin(coord, atype, box_3d, rcut) + assert int(np.asarray(ng_vesin.edge_mask).sum()) < int( + np.asarray(ng_plain.edge_mask).sum() + ) + + +def test_vesin_nlist_edges_pair_excl_raises(): + """VesinNeighborList.build with return_mode='edges' and pair_excl raises NotImplementedError.""" + from deepmd.pt_expt.utils.vesin_neighbor_list import ( + VesinNeighborList, + ) + + coord = torch.zeros((1, 4, 3), dtype=torch.float64) + atype = torch.zeros((1, 4), dtype=torch.int64) + pe = PairExcludeMask(2, [(0, 1)]) + nl = VesinNeighborList() + with pytest.raises(NotImplementedError, match="return_mode='edges'"): + nl.build(coord, atype, None, 2.0, [4], return_mode="edges", pair_excl=pe)