diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 2cf3e4c12b..b3ce544377 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -463,7 +463,8 @@ def _call_common_graph( ng = build_neighbor_graph_ase(cc, atype, bb, self.get_rcut()) else: raise ValueError( - f"unknown neighbor_graph_method {method!r}; use 'dense' or 'ase'" + f"unknown neighbor_graph_method {method!r}; the dpmodel/jax backend " + "supports 'dense'/'ase' only ('vesin'/'nv' require the pt_expt backend)." ) xp = array_api_compat.array_namespace(atype) nf, nloc = atype.shape[:2] diff --git a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py index bc7312fcab..3b00ee6fac 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py @@ -58,7 +58,8 @@ def build_neighbor_graph_ase( coord (nf, nloc, 3) local coordinates. atype - (nf, nloc) local atom types (unused for the search; carried for API parity). + (nf, nloc) local atom types; ``type < 0`` marks a virtual atom, excluded + as center and neighbor (the search itself is type-blind). box (nf, 3, 3) simulation cell, or ``None`` for non-periodic. rcut @@ -127,6 +128,14 @@ def _to_cpu_numpy(x: Any) -> np.ndarray: np.concatenate(nframe_parts) if nframe_parts else np.zeros((0,), dtype=np.int64) ) + # virtual atoms (atype < 0) are excluded as centers AND neighbors -- the + # World-2 builder contract shared with the dense reference builder; the + # geometric search above cannot know about them. + atype_np = _to_cpu_numpy(atype).reshape(nf, nloc) + keep = (atype_np[nframe_all, i_all] >= 0) & (atype_np[nframe_all, j_all] >= 0) + 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( i_all, j_all, S_all, coord, box, nframe_all, nloc, layout=layout ) diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index b7143a9ac6..9109561b32 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -33,9 +33,6 @@ EdgeNeighborList, NeighborList, ) -from deepmd.pt.utils.region import ( - normalize_coord, -) from deepmd.pt_expt.utils.edge_schema import ( edge_schema_from_neighbor_matrix, ) @@ -164,74 +161,24 @@ def build( See :meth:`deepmd.dpmodel.utils.neighbor_list.NeighborList.build`. The returned ``nlist`` is distance-sorted and truncated to ``sum(sel)``. """ - from nvalchemiops.torch.neighbors import ( - neighbor_list, + device = coord.device + nf, nloc = atype.shape[:2] + target_neighbors = int(sum(sel)) + coord = coord.reshape(nf, nloc, 3) + + # Delegate the raw search to the shared helper in nv_graph_builder. + # Function-level import avoids a module-level pt -> pt_expt cycle while + # keeping the search logic in exactly one place (graph-builder primary, + # legacy strategy is the deprecation-bound caller). + from deepmd.pt_expt.utils.nv_graph_builder import ( + nv_search_matrix, ) - device = coord.device - with _input_device_context(device): - nf, nloc = atype.shape[:2] - target_neighbors = int(sum(sel)) - search_capacity = target_neighbors - total_atoms = nf * nloc - coord = coord.reshape(nf, nloc, 3) - periodic = box is not None - if not periodic: - cell = None - pbc = None - else: - cell = box.reshape(nf, 3, 3).to(device=device, dtype=coord.dtype) - coord = normalize_coord(coord, cell) - pbc = torch.ones((nf, 3), dtype=torch.bool, device=device) - positions_for_nlist = coord.reshape(total_atoms, 3).detach() - batch_idx = torch.arange( - nf, dtype=torch.int32, device=device - ).repeat_interleave(nloc) - batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc - method = choose_nv_nlist_method(nloc, periodic=periodic, device=device) - - # ``batch_naive`` otherwise derives ``max_atoms_per_system`` from - # ``batch_ptr`` with a ``.max().item()`` device->host sync on every - # call. Our batches are homogeneous (``nloc`` atoms per frame), so the - # value is known on the host; passing it explicitly removes that - # per-call sync. ``batch_cell_list`` neither accepts the argument nor - # has a ``**kwargs`` catch-all, so the override is guarded on method. - extra_nl_kwargs: dict[str, Any] = {} - if method == "batch_naive": - extra_nl_kwargs["max_atoms_per_system"] = int(nloc) - - # Grow the search capacity until all neighbors fit so the distance-sort - # below selects the true nearest ``sum(sel)``. - while True: - nlist_result = neighbor_list( - positions_for_nlist, - float(rcut), - cell=cell, - pbc=pbc, - batch_idx=batch_idx, - batch_ptr=batch_ptr, - method=method, - max_neighbors=int(search_capacity), - return_neighbor_list=False, - wrap_positions=False, - **extra_nl_kwargs, - ) - if len(nlist_result) == 2: - neighbor_matrix, num_neighbors = nlist_result - shifts = torch.zeros( - (*neighbor_matrix.shape, 3), - dtype=torch.int32, - device=device, - ) - else: - neighbor_matrix, num_neighbors, shifts = nlist_result - max_found = ( - int(num_neighbors.max().item()) if num_neighbors.numel() > 0 else 0 - ) - if max_found <= search_capacity: - break - search_capacity = max(max_found, _grow_search_capacity(search_capacity)) + coord, cell, neighbor_matrix, num_neighbors, shifts = nv_search_matrix( + coord, box, rcut, start_capacity=target_neighbors + ) + with _input_device_context(device): if return_mode == "edges": return edge_schema_from_neighbor_matrix( coord=coord, diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 50c8e6cae6..89a724fc2a 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -65,6 +65,10 @@ if TYPE_CHECKING: import ase.neighborlist + from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + ) + # Public output keys emitted by the graph-form AOTI forward # (``forward_lower_graph_exportable``) keyed by the output-variable category that @@ -125,6 +129,19 @@ class DeepEval(DeepEvalBackend): neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional The ASE neighbor list class to produce the neighbor list. If None, the neighbor list will be built natively in the model. + nlist_backend : str, default: "auto" + Neighbor-list builder for the NLIST/extended lower path (``.pte`` and + nlist-form ``.pt2``): ``"auto"`` / ``"vesin"`` / ``"native"``. Not + used by graph-form ``.pt2`` artifacts. + neighbor_graph_method : str, default: "dense" + Carry-all graph builder for GRAPH-FORM ``.pt2`` artifacts ONLY + (``metadata["lower_input_kind"] == "graph"``): ``"dense"`` / ``"ase"`` + (backend-agnostic) or ``"vesin"`` / ``"nv"`` (on-device O(N)). A + non-default value on any other artifact raises at construction — the + knob would silently do nothing there; use ``nlist_backend`` for the + nlist path instead. All builders emit the same neighbor set, so the + choice is performance-only. Consolidating the two knobs into a single + backend-selection API is deferred to the dense-nlist deprecation. **kwargs : dict Keyword arguments. """ @@ -137,11 +154,15 @@ def __init__( auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, nlist_backend: str = "auto", + neighbor_graph_method: str = "dense", **kwargs: Any, ) -> None: self.output_def = output_def self.model_path = model_file self.neighbor_list = neighbor_list + # World-2 graph-form ``.pt2`` (lower_input_kind == "graph") builder select: + # "dense"/"ase" (backend-agnostic) or "vesin"/"nv" (on-device O(N)). + self._neighbor_graph_method = neighbor_graph_method self._is_pt2 = model_file.endswith(".pt2") if self._is_pt2: @@ -157,6 +178,20 @@ def __init__( "`.pt` (training checkpoint)." ) + # neighbor_graph_method is consumed ONLY by graph-form .pt2 eval + # (_eval_model_graph); fail fast instead of silently ignoring it on + # nlist-form artifacts (there, the builder knob is nlist_backend). + if ( + neighbor_graph_method != "dense" + and getattr(self, "metadata", {}).get("lower_input_kind") != "graph" + ): + raise ValueError( + f"neighbor_graph_method={neighbor_graph_method!r} only applies to " + "graph-form .pt2 artifacts (lower_input_kind == 'graph'); this " + f"model is not graph-form. Use nlist_backend to select the " + "neighbor-list builder for the nlist path." + ) + self._setup_nlist_backend(nlist_backend) if isinstance(auto_batch_size, bool): @@ -1660,9 +1695,6 @@ def _eval_model_graph( forward returns the LOCAL public keys directly, so results are reshaped without ``communicate_extended_output``. """ - from deepmd.dpmodel.utils.neighbor_graph import ( - build_neighbor_graph, - ) from deepmd.pt_expt.utils.env import ( DEVICE, ) @@ -1678,28 +1710,19 @@ def _eval_model_graph( box_input = cells.reshape(nframes, 9) if cells is not None else None # Dynamic edge axis (B2.0): build the carry-all graph at its exact edge # count (no static padding); the AOTI artifact accepts any E. - graph = build_neighbor_graph( - coord_input, - atom_types, - box_input, - self._rcut, - ) + graph = self._build_eval_graph(coord_input, atom_types, box_input, DEVICE) atype_t = torch.tensor( np.asarray(atom_types).reshape(-1), dtype=torch.int64, device=DEVICE ) - n_node_t = torch.tensor( - np.asarray(graph.n_node), dtype=torch.int64, device=DEVICE - ) - edge_index_t = torch.tensor( - np.asarray(graph.edge_index), dtype=torch.int64, device=DEVICE - ) - edge_vec_t = torch.tensor( - np.asarray(graph.edge_vec), dtype=torch.float64, device=DEVICE - ) - edge_mask_t = torch.tensor( - np.asarray(graph.edge_mask), dtype=torch.bool, device=DEVICE + # graph fields may be numpy (dense/ase) or torch, possibly on CUDA + # (vesin/nv) -- torch.as_tensor handles both and moves to DEVICE. + n_node_t = torch.as_tensor(graph.n_node, dtype=torch.int64, device=DEVICE) + edge_index_t = torch.as_tensor( + graph.edge_index, dtype=torch.int64, device=DEVICE ) + edge_vec_t = torch.as_tensor(graph.edge_vec, dtype=torch.float64, device=DEVICE) + edge_mask_t = torch.as_tensor(graph.edge_mask, dtype=torch.bool, device=DEVICE) fparam_t, aparam_t = self._prepare_optional_lower_inputs( fparam, aparam, nframes, natoms, DEVICE @@ -1734,6 +1757,61 @@ def _eval_model_graph( ) return tuple(results) + def _build_eval_graph( + self, + coord_input: np.ndarray, + atom_types: np.ndarray, + box_input: np.ndarray | None, + device: "torch.device", + ) -> "NeighborGraph": + """Build the carry-all NeighborGraph for graph-form ``.pt2`` inference. + + Dispatches on ``self._neighbor_graph_method``: ``dense``/``ase`` run + backend-agnostic (numpy); ``vesin``/``nv`` run on-device (torch, O(N)). + All backends emit the SAME neighbor set (carry-all, sel-free), so the + selection is a pure performance choice and results are unchanged. + """ + method = self._neighbor_graph_method + 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) + 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 + ) + if method in ("vesin", "nv"): + cc = torch.as_tensor(coord_input, dtype=torch.float64, device=device) + aa = torch.as_tensor( + np.asarray(atom_types), dtype=torch.int64, device=device + ) + bb = ( + torch.as_tensor(box_input, dtype=torch.float64, device=device) + if box_input is not None + else None + ) + if method == "vesin": + from deepmd.pt_expt.utils.vesin_graph_builder import ( + build_neighbor_graph_vesin, + ) + + return build_neighbor_graph_vesin(cc, aa, bb, self._rcut) + from deepmd.pt_expt.utils.nv_graph_builder import ( + build_neighbor_graph_nv, + ) + + return build_neighbor_graph_nv(cc, aa, bb, self._rcut) + raise ValueError( + f"unknown neighbor_graph_method {method!r}; " + "use 'dense', 'ase', 'vesin', or 'nv'" + ) + 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 5b19cb63f1..928149ca94 100644 --- a/deepmd/pt_expt/model/make_model.py +++ b/deepmd/pt_expt/model/make_model.py @@ -456,14 +456,38 @@ def _call_common_graph( build_neighbor_graph_ase, ) + # mirror the dpmodel guard: _resolve_graph_method's eligibility + # check only protects the default (None) path; an EXPLICIT + # neighbor_graph_method would otherwise reach the builders for + # descriptors without a graph lower. + descriptor = getattr(self.atomic_model, "descriptor", None) + uses_graph_lower = getattr(descriptor, "uses_graph_lower", lambda: False) + if not (self.mixed_types() and uses_graph_lower()): + raise NotImplementedError( + "neighbor_graph_method requires a mixed_types descriptor with a " + "graph lower (e.g. dpa1 attn_layer=0)" + ) rcut = self.get_rcut() if method == "dense": ng = build_neighbor_graph(cc, atype, bb, rcut) elif method == "ase": ng = build_neighbor_graph_ase(cc, atype, bb, rcut) + 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) + 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) else: raise ValueError( - f"unknown neighbor_graph_method {method!r}; use 'dense' or 'ase'" + f"unknown neighbor_graph_method {method!r}; " + "use 'dense', 'ase', 'vesin', or 'nv'" ) nf, nloc = atype.shape[:2] atype_flat = atype.reshape(nf * nloc) diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py new file mode 100644 index 0000000000..423a315a49 --- /dev/null +++ b/deepmd/pt_expt/utils/nv_graph_builder.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Carry-all NeighborGraph builder backed by nvalchemiops (GPU cell list). + +World-2 counterpart of :mod:`deepmd.pt.utils.nv_nlist`: instead of building the +dense extended quartet, it decodes nvalchemiops' dense +``(total_atoms, max_neighbors)`` neighbor matrix into flat per-frame local +``(i, j, S, nframe_id)`` and delegates to the array-API +:func:`~deepmd.dpmodel.utils.neighbor_graph.neighbor_graph_from_ijs`, which +recomputes ``edge_vec`` differentiably from the (normalized) coordinates. + +Unlike the vesin builder, nvalchemiops batches natively over frames via +``batch_idx``/``batch_ptr`` -- a single GPU kernel handles all ``nf`` frames, +so there is NO per-frame Python loop. CUDA-only ⇒ this module lives in pt_expt. + +The matrix decode mirrors :func:`deepmd.pt.utils.nv_nlist._matrix_to_extended_inputs` +(the authoritative, tested extraction) but stops at the sparse ``(i, j, S)`` +edge list rather than materializing the extended-atom contract. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + GraphLayout, + NeighborGraph, + neighbor_graph_from_ijs, +) +from deepmd.pt.utils.nv_nlist import ( + _input_device_context, + choose_nv_nlist_method, + is_nv_available, +) +from deepmd.pt.utils.region import ( + normalize_coord, +) + + +def _grow_search_capacity(capacity: int) -> int: + """Increase Toolkit-Ops capacity by 1.25x, rounded up (mirror nv_nlist).""" + return (capacity * 5 + 3) // 4 + + +def nv_matrix_to_ijs( + neighbor_matrix: torch.Tensor, + num_neighbors: torch.Tensor, + shifts: torch.Tensor, + nloc: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Decode nvalchemiops' dense neighbor matrix to a sparse edge list. + + Pure torch and device-agnostic (CPU-runnable), so the regression-prone + index arithmetic is unit-testable on the default CI without CUDA — the + GPU ``neighbor_list`` search itself stays behind the opt-in CUDA suite. + Step 1 of :func:`deepmd.pt.utils.nv_nlist._matrix_to_extended_inputs`. + + Parameters + ---------- + neighbor_matrix + (total_atoms, max_neighbors) int; ``neighbor_matrix[dst, slot] = src``, + both flattened batch indices in ``[0, total_atoms)``. Frames are + batch-isolated: a neighbor always shares its center's frame. + num_neighbors + (total_atoms,) int, valid slot count per center. + shifts + (total_atoms, max_neighbors, 3) int periodic image shifts per slot. + nloc + Atoms per frame (``total_atoms = nf * nloc``). + + Returns + ------- + center_local + (E,) int64 per-frame local center index ``i`` (``dst % nloc``). + src_local + (E,) int64 per-frame local neighbor index ``j`` (``src % nloc``). + shift + (E, 3) int64 periodic image shift ``S``. + frame_idx + (E,) int64 frame of each edge (``dst // nloc``). + """ + device = neighbor_matrix.device + total_atoms, max_neighbors = neighbor_matrix.shape + slot = torch.arange(max_neighbors, dtype=torch.long, device=device).expand( + total_atoms, max_neighbors + ) + valid = (slot < num_neighbors.unsqueeze(1)).reshape(-1) + edge_idx = torch.nonzero(valid, as_tuple=False).flatten() + + dst = edge_idx // max_neighbors # flattened center + src = neighbor_matrix.reshape(-1).index_select(0, edge_idx).to(torch.int64) + shift = shifts.reshape(-1, 3).index_select(0, edge_idx).to(torch.int64) + frame_idx = (dst // nloc).to(torch.int64) # frame of the edge + center_local = (dst % nloc).to(torch.int64) # i = center + src_local = (src % nloc).to(torch.int64) # j = neighbor + return center_local, src_local, shift, frame_idx + + +def nv_search_matrix( + coord: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + start_capacity: int, +) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Run the nvalchemiops neighbor search and return the raw matrix output. + + Encapsulates the full search pipeline: ``_input_device_context`` pinning, + periodic coordinate normalization, batch tensor construction, and the + grow-until-fit capacity loop. This is the single authoritative nv search; + :class:`~deepmd.pt.utils.nv_nlist.NvNeighborList` delegates here so the + search logic is maintained in exactly one place. + + Parameters + ---------- + coord : (nf, nloc, 3) local coordinates (already reshaped). + box : (nf, 3, 3) simulation cell, or ``None`` for non-periodic. + rcut : cutoff radius. + start_capacity : initial max-neighbor capacity; grown automatically when + any atom has more neighbors than the current capacity. + + Returns + ------- + coord : (nf, nloc, 3) coordinates, normalized in-cell if periodic. + cell : (nf, 3, 3) float box, or ``None`` for non-periodic. + neighbor_matrix : (total_atoms, capacity) int neighbor matrix. + num_neighbors : (total_atoms,) valid neighbor count per center. + shifts : (total_atoms, capacity, 3) int periodic image shifts. + """ + from nvalchemiops.torch.neighbors import ( + neighbor_list, + ) + + device = coord.device + nf = coord.shape[0] + nloc = coord.shape[1] + periodic = box is not None + + with _input_device_context(device): + if periodic: + cell = box.reshape(nf, 3, 3).to(device=device, dtype=coord.dtype) + coord = normalize_coord(coord, cell) + pbc = torch.ones((nf, 3), dtype=torch.bool, device=device) + else: + cell = None + pbc = None + + total_atoms = nf * nloc + positions = coord.reshape(total_atoms, 3).detach() + batch_idx = torch.arange( + nf, dtype=torch.int32, device=device + ).repeat_interleave(nloc) + batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc + method = choose_nv_nlist_method(nloc, periodic=periodic, device=device) + extra_nl_kwargs: dict[str, Any] = {} + if method == "batch_naive": + extra_nl_kwargs["max_atoms_per_system"] = int(nloc) + + search_capacity = start_capacity + while True: + nlist_result = neighbor_list( + positions, + float(rcut), + cell=cell, + pbc=pbc, + batch_idx=batch_idx, + batch_ptr=batch_ptr, + method=method, + max_neighbors=int(search_capacity), + return_neighbor_list=False, + wrap_positions=False, + **extra_nl_kwargs, + ) + if len(nlist_result) == 2: + neighbor_matrix, num_neighbors = nlist_result + shifts = torch.zeros( + (*neighbor_matrix.shape, 3), dtype=torch.int32, device=device + ) + else: + neighbor_matrix, num_neighbors, shifts = nlist_result + max_found = ( + int(num_neighbors.max().item()) if num_neighbors.numel() > 0 else 0 + ) + if max_found <= search_capacity: + break + search_capacity = max(max_found, _grow_search_capacity(search_capacity)) + + return coord, cell, neighbor_matrix, num_neighbors, shifts + + +def build_neighbor_graph_nv( + coord: Any, + atype: Any, + box: Any | None, + rcut: float, + layout: GraphLayout | None = None, +) -> NeighborGraph: + """Build a CARRY-ALL NeighborGraph using nvalchemiops' GPU cell list. + + Parameters + ---------- + coord + (nf, nloc, 3) or (nf, nloc*3) local coordinates (CUDA tensor). + atype + (nf, nloc) local atom types (carried for API parity). + 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. + + Returns + ------- + graph + The carry-all :class:`NeighborGraph` over the LOCAL atoms, ``edge_vec`` + recomputed differentiably from the (normalized) ``coord``/``box``. + + Raises + ------ + ImportError + if ``nvalchemi-toolkit-ops`` (CUDA) is not installed. + """ + if not is_nv_available(): + raise ImportError( + "build_neighbor_graph_nv requires nvalchemi-toolkit-ops (CUDA); " + "install with `pip install nvalchemi-toolkit-ops` or use " + "neighbor_graph_method='dense'." + ) + + device = coord.device + nf = coord.shape[0] if coord.ndim == 3 else 1 + coord = coord.reshape(nf, -1, 3) + nloc = coord.shape[1] + + if nloc == 0: + empty_i = torch.zeros((0,), dtype=torch.int64, device=device) + empty_S = torch.zeros((0, 3), dtype=torch.int64, device=device) + return neighbor_graph_from_ijs( + empty_i, empty_i, empty_S, coord, box, empty_i, nloc, layout=layout + ) + + # Carry-all: grow capacity until every neighbor fits (no sel cap). + # NOTE: unlike the vesin builder (which searches the ORIGINAL coords -- + # vesin handles unwrapped positions natively), nvalchemiops requires + # in-cell positions, so BOTH the search and the edge_vec recomputation use + # the normalized coords; S then matches the coords the search actually saw. + coord, cell, neighbor_matrix, num_neighbors, shifts = nv_search_matrix( + coord, box, rcut, start_capacity=max(64, nloc) + ) + box_out = cell # edge_vec is recomputed from these (normalized) coords + + # Decode the dense matrix to a sparse (i, j, S) edge list (CPU-testable + # helper; see nv_matrix_to_ijs). + center_local, src_local, shift, frame_idx = nv_matrix_to_ijs( + neighbor_matrix, num_neighbors, shifts, nloc + ) + + # virtual atoms (atype < 0) are excluded as centers AND neighbors — the + # World-2 builder contract shared with the dense reference builder; the + # geometric search above cannot know about them. + at = torch.as_tensor(atype, device=device).reshape(nf, nloc) + keep = (at[frame_idx, center_local] >= 0) & (at[frame_idx, src_local] >= 0) + center_local, src_local = center_local[keep], src_local[keep] + shift, frame_idx = shift[keep], frame_idx[keep] + + return neighbor_graph_from_ijs( + center_local, src_local, shift, coord, box_out, frame_idx, nloc, layout=layout + ) diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py new file mode 100644 index 0000000000..3f86fc7b75 --- /dev/null +++ b/deepmd/pt_expt/utils/vesin_graph_builder.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Carry-all NeighborGraph builder backed by vesin.torch (O(N) cell list). + +World-2 counterpart of vesin_neighbor_list.py: instead of building the dense +quartet, it returns per-frame local (i, j, S), then delegates to the array-API +``neighbor_graph_from_ijs`` (which recomputes ``edge_vec`` differentiably from +the ORIGINAL grad-carrying coords). torch-only => lives in pt_expt. + +Scope note: ``vesin.torch``'s API is single-system, so this builder LOOPS over +frames in Python (~1 ms/frame call overhead measured on GPU). It is intended +for ``nf == 1`` inference and CPU use. It is never on a default hot path: +``neighbor_graph_method=None`` resolves to the ``"dense"`` converter, and +vesin is explicit opt-in only. For batched multi-frame GPU work prefer +``nv`` (:mod:`.nv_graph_builder`), which batches all frames in one kernel. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import array_api_compat +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + GraphLayout, + NeighborGraph, + neighbor_graph_from_ijs, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, +) + + +def vesin_search_ijs( + positions: torch.Tensor, + cell: torch.Tensor | None, + periodic: bool, + rcut: float, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Raw ``vesin.torch`` neighbor search returning ``(ii, jj, ss)`` as int64. + + The caller is responsible for ensuring ``vesin.torch`` is importable (check + :func:`~deepmd.pt_expt.utils.vesin_neighbor_list.is_vesin_torch_available` + before calling). ``positions`` must be detached (the search is + non-differentiable). + + Parameters + ---------- + positions : (nloc, 3) local-frame coordinates, already detached. + cell : (3, 3) box matrix for periodic systems, or ``None`` for non-periodic. + For non-periodic systems a zero box is constructed internally. + periodic : whether the system is periodic. + rcut : neighbor cutoff radius. + device : device to pin as the ambient default. ``vesin.torch`` allocates + some internal tensors on the ambient default device, which may be a + fake/other device in some test contexts (e.g. a placeholder CUDA + default); pinning it here prevents spurious CUDA initializations. + + Returns + ------- + ii : (E,) int64 center local indices. + jj : (E,) int64 neighbor local indices. + ss : (E, 3) int64 periodic image shifts. + """ + import vesin.torch as _vesin_torch + + box = ( + cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) + ) + nl = _vesin_torch.NeighborList(cutoff=float(rcut), full_list=True) + with torch.device(device): + ii, jj, ss = nl.compute( + points=positions, + box=box, + periodic=periodic, + quantities="ijS", + ) + return ii.to(torch.int64), jj.to(torch.int64), ss.to(torch.int64).reshape(-1, 3) + + +def build_neighbor_graph_vesin( + coord: Any, + atype: Any, + box: Any | None, + rcut: float, + layout: GraphLayout | None = None, +) -> 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``. + """ + if not is_vesin_torch_available(): + raise ImportError( + "build_neighbor_graph_vesin requires vesin[torch]; " + "install with `pip install vesin[torch]` or use neighbor_graph_method='dense'." + ) + + xp = array_api_compat.array_namespace(coord) + dev = array_api_compat.device(coord) + nf = coord.shape[0] if coord.ndim == 3 else 1 + coord = xp.reshape(coord, (nf, -1, 3)) + nloc = coord.shape[1] + periodic = box is not None + if periodic: + box = xp.reshape(box, (nf, 3, 3)) + + if nloc == 0: + empty_i = torch.zeros((0,), dtype=torch.int64, device=dev) + empty_S = torch.zeros((0, 3), dtype=torch.int64, device=dev) + empty_nf = torch.zeros((0,), dtype=torch.int64, device=dev) + return neighbor_graph_from_ijs( + empty_i, empty_i, empty_S, coord, box, empty_nf, nloc, layout=layout + ) + + i_parts, j_parts, S_parts, nf_parts = [], [], [], [] + for f in range(nf): + pts = coord[f].detach() + cell_f = box[f].detach() if periodic else None + ii, jj, ss = vesin_search_ijs(pts, cell_f, periodic, rcut, dev) + i_parts.append(ii) + j_parts.append(jj) + S_parts.append(ss) + nf_parts.append(torch.full((ii.shape[0],), f, dtype=torch.int64, device=dev)) + + # guard torch.cat against empty part lists (nf == 0), mirroring ase_builder + i_all = ( + torch.cat(i_parts) + if i_parts + else torch.zeros((0,), dtype=torch.int64, device=dev) + ) + j_all = ( + torch.cat(j_parts) + if j_parts + else torch.zeros((0,), dtype=torch.int64, device=dev) + ) + S_all = ( + torch.cat(S_parts) + if S_parts + else torch.zeros((0, 3), dtype=torch.int64, device=dev) + ) + nf_all = ( + torch.cat(nf_parts) + if nf_parts + else torch.zeros((0,), dtype=torch.int64, device=dev) + ) + + # virtual atoms (atype < 0) are excluded as centers AND neighbors — the + # World-2 builder contract shared with the dense reference builder; the + # geometric search above cannot know about them. + at = torch.as_tensor(atype, device=dev).reshape(nf, nloc) + keep = (at[nf_all, i_all] >= 0) & (at[nf_all, j_all] >= 0) + i_all, j_all, nf_all = i_all[keep], j_all[keep], nf_all[keep] + S_all = S_all[keep] + + # i = center (dst), j = neighbor (src); pass ORIGINAL coord/box + # (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( + i_all, j_all, S_all, coord, box, nf_all, nloc, layout=layout + ) diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py index 6b1a165b98..39a0dff883 100644 --- a/deepmd/pt_expt/utils/vesin_neighbor_list.py +++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py @@ -169,8 +169,6 @@ def _build_single( non-differentiable); the returned ``extended_coord`` is rebuilt from ``positions`` so gradients flow to the local atoms and box. """ - import vesin.torch - device = positions.device nsel = sum(sel) nloc = positions.shape[0] @@ -191,20 +189,18 @@ def _build_single( cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) ) - # Pin the default device to the input's device: vesin.torch allocates some - # internal tensors on the ambient default device, which may be a fake/other - # device in some contexts (e.g. tests set a placeholder CUDA default). The - # search runs on detached inputs -- it is non-differentiable. - nl = vesin.torch.NeighborList(cutoff=rcut, full_list=True) - with torch.device(device): - ii, jj, ss = nl.compute( - points=positions.detach(), - box=box.detach(), - periodic=periodic, - quantities="ijS", - ) - ii = ii.to(torch.int64) - jj = jj.to(torch.int64) + # Delegate the raw search to the shared helper in vesin_graph_builder + # (function-level import: legacy module depends on graph module lazily to + # avoid a module-level cycle — vesin_graph_builder imports + # is_vesin_torch_available from this module). + from deepmd.pt_expt.utils.vesin_graph_builder import ( + vesin_search_ijs, + ) + + ii, jj, ss = vesin_search_ijs( + positions.detach(), cell if periodic else None, periodic, rcut, device + ) + # ss is int64 from the helper; cast to float here for later ``ss @ box`` math. ss = ss.to(positions.dtype) # ghost atoms: neighbors reached through a non-zero periodic shift. Rebuild @@ -271,8 +267,6 @@ def _build_single_edges( sel: list[int], ) -> EdgeNeighborList: """Single-frame ``vesin`` output converted directly to edge vectors.""" - import vesin.torch - device = positions.device nsel = sum(sel) nloc = positions.shape[0] @@ -288,21 +282,18 @@ def _build_single_edges( ) periodic = cell is not None - box = ( - cell if periodic else torch.zeros((3, 3), dtype=positions.dtype, device=device) + from deepmd.pt_expt.utils.vesin_graph_builder import ( + vesin_search_ijs, ) - nl = vesin.torch.NeighborList(cutoff=rcut, full_list=True) - with torch.device(device): - ii, jj, ss = nl.compute( - points=positions.detach(), - box=box.detach(), - periodic=periodic, - quantities="ijS", - ) + + ii, jj, ss = vesin_search_ijs( + positions.detach(), cell if periodic else None, periodic, rcut, device + ) + # ss is int64 from the helper; edge_schema_from_ij_shifts accepts int shifts. return edge_schema_from_ij_shifts( positions=positions, atype=atype, - cell=box if periodic else None, + cell=cell, ii=ii, jj=jj, shifts=ss, diff --git a/source/tests/common/dpmodel/test_from_ijs.py b/source/tests/common/dpmodel/test_from_ijs.py index bab616e452..dc0af94037 100644 --- a/source/tests/common/dpmodel/test_from_ijs.py +++ b/source/tests/common/dpmodel/test_from_ijs.py @@ -88,3 +88,28 @@ def test_ase_matches_intree_carry_all_nonperiodic(self) -> None: if __name__ == "__main__": unittest.main() + + def test_ase_excludes_virtual_atoms_like_dense(self) -> None: + """Virtual atoms (atype < 0) excluded as center AND neighbor. + + The dense reference builder filters virtual atoms during construction; + the geometric ASE search is type-blind, so the builder must post-filter + to keep the World-2 "same neighbor set" contract (OutisLi, #5714). + """ + pytest.importorskip("ase") + from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, + build_neighbor_graph_ase, + ) + + rng = np.random.default_rng(5) + coord = rng.normal(size=(1, 8, 3)) * 2.0 + atype = np.array([[0, 1, -1, 0, 1, -1, 0, 1]], dtype=np.int64) + box = np.eye(3)[None] * 8.0 + ng_ase = build_neighbor_graph_ase(coord, atype, box, rcut=4.0) + ng_ref = build_neighbor_graph(coord, atype, box, rcut=4.0) + self.assertEqual(self._sets(ng_ase, 8), self._sets(ng_ref, 8)) + # no real edge touches a virtual atom + ei = ng_ase.edge_index[:, ng_ase.edge_mask] + flat_atype = atype.reshape(-1) + assert np.all(flat_atype[ei[0]] >= 0) and np.all(flat_atype[ei[1]] >= 0) diff --git a/source/tests/pt_expt/infer/test_graph_deepeval.py b/source/tests/pt_expt/infer/test_graph_deepeval.py index 7fc83b677f..4e7b929391 100644 --- a/source/tests/pt_expt/infer/test_graph_deepeval.py +++ b/source/tests/pt_expt/infer/test_graph_deepeval.py @@ -32,6 +32,9 @@ from deepmd.pt_expt.utils.serialization import ( deserialize_to_file, ) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, +) # dpa1 with attn_layer == 0 -- the energy model exercised by the graph path. DPA1_CONFIG = { @@ -237,3 +240,27 @@ def test_graph_pt2_deepeval_parity(graph_pt2, pbc, system) -> None: atol=1e-10, err_msg="atom_virial", ) + + +@pytest.mark.skipif(not is_vesin_torch_available(), reason="vesin[torch] not installed") +@pytest.mark.parametrize("pbc", [True, False]) # periodic vs non-periodic +def test_graph_pt2_deepeval_vesin_matches_dense(graph_pt2, pbc) -> None: + """Selecting neighbor_graph_method='vesin' at DeepEval yields identical + energy/force/virial to the default 'dense' builder on the SAME graph ``.pt2`` + (the builder is a pure perf choice; neighbor sets are equal). + """ + pt2_path, _ = graph_pt2 + coords, cells, atype = _build_system(**_SYSTEMS["small_8"]) + box = cells if pbc else None + max_nn = _max_neighbors(coords, box, atype) + assert max_nn < SEL, "test system must be non-binding for carry-all parity" + + dp_dense = DeepPot(pt2_path) # default neighbor_graph_method == "dense" + dp_vesin = DeepPot(pt2_path, neighbor_graph_method="vesin") + assert dp_vesin.deep_eval._neighbor_graph_method == "vesin" + + e_d, f_d, v_d = dp_dense.eval(coords, box, atype) + e_v, f_v, v_v = dp_vesin.eval(coords, box, atype) + np.testing.assert_allclose(e_v, e_d, rtol=1e-10, atol=1e-10, err_msg="energy") + np.testing.assert_allclose(f_v, f_d, rtol=1e-10, atol=1e-10, err_msg="force") + np.testing.assert_allclose(v_v, v_d, rtol=1e-10, atol=1e-10, err_msg="virial") diff --git a/source/tests/pt_expt/model/test_graph_builder_dispatch.py b/source/tests/pt_expt/model/test_graph_builder_dispatch.py new file mode 100644 index 0000000000..6ec259c177 --- /dev/null +++ b/source/tests/pt_expt/model/test_graph_builder_dispatch.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""neighbor_graph_method dispatch: 'vesin' routes to the vesin builder and is a +perf-only equivalent of 'dense' (same energy + force); dpmodel/jax fail-fast. +""" + +import numpy as np +import pytest +import torch + +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.nv_nlist import ( + is_nv_available, +) +from deepmd.pt_expt.descriptor.dpa1 import ( + DescrptDPA1, +) +from deepmd.pt_expt.fitting.invar_fitting import ( + InvarFitting, +) +from deepmd.pt_expt.model.ener_model import ( + EnergyModel, +) +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, +) + +GLOBAL_SEED = 20240101 + + +def _make_model(): + rcut, rcut_smth, sel, nt = 6.0, 2.0, 20, 2 + ds = DescrptDPA1( + rcut, + rcut_smth, + sel, + nt, + neuron=[3, 6], + axis_neuron=2, + attn=4, + attn_layer=0, # graph lower only supports attn_layer == 0 + attn_dotr=True, + attn_mask=False, + activation_function="tanh", + set_davg_zero=False, + type_one_side=True, + precision="float64", + seed=GLOBAL_SEED, + ).to(env.DEVICE) + ft = InvarFitting( + "energy", + nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + precision="float64", + seed=GLOBAL_SEED, + ).to(env.DEVICE) + return EnergyModel(ds, ft, type_map=["O", "H"]).to(env.DEVICE) + + +def _eval(model, method): + rng = np.random.default_rng(0) + coord = torch.tensor( + rng.random((1, 6, 3)) * 4.0, dtype=torch.float64, device=env.DEVICE + ) + atype = torch.tensor([[0, 1, 1, 0, 1, 1]], dtype=torch.int64, device=env.DEVICE) + box = (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 6.0).reshape(1, 3, 3) + ret = model.forward_common(coord, atype, box, neighbor_graph_method=method) + # graph path returns the output-agnostic dict (no translated force/virial); + # energy_redu = total energy, energy_derv_r = d energy / d coord (force parity) + return ret["energy_redu"], ret["energy_derv_r"] + + +@pytest.mark.skipif(not is_vesin_torch_available(), reason="vesin[torch] not installed") +def test_vesin_matches_dense_energy_force(): + torch.manual_seed(0) + model = _make_model() + e_d, f_d = _eval(model, "dense") + e_v, f_v = _eval(model, "vesin") + tol = 1e-12 if env.DEVICE.type == "cpu" else 1e-10 + torch.testing.assert_close(e_v, e_d, rtol=tol, atol=tol) + torch.testing.assert_close(f_v, f_d, rtol=tol, atol=tol) + + +@pytest.mark.skipif( + not (torch.cuda.is_available() and is_nv_available()), + reason="nvalchemiops requires CUDA + nvalchemi-toolkit-ops", +) +def test_nv_matches_dense_energy_force(): + torch.manual_seed(0) + model = _make_model() + e_d, f_d = _eval(model, "dense") + e_n, f_n = _eval(model, "nv") + tol = 1e-10 # CUDA fp64: absorbs scatter-atomic / index_add nondeterminism + torch.testing.assert_close(e_n, e_d, rtol=tol, atol=tol) + torch.testing.assert_close(f_n, f_d, rtol=tol, atol=tol) + + +def test_dpmodel_backend_rejects_vesin(): + """dpmodel/jax fail-fast names vesin/nv as pt_expt-only.""" + from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DPDescrptDPA1 + from deepmd.dpmodel.fitting.invar_fitting import InvarFitting as DPInvarFitting + from deepmd.dpmodel.model.ener_model import EnergyModel as DPEnergyModel + + rcut, rcut_smth, sel, nt = 6.0, 2.0, 20, 2 + ds = DPDescrptDPA1( + rcut, + rcut_smth, + sel, + nt, + neuron=[3, 6], + axis_neuron=2, + attn=4, + attn_layer=0, + attn_dotr=True, + attn_mask=False, + activation_function="tanh", + set_davg_zero=False, + type_one_side=True, + precision="float64", + seed=GLOBAL_SEED, + ) + ft = DPInvarFitting( + "energy", + nt, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + precision="float64", + seed=GLOBAL_SEED, + ) + model = DPEnergyModel(ds, ft, type_map=["O", "H"]) + coord = np.random.default_rng(0).random((1, 6, 3)) * 4.0 + atype = np.array([[0, 1, 1, 0, 1, 1]], dtype=np.int64) + box = (np.eye(3) * 6.0).reshape(1, 3, 3) + with pytest.raises(ValueError, match="pt_expt backend"): + model.call_common(coord, atype, box, neighbor_graph_method="vesin") + with pytest.raises(ValueError, match="pt_expt backend"): + model.call_common(coord, atype, box, neighbor_graph_method="nv") + + +def test_explicit_method_fails_fast_for_ineligible_descriptor(): + """An EXPLICIT neighbor_graph_method must fail fast when the descriptor + has no graph lower (mirrors the dpmodel guard; the default-path check in + _resolve_graph_method does not protect explicit methods). Regression for + OutisLi review on #5714. + """ + from deepmd.pt_expt.descriptor.se_e2_a import ( + DescrptSeA, + ) + + # se_e2_a: mixed_types() is False and there is no graph lower + ds = DescrptSeA( + 6.0, + 2.0, + [10, 10], + neuron=[3, 6], + axis_neuron=2, + precision="float64", + seed=GLOBAL_SEED, + ).to(env.DEVICE) + ft = InvarFitting( + "energy", + 2, + ds.get_dim_out(), + 1, + mixed_types=ds.mixed_types(), + precision="float64", + seed=GLOBAL_SEED, + ).to(env.DEVICE) + model = EnergyModel(ds, ft, type_map=["O", "H"]).to(env.DEVICE) + coord = torch.rand(1, 4, 3, dtype=torch.float64, device=env.DEVICE) * 3 + atype = torch.zeros(1, 4, dtype=torch.int64, device=env.DEVICE) + box = (torch.eye(3, dtype=torch.float64, device=env.DEVICE) * 6).reshape(1, 9) + for method in ("dense", "ase", "vesin", "nv"): + with pytest.raises(NotImplementedError, match="graph lower"): + model.call_common(coord, atype, box, neighbor_graph_method=method) diff --git a/source/tests/pt_expt/test_plugin.py b/source/tests/pt_expt/test_plugin.py index a59242e592..ddee7a6876 100644 --- a/source/tests/pt_expt/test_plugin.py +++ b/source/tests/pt_expt/test_plugin.py @@ -22,12 +22,36 @@ def fake_entry_points(*, group=None): return [_FakeEntryPoint(calls)] monkeypatch.setattr(importlib.metadata, "entry_points", fake_entry_points) + + # Snapshot the deepmd.pt_expt module tree BEFORE re-importing. Just popping + # "deepmd.pt_expt" and leaving its submodules cached poisons sys.modules + # for the rest of the pytest process: a later import of a cached submodule + # (e.g. deepmd.pt_expt.infer.deep_eval) re-creates a BARE parent package + # whose submodule attributes (utils/infer/...) are never rebound, and + # mock.patch("deepmd.pt_expt.utils...") then fails with AttributeError on + # py3.10 (shard-order dependent CI failure). + saved = { + k: v + for k, v in sys.modules.items() + if k == "deepmd.pt_expt" or k.startswith("deepmd.pt_expt.") + } + deepmd_pkg = sys.modules.get("deepmd") sys.modules.pop("deepmd.pt_expt", None) try: importlib.import_module("deepmd.pt_expt") finally: - sys.modules.pop("deepmd.pt_expt", None) + # drop everything the fresh import created, then restore the snapshot + # (including the parent-package attribute binding). + for k in [ + m + for m in list(sys.modules) + if m == "deepmd.pt_expt" or m.startswith("deepmd.pt_expt.") + ]: + sys.modules.pop(k, None) + sys.modules.update(saved) + if deepmd_pkg is not None and "deepmd.pt_expt" in saved: + deepmd_pkg.pt_expt = saved["deepmd.pt_expt"] assert groups == ["deepmd.pt_expt"] assert calls == ["load"] diff --git a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py index 17aef4d671..a541f744cc 100644 --- a/source/tests/pt_expt/utils/test_graph_pt2_metadata.py +++ b/source/tests/pt_expt/utils/test_graph_pt2_metadata.py @@ -107,3 +107,28 @@ def test_dense_pt2_has_lower_input_kind_nlist(dpa1_dpmodel_data) -> None: assert meta["lower_input_kind"] == "nlist" # edge_capacity is a graph-only artifact constant; the dense path omits it. assert "edge_capacity" not in meta + + +def test_neighbor_graph_method_rejected_on_nlist_artifact(dpa1_dpmodel_data) -> None: + """A non-default ``neighbor_graph_method`` on a NLIST-form artifact raises. + + The knob is consumed only by graph-form ``.pt2`` eval; silently ignoring + it on nlist-form artifacts misled users into thinking they selected an + O(N) builder (OutisLi review, #5714). The nlist-path knob is + ``nlist_backend``. + """ + from deepmd.infer import ( + DeepPot, + ) + + with tempfile.TemporaryDirectory() as d: + p = os.path.join(d, "m_dense.pt2") + deserialize_to_file( + p, + copy.deepcopy(dpa1_dpmodel_data), + do_atomic_virial=True, + ) + with pytest.raises(ValueError, match="graph-form"): + DeepPot(p, neighbor_graph_method="vesin") + # the default stays accepted (no behavior change) + DeepPot(p) diff --git a/source/tests/pt_expt/utils/test_nv_graph_builder.py b/source/tests/pt_expt/utils/test_nv_graph_builder.py new file mode 100644 index 0000000000..3da0da5931 --- /dev/null +++ b/source/tests/pt_expt/utils/test_nv_graph_builder.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""nvalchemiops carry-all NeighborGraph builder: neighbor SET must equal the +in-tree ``dense`` carry-all reference. CUDA + nvalchemi-toolkit-ops only. +""" + +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, +) + +nv_builder = pytest.importorskip("deepmd.pt_expt.utils.nv_graph_builder") +from deepmd.pt.utils.nv_nlist import ( + is_nv_available, +) + +pytestmark = pytest.mark.skipif( + not (torch.cuda.is_available() and is_nv_available()), + reason="nvalchemiops requires CUDA + nvalchemi-toolkit-ops", +) + + +def _sets(ng, nloc): + """Per-center set of (src_local, rounded edge_vec) over real edges.""" + ei = np.asarray(ng.edge_index.cpu()) + ev = np.asarray(ng.edge_vec.detach().cpu()) + em = np.asarray(ng.edge_mask.cpu()) + out = {c: set() for c in range(nloc)} + for e in range(ei.shape[1]): + if em[e]: + out[int(ei[1, e])].add((int(ei[0, e]), tuple(np.round(ev[e], 6)))) + return out + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC and PBC +def test_nv_matches_intree_carry_all(periodic): + dev = torch.device("cuda") + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [0.0, 1.1, 0.0], [1.8, 1.8, 0.0]]], + dtype=torch.float64, + device=dev, + ) + box = ( + (torch.eye(3, dtype=torch.float64, device=dev) * 3.0).reshape(1, 3, 3) + if periodic + else None + ) + atype = torch.tensor([[0, 1, 0, 1]], dtype=torch.int64, device=dev) + ng_ref = build_neighbor_graph(coord, atype, box, 2.0) + ng = nv_builder.build_neighbor_graph_nv(coord, atype, box, 2.0) + assert _sets(ng, 4) == _sets(ng_ref, 4) + + +def test_nv_batches_frames_without_python_loop(): + """Multi-frame: nv searches all frames in one kernel (no per-frame loop).""" + dev = torch.device("cuda") + rng = np.random.default_rng(0) + coord = torch.tensor(rng.random((3, 5, 3)) * 3.0, dtype=torch.float64, device=dev) + box = ( + (torch.eye(3, dtype=torch.float64, device=dev) * 4.0) + .reshape(1, 3, 3) + .repeat(3, 1, 1) + ) + atype = torch.tensor( + [[0, 1, 0, 1, 0], [1, 0, 1, 0, 1], [0, 0, 1, 1, 0]], + dtype=torch.int64, + device=dev, + ) + ng_ref = build_neighbor_graph(coord, atype, box, 2.0) + ng = nv_builder.build_neighbor_graph_nv(coord, atype, box, 2.0) + # per-frame node offset: frame f centers occupy nodes [f*5, (f+1)*5) + for f in range(3): + s_ref = { + ( + int(ng_ref.edge_index[0, e]), + tuple(np.round(np.asarray(ng_ref.edge_vec[e].detach().cpu()), 6)), + ) + for e in range(ng_ref.edge_index.shape[1]) + if bool(ng_ref.edge_mask[e]) + and f * 5 <= int(ng_ref.edge_index[1, e]) < (f + 1) * 5 + } + s = { + ( + int(ng.edge_index[0, e]), + tuple(np.round(np.asarray(ng.edge_vec[e].detach().cpu()), 6)), + ) + for e in range(ng.edge_index.shape[1]) + if bool(ng.edge_mask[e]) and f * 5 <= int(ng.edge_index[1, e]) < (f + 1) * 5 + } + assert s == s_ref, f"frame {f} neighbor set mismatch" + + +def test_nv_edge_vec_is_differentiable(): + dev = torch.device("cuda") + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [0.0, 1.1, 0.0], [1.8, 1.8, 0.0]]], + dtype=torch.float64, + device=dev, + ).requires_grad_(True) + box = (torch.eye(3, dtype=torch.float64, device=dev) * 3.0).reshape(1, 3, 3) + atype = torch.tensor([[0, 1, 0, 1]], dtype=torch.int64, device=dev) + ng = nv_builder.build_neighbor_graph_nv(coord, atype, box, 2.0) + (ng.edge_vec**2).sum().backward() + assert coord.grad is not None and torch.any(coord.grad != 0) + + +def test_nv_excludes_virtual_atoms_like_dense(): + """Virtual atoms (atype < 0) excluded as center AND neighbor (dense contract).""" + dev = torch.device("cuda") + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [0.0, 1.1, 0.0], [1.8, 1.8, 0.0]]], + dtype=torch.float64, + device=dev, + ) + box = (torch.eye(3, dtype=torch.float64, device=dev) * 3.0).reshape(1, 3, 3) + atype = torch.tensor([[0, -1, 0, 1]], dtype=torch.int64, device=dev) # 1 virtual + ng_ref = build_neighbor_graph(coord, atype, box, 2.0) + ng = nv_builder.build_neighbor_graph_nv(coord, atype, box, 2.0) + assert _sets(ng, 4) == _sets(ng_ref, 4) + ei = np.asarray(ng.edge_index.cpu())[:, np.asarray(ng.edge_mask.cpu())] + at = atype.reshape(-1).cpu().numpy() + assert np.all(at[ei[0]] >= 0) and np.all(at[ei[1]] >= 0) diff --git a/source/tests/pt_expt/utils/test_nv_matrix_decode.py b/source/tests/pt_expt/utils/test_nv_matrix_decode.py new file mode 100644 index 0000000000..5eccfa8277 --- /dev/null +++ b/source/tests/pt_expt/utils/test_nv_matrix_decode.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CPU unit tests for the nv dense-matrix -> (i, j, S) decode. + +The GPU ``neighbor_list`` search in ``build_neighbor_graph_nv`` is CUDA-only +and stays behind the opt-in CUDA suite (test_nv_graph_builder.py); the decode +(``nv_matrix_to_ijs``) is pure torch index arithmetic, so its regression-prone +parts (``// max_neighbors``, ``% nloc``, frame isolation, slot-validity mask) +are pinned here on the default CI with synthetic inputs. +""" + +import numpy as np +import torch + +from deepmd.pt_expt.utils.nv_graph_builder import ( + nv_matrix_to_ijs, +) + + +def _edge_set(i, j, s, f): + return { + (int(f[e]), int(i[e]), int(j[e]), tuple(int(x) for x in s[e])) + for e in range(i.shape[0]) + } + + +class TestNvMatrixDecode: + def test_two_frames_hand_checked(self) -> None: + """nf=2, nloc=3, max_neighbors=2; edges and shifts checked by hand.""" + nloc = 3 + # flattened centers 0..5; frame 0 = atoms 0-2, frame 1 = atoms 3-5. + # matrix[dst, slot] = src (flattened); only the first num_neighbors + # slots are valid, the rest is stale garbage that MUST be ignored. + neighbor_matrix = torch.tensor( + [ + [1, 2], # center 0: neighbors 1, 2 + [0, 9], # center 1: neighbor 0 (slot 1 = garbage) + [0, 9], # center 2: neighbor 0 (slot 1 = garbage) + [4, 9], # center 3 (frame 1, local 0): neighbor 4 (local 1) + [3, 9], # center 4 (frame 1, local 1): neighbor 3 (local 0) + [9, 9], # center 5: no neighbors (all garbage) + ], + dtype=torch.int32, + ) + num_neighbors = torch.tensor([2, 1, 1, 1, 1, 0], dtype=torch.int32) + shifts = torch.zeros((6, 2, 3), dtype=torch.int32) + shifts[0, 1] = torch.tensor([1, 0, -1], dtype=torch.int32) # edge 0->2 + + i, j, s, f = nv_matrix_to_ijs(neighbor_matrix, num_neighbors, shifts, nloc) + + assert i.dtype == j.dtype == s.dtype == f.dtype == torch.int64 + assert _edge_set(i, j, s, f) == { + (0, 0, 1, (0, 0, 0)), + (0, 0, 2, (1, 0, -1)), + (0, 1, 0, (0, 0, 0)), + (0, 2, 0, (0, 0, 0)), + (1, 0, 1, (0, 0, 0)), # frame 1: local indices via % nloc + (1, 1, 0, (0, 0, 0)), + } + + def test_empty_no_neighbors(self) -> None: + """All-zero num_neighbors yields zero edges (no garbage leaks).""" + neighbor_matrix = torch.full((4, 3), 7, dtype=torch.int32) + num_neighbors = torch.zeros((4,), dtype=torch.int32) + shifts = torch.zeros((4, 3, 3), dtype=torch.int32) + i, j, s, f = nv_matrix_to_ijs(neighbor_matrix, num_neighbors, shifts, 2) + assert i.shape == (0,) and j.shape == (0,) + assert s.shape == (0, 3) and f.shape == (0,) + + def test_random_vs_oracle(self) -> None: + """Random matrices match a brute-force python oracle.""" + rng = np.random.default_rng(11) + nf, nloc, mn = 3, 4, 5 + total = nf * nloc + num = rng.integers(0, mn + 1, size=total) + mat = np.zeros((total, mn), dtype=np.int64) + shf = rng.integers(-2, 3, size=(total, mn, 3)) + oracle = set() + for dst in range(total): + frame = dst // nloc + for slot in range(mn): + # batch isolation: valid srcs share the center's frame + src = int(rng.integers(frame * nloc, (frame + 1) * nloc)) + mat[dst, slot] = src + if slot < num[dst]: + oracle.add( + ( + frame, + dst % nloc, + src % nloc, + tuple(int(x) for x in shf[dst, slot]), + ) + ) + i, j, s, f = nv_matrix_to_ijs( + torch.from_numpy(mat).to(torch.int32), + torch.from_numpy(num).to(torch.int32), + torch.from_numpy(shf).to(torch.int32), + nloc, + ) + assert _edge_set(i, j, s, f) == oracle diff --git a/source/tests/pt_expt/utils/test_vesin_graph_builder.py b/source/tests/pt_expt/utils/test_vesin_graph_builder.py new file mode 100644 index 0000000000..dea25adab3 --- /dev/null +++ b/source/tests/pt_expt/utils/test_vesin_graph_builder.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import pytest +import torch + +from deepmd.dpmodel.utils.neighbor_graph import ( + build_neighbor_graph, +) + +vesin_builder = pytest.importorskip("deepmd.pt_expt.utils.vesin_graph_builder") +from deepmd.pt_expt.utils.vesin_neighbor_list import ( + is_vesin_torch_available, +) + +pytestmark = pytest.mark.skipif( + not is_vesin_torch_available(), reason="vesin[torch] not installed" +) + + +def _sets(ng, nloc): + """Per-center set of (src_local, rounded edge_vec) over real edges.""" + ei = np.asarray(ng.edge_index) + ev = np.asarray(ng.edge_vec) + em = np.asarray(ng.edge_mask) + out = {c: set() for c in range(nloc)} + for e in range(ei.shape[1]): + if not em[e]: + continue + src, dst = int(ei[0, e]), int(ei[1, e]) + out[dst].add((src, tuple(np.round(ev[e], 6)))) + return out + + +def _system(periodic): + coord = torch.tensor( + [[0.0, 0.0, 0.0], [0.9, 0.0, 0.0], [0.0, 1.1, 0.0], [1.8, 1.8, 0.0]], + dtype=torch.float64, + ) + box = torch.eye(3, dtype=torch.float64) * 3.0 if periodic else None + # atype is (nf, nloc) = (1, 4); build_neighbor_graph requires 2-D atype + atype = torch.tensor([[0, 1, 0, 1]], dtype=torch.int64) + return coord, atype, box + + +@pytest.mark.parametrize("periodic", [False, True]) # non-PBC and PBC +def test_vesin_matches_intree_carry_all(periodic): + coord, atype, box = _system(periodic) + rcut = 2.0 + ng_ref = build_neighbor_graph( + coord.reshape(1, 4, 3), + atype, + None if box is None else box.reshape(1, 3, 3), + rcut, + ) + ng = vesin_builder.build_neighbor_graph_vesin( + coord.reshape(1, 4, 3), + atype, + None if box is None else box.reshape(1, 3, 3), + rcut, + ) + assert _sets(ng, 4) == _sets(ng_ref, 4) + + +def test_vesin_outputs_on_input_device(): + coord, atype, box = _system(True) + ng = vesin_builder.build_neighbor_graph_vesin( + coord.reshape(1, 4, 3), atype, box.reshape(1, 3, 3), 2.0 + ) + assert ng.edge_index.device.type == coord.device.type + assert ng.edge_vec.device.type == coord.device.type + + +def test_vesin_empty_system(): + coord = torch.zeros((1, 0, 3), dtype=torch.float64) + atype = torch.zeros((0,), dtype=torch.int64) + ng = vesin_builder.build_neighbor_graph_vesin(coord, atype, None, 2.0) + assert bool(ng.edge_mask.any()) is False # only min_edges guard edges + + +def test_vesin_edge_vec_is_differentiable(): + coord, atype, box = _system(True) + coord = coord.reshape(1, 4, 3).requires_grad_(True) + ng = vesin_builder.build_neighbor_graph_vesin( + coord, atype, box.reshape(1, 3, 3), 2.0 + ) + # Use squared sum: with full_list=True every edge (i,j,S) has a reverse (j,i,-S) + # so edge_vec.sum() = 0 and its gradient is identically zero. The squared + # loss is asymmetric and gives a non-trivial, non-cancelling gradient. + (ng.edge_vec**2).sum().backward() + assert coord.grad is not None and torch.any(coord.grad != 0) + + +def test_vesin_excludes_virtual_atoms_like_dense(): + """Virtual atoms (atype < 0) excluded as center AND neighbor (dense contract).""" + coord, _, box = _system(periodic=True) + atype = torch.tensor([[0, -1, 0, 1]], dtype=torch.int64) # atom 1 virtual + rcut = 2.0 + ng_ref = build_neighbor_graph( + coord.reshape(1, 4, 3), atype, box.reshape(1, 3, 3), rcut + ) + ng = vesin_builder.build_neighbor_graph_vesin( + coord.reshape(1, 4, 3), atype, box.reshape(1, 3, 3), rcut + ) + assert _sets(ng, 4) == _sets(ng_ref, 4) + 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)