diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index a9af7a50e5..0179543dd4 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -18,6 +18,17 @@ is_lmdb, make_neighbor_stat_data, ) +from .neighbor_graph import ( + GraphLayout, + NeighborGraph, + build_neighbor_graph, + edge_force_virial, + from_dense_quartet, + node_validity_mask, + pad_and_guard_edges, + segment_mean, + segment_sum, +) from .neighbor_list import ( NeighborList, ) @@ -64,20 +75,25 @@ "EmbeddingNet", "EnvMat", "FittingNet", + "GraphLayout", "LmdbDataReader", "LmdbTestData", "LmdbTestDataNlocView", "NativeLayer", "NativeNet", + "NeighborGraph", "NeighborList", "NetworkCollection", "PairExcludeMask", "SameNlocBatchSampler", "aggregate", "build_multiple_neighbor_list", + "build_neighbor_graph", "build_neighbor_list", "compute_total_numb_batch", + "edge_force_virial", "extend_coord_with_ghosts", + "from_dense_quartet", "get_graph_index", "get_multiple_nlist_key", "inter2phys", @@ -88,11 +104,15 @@ "make_multilayer_network", "make_neighbor_stat_data", "nlist_distinguish_types", + "node_validity_mask", "normalize_coord", + "pad_and_guard_edges", "phys2inter", "resolve_model_prob", "resolve_model_prob_from_epochs", "save_dp_model", + "segment_mean", + "segment_sum", "to_face_distance", "traverse_model_dict", ] diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py new file mode 100644 index 0000000000..08b165f861 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""NeighborGraph: backend-agnostic edge-graph neighbor-list subsystem. + +The unified edge/graph neighbor-list contract and its supporting machinery: +``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity ++ edge padding), ``builder`` (the carry-all ``build_neighbor_graph`` dispatcher + +the ``from_dense_quartet`` legacy converter), ``segment`` (mask-aware +segment-reduction toolkit), and ``derivatives`` (edge force/virial assembly). +See memory/spec_unified_edge_nlist.md. +""" + +from .builder import ( + build_neighbor_graph, + from_dense_quartet, +) +from .derivatives import ( + edge_force_virial, +) +from .graph import ( + GraphLayout, + NeighborGraph, + node_validity_mask, + pad_and_guard_edges, +) +from .segment import ( + segment_mean, + segment_sum, +) + +__all__ = [ + "GraphLayout", + "NeighborGraph", + "build_neighbor_graph", + "edge_force_virial", + "from_dense_quartet", + "node_validity_mask", + "pad_and_guard_edges", + "segment_mean", + "segment_sum", +] diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py new file mode 100644 index 0000000000..9a10d3f805 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Builders/converters that produce a :class:`NeighborGraph`. + +Two distinct groups (see memory/spec_unified_edge_nlist.md decision #17), kept +separate so a consumer can never assume completeness while a function silently +truncated: + +1. **Dispatcher (compute from raw geometry).** ``build_neighbor_graph`` takes + coordinates/box/types -- *no pre-existing list* -- and SEARCHES for neighbors, + returning a CARRY-ALL graph: every neighbor within ``rcut``. ``sel`` is + normalization-only (consumed downstream by the descriptor) and is NEVER a + cutoff here. This module ships the ``dense`` (all-pairs, O(N^2) reference) + search; O(N) ``vesin``/``ase`` backends land later behind a ``method`` key. + +2. **Converters (adapt an already-built list).** ``from_dense_quartet`` adapts an + existing extended quartet (extended_coord, nlist, mapping) into a graph. It + performs NO search and therefore INHERITS that quartet's ``sel`` truncation -- + it is the backward-compat bridge to the legacy dense nlist (World 1) and the + test oracle, NOT a carry-all path. The ``(i,j,S)`` converter (``from_ijs``, + fed by ASE/vesin/LAMMPS) lands with the dispatcher's O(N) backends. + +The dispatcher and the converters share the format-conversion code (a search +backend = search + its converter as the final step); they are separate only on +the question "did I get raw geometry, or an already-built list?". + +Both are fully vectorized over the frame axis (no Python frame loop): per-slot +``(frame, center, neighbor)`` index grids are flattened, masked, and gathered in +one shot, with cross-frame gathers done through ``frame * nall + idx`` flat +indices. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) + +import array_api_compat + +from .graph import ( + GraphLayout, + NeighborGraph, + pad_and_guard_edges, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + + +def from_dense_quartet( + extended_coord: Array, + nlist: Array, + mapping: Array, + layout: GraphLayout | None = None, +) -> NeighborGraph: + """Convert a legacy extended quartet into a ghost-free NeighborGraph (CONVERTER). + + This is a backward-compat CONVERTER (World 1 -> graph): it performs NO neighbor + search and INHERITS the ``sel`` truncation already baked into ``nlist``. Use it + only when a caller (an MD code, or the legacy dense path) already holds a + built quartet; for the carry-all graph use :func:`build_neighbor_graph`. + + For each valid neighbor slot it emits one edge with ``src = mapping[neighbor]`` + (the neighbor's LOCAL owner -> ghost-free index), ``dst = center`` (local), and + ``edge_vec = extended_coord[neighbor] - extended_coord[center]`` (the ghost + coordinate already carries the periodic shift). Invalid slots (``nlist == -1``) + are dropped. Nodes are flattened with a ``frame * nloc`` offset; the edge axis + is padded/guarded via ``pad_and_guard_edges``. + + Because every neighbor maps to a LOCAL owner, the resulting graph is ghost-free: + forces scatter to local atoms (periodic images of the same atom sum to one owner + through the ``src`` index), so no ``edge_scatter_index`` is needed (single-rank). + + Parameters + ---------- + extended_coord + (nf, nall, 3) extended (local + ghost) coordinates. + nlist + (nf, nloc, nsel) neighbor list into the extended atoms; -1 is padding. + mapping + (nf, nall) extended -> local-owner index (local atoms map to themselves). + layout + edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + """ + if layout is None: + layout = GraphLayout() + xp = array_api_compat.array_namespace(extended_coord, nlist, mapping) + dev = array_api_compat.device(extended_coord) + nf, nloc, nsel = nlist.shape + nall = extended_coord.shape[1] + # per-slot (nf, nloc, nsel) index grids, flattened frame-major + ff_grid = xp.broadcast_to( + xp.reshape(xp.arange(nf, dtype=xp.int64, device=dev), (nf, 1, 1)), + (nf, nloc, nsel), + ) + center_grid = xp.broadcast_to( + xp.reshape(xp.arange(nloc, dtype=xp.int64, device=dev), (1, nloc, 1)), + (nf, nloc, nsel), + ) + ff_flat = xp.reshape(ff_grid, (-1,)) + center_flat = xp.reshape(center_grid, (-1,)) + nl_flat = xp.reshape(nlist, (-1,)) + keep = xp.reshape(xp.nonzero(nl_flat >= 0)[0], (-1,)) + ff_k = xp.take(ff_flat, keep, axis=0) + dst_local = xp.take(center_flat, keep, axis=0) # center index in [0, nloc) + j_ext = xp.take(nl_flat, keep, axis=0) # neighbor index in [0, nall) + # cross-frame gathers via flat (frame * nall + idx) indices; centers are the + # first nloc extended atoms (local atoms precede ghosts). + ec_flat = xp.reshape(extended_coord, (nf * nall, 3)) + map_flat = xp.reshape(mapping, (nf * nall,)) + g_nei = ff_k * nall + j_ext + g_cen = ff_k * nall + dst_local + src_local = xp.take(map_flat, g_nei, axis=0) # local owner of the neighbor + edge_vec = xp.take(ec_flat, g_nei, axis=0) - xp.take(ec_flat, g_cen, axis=0) + edge_index = xp.astype( + xp.stack([ff_k * nloc + src_local, ff_k * nloc + dst_local], axis=0), xp.int64 + ) + edge_index, edge_vec, edge_mask = pad_and_guard_edges( + edge_index, edge_vec, layout.edge_capacity, layout.min_edges + ) + n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) + + +def build_neighbor_graph( + coord: Array, + atype: Array, + box: Array | None, + rcut: float, + layout: GraphLayout | None = None, +) -> NeighborGraph: + """Build a CARRY-ALL NeighborGraph DIRECTLY from coordinates (``dense`` search). + + This is the dispatcher's reference ``dense`` backend: it SEARCHES for neighbors + from raw geometry and emits EVERY neighbor within ``rcut``. It is **sel-free** -- + there is intentionally no ``sel`` parameter, because ``sel`` is normalization-only + (consumed by the descriptor downstream) and never a cutoff. It does NOT route + through the legacy dense nlist / :func:`from_dense_quartet`, so it carries no + ``sel`` truncation. + + Implementation: reuse the tested periodic ghosting + (:func:`~deepmd.dpmodel.utils.nlist.extend_coord_with_ghosts`) to materialise all + periodic images within ``rcut``, then enumerate all center-neighbor pairs within + ``rcut`` UNCAPPED, vectorized over frames. This is an O(N^2) reference search + (correctness oracle); the O(N) ``vesin``/``ase`` backends arrive later behind a + ``method`` key. Edges map every neighbor to its LOCAL owner + (``src = mapping[neighbor]``), so the graph is ghost-free. + + Parameters + ---------- + coord + (nf, nloc, 3) or (nf, nloc*3) local coordinates. + atype + (nf, nloc) local atom types; ``type < 0`` marks a virtual atom (excluded + as both a center and a neighbor). + box + (nf, 3, 3) or (nf, 9) simulation cell; ``None`` for non-periodic. + rcut + cutoff radius (neighbors kept where ``0 < |edge_vec| <= rcut``, matching the + legacy nlist convention so this coincides with :func:`from_dense_quartet` + at non-binding ``sel``). + layout + edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards. + """ + from deepmd.dpmodel.utils.nlist import ( + extend_coord_with_ghosts, + ) + from deepmd.dpmodel.utils.region import ( + normalize_coord, + ) + + if layout is None: + layout = GraphLayout() + xp = array_api_compat.array_namespace(coord, atype) + dev = array_api_compat.device(coord) + nf, nloc = atype.shape[:2] + coord = xp.reshape(coord, (nf, nloc, 3)) + if box is not None: + box = xp.reshape(box, (nf, 3, 3)) + coord = normalize_coord(coord, box) + extended_coord, extended_atype, mapping = extend_coord_with_ghosts( + coord, atype, box, rcut + ) + extended_coord = xp.reshape(extended_coord, (nf, -1, 3)) + nall = extended_coord.shape[1] + # all center-neighbor displacements: (nf, nloc, nall, 3) = ext[j] - center[i] + centers = extended_coord[:, :nloc, :] + diff = extended_coord[:, None, :, :] - centers[:, :, None, :] + dist = xp.linalg.vector_norm(diff, axis=-1) # (nf, nloc, nall) + # per-slot (nf, nloc, nall) index grids + ff_grid = xp.broadcast_to( + xp.reshape(xp.arange(nf, dtype=xp.int64, device=dev), (nf, 1, 1)), + (nf, nloc, nall), + ) + i_grid = xp.broadcast_to( + xp.reshape(xp.arange(nloc, dtype=xp.int64, device=dev), (1, nloc, 1)), + (nf, nloc, nall), + ) + j_grid = xp.broadcast_to( + xp.reshape(xp.arange(nall, dtype=xp.int64, device=dev), (1, 1, nall)), + (nf, nloc, nall), + ) + # keep neighbors within rcut, dropping: the self extended atom (i==j; a periodic + # IMAGE of i has j!=i and is kept), virtual neighbors, and virtual centers. + not_self = j_grid != i_grid + vir_nei = xp.broadcast_to((extended_atype < 0)[:, None, :], (nf, nloc, nall)) + vir_cen = xp.broadcast_to((atype < 0)[:, :, None], (nf, nloc, nall)) + keep_mask = ( + (dist <= rcut) & not_self & xp.logical_not(vir_nei) & xp.logical_not(vir_cen) + ) + keep = xp.reshape(xp.nonzero(xp.reshape(keep_mask, (-1,)))[0], (-1,)) + ff_k = xp.take(xp.reshape(ff_grid, (-1,)), keep, axis=0) + dst_local = xp.take(xp.reshape(i_grid, (-1,)), keep, axis=0) # local center + j_ext = xp.take(xp.reshape(j_grid, (-1,)), keep, axis=0) # extended neighbor + edge_vec = xp.take(xp.reshape(diff, (nf * nloc * nall, 3)), keep, axis=0) + # cross-frame neighbor-owner gather via flat (frame * nall + idx) + map_flat = xp.reshape(mapping, (nf * nall,)) + src_local = xp.take(map_flat, ff_k * nall + j_ext, axis=0) + edge_index = xp.astype( + xp.stack([ff_k * nloc + src_local, ff_k * nloc + dst_local], axis=0), xp.int64 + ) + edge_index, edge_vec, edge_mask = pad_and_guard_edges( + edge_index, edge_vec, layout.edge_capacity, layout.min_edges + ) + n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py new file mode 100644 index 0000000000..1c0bafc234 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Assemble per-node force and virial from a per-edge gradient g_e = dE/d(edge_vec). + +The autograd that produces g_e (grad(E, edge_vec)) is wired in the torch/jax +backend later; this pure-array-API assembly is shared by all backends. + +Conventions (see memory/spec_unified_edge_nlist.md): + edge_vec_e = r_src - r_dst ; F_k = sum_{dst=k} g - sum_{src=k} g + per-edge virial w_e = -g_e (x) edge_vec_e + atom virial attributed FULL-TO-src (canonical TF==pt-legacy convention) + per-frame virial = sum over the edges of that frame of w_e (DeePMD virials + are per frame; a multi-frame NeighborGraph must NOT collapse frames) +Padding/guard edges (edge_mask == 0) are zeroed before any scatter. +""" + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) + +from .segment import ( + segment_sum, +) + + +def edge_force_virial( + g_e: Array, + edge_vec: Array, + edge_index: Array, + edge_mask: Array, + n_node: Array, + node_capacity: int | None = None, +) -> tuple[Array, Array, Array]: + """Assemble per-node force/atom-virial and PER-FRAME virial from ``g_e``. + + Handles the fully general layout: multi-frame, RAGGED (different per-frame + node and edge counts), padding/guard EDGES (``edge_mask == 0``), and a padded + NODE axis (``node_capacity`` > ``sum(n_node)``). + + Parameters + ---------- + n_node + (nf,) per-frame REAL node counts. Real nodes occupy the compact prefix + ``[0, sum(n_node))`` frame-major; ``nf = n_node.shape[0]``. + node_capacity + Size of the (possibly padded) node axis ``N``. ``None`` => ``sum(n_node)`` + (no node padding — the torch/eager case). When set (jax static ``N_max``), + force/atom_virial are sized to it; padding nodes (never referenced by an + edge) get zero. Frame assignment is unaffected (padding nodes are not + ``dst`` of any real edge). + + Returns + ------- + force + (N, 3) per-node force. + atom_virial + (N, 3, 3) per-node virial, full-to-``src`` attribution. + virial + (nf, 3, 3) PER-FRAME virial. A multi-frame graph keeps each frame's + virial separate (DeePMD virials are per frame); edges are assigned to a + frame via the frame of their ``dst`` node. + """ + xp = array_api_compat.array_namespace(g_e) + n_real = int(xp.sum(n_node)) # real node count + n_out = n_real if node_capacity is None else int(node_capacity) # node-axis size + nf = n_node.shape[0] + # zero padding/guard contributions; cast mask to g's dtype (array-API pure, + # CLAUDE.md mask-multiply guideline — avoids bool*float under array_api_strict) + g = g_e * xp.astype(edge_mask[:, None], g_e.dtype) + src = edge_index[0] + dst = edge_index[1] + # force (output sized to the node axis, incl. any padding tail) + force = segment_sum(g, dst, n_out) - segment_sum(g, src, n_out) + # per-edge virial w_e[k, j] = -g_e[k] * edge_vec[j] (broadcast, no einsum) + w_edge = -(g[:, :, None] * edge_vec[:, None, :]) # (E, 3, 3) + # atom virial: full-to-src + atom_virial = segment_sum(w_edge, src, n_out) # (N, 3, 3) + # per-frame virial: assign each edge to the frame of its dst node. Node + # ``k`` belongs to frame ``searchsorted(cumsum(n_node), k, "right")`` because + # real nodes are compact frame-major (frame f owns a contiguous block). + boundaries = xp.cumulative_sum(n_node) # (nf,) per-frame node upper bounds + edge_frame = xp.astype( + xp.searchsorted(boundaries, dst, side="right"), xp.int64 + ) # (E,) in [0, nf) + virial = segment_sum(w_edge, edge_frame, nf) # (nf, 3, 3) + return force, atom_virial, virial diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py new file mode 100644 index 0000000000..232145bda0 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-agnostic edge-graph neighbor-list contract (NeighborGraph) and its +length policy (GraphLayout). See memory/spec_unified_edge_nlist.md. + +Node validity (real vs padding) is NOT a stored field: it is derived as +``arange(N) < sum(n_node)`` because ``n_node`` already encodes the real-node +count and the layout is compact-prefix (real nodes first, padding suffix). +``edge_mask`` IS stored because there is no per-axis edge count to derive it from. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import array_api_compat + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) + + +@dataclass +class NeighborGraph: + """Edge-graph neighbor list. Node axis is flat ``N = sum(n_node)``. + + Geometry enters the model ONLY through ``edge_vec`` (the single autograd + leaf). ``edge_index``/``angle_index`` use the SoA ``(2, .)`` layout so the + src/dst index vectors are contiguous. + """ + + n_node: Array + """(nf,) int nodes per frame (single-rank: local atoms; multi-rank: local+halo).""" + edge_index: Array + """(2, E) int [src, dst]; src = neighbor, dst = center; both in [0, N).""" + edge_vec: Array + """(E, 3) float r_src - r_dst (neighbor - center); the only geometry / autograd leaf.""" + edge_mask: Array + """(E,) bool real (1) vs padding (0). Always stored (no n_edge to derive from).""" + n_local: Array | None = None + """(nf,) int multi-rank owned-vs-halo split; owned = first n_local[f]. None = all local.""" + angle_index: Array | None = None + """(2, A) int [edge_a, edge_b] sharing a center; into [0, E). None if no angles.""" + angle_mask: Array | None = None + """(A,) bool real vs padding on the angle axis. None if no angles.""" + + +@dataclass +class GraphLayout: + """Length policy: the only torch/jax difference. None => dynamic axis (torch); + int => static capacity (jax/paddle padding target). + """ + + edge_capacity: int | None = None + angle_capacity: int | None = None + node_capacity: int | None = None + frame_capacity: int | None = None + min_edges: int = 2 + + +def pad_and_guard_edges( + edge_index: Array, + edge_vec: Array, + capacity: int | None, + min_edges: int = 2, + pad_value: int = 0, +) -> tuple[Array, Array, Array]: + """Append padding/guard edges as a contiguous suffix and build edge_mask. + + Real edges (``edge_index``/``edge_vec``) stay at the front (compact layout). + - ``capacity is None`` (torch dynamic): append exactly ``min_edges`` masked + dummy edges so the edge axis has a known lower bound and shape-stable + guards for export. + - ``capacity`` set (jax static): pad to ``E_max = capacity``; raise on overflow. + Dummy edges point at node ``pad_value`` (in-range) with zero ``edge_vec``. + """ + xp = array_api_compat.array_namespace(edge_index) + dev = array_api_compat.device(edge_index) + e_real = edge_index.shape[1] + if capacity is None: + target = e_real + min_edges + else: + if e_real > capacity: + raise ValueError( + f"edge overflow: {e_real} real edges > edge_capacity {capacity}" + ) + target = capacity + n_pad = target - e_real + pad_idx = xp.full((2, n_pad), pad_value, dtype=edge_index.dtype, device=dev) + pad_vec = xp.zeros((n_pad, 3), dtype=edge_vec.dtype, device=dev) + ei = xp.concat([edge_index, pad_idx], axis=1) + ev = xp.concat([edge_vec, pad_vec], axis=0) + arange = xp.arange(target, dtype=edge_index.dtype, device=dev) + edge_mask = arange < e_real + return ei, ev, edge_mask + + +def node_validity_mask(n_node: Array, n_total: int) -> Array: + """Derive the (n_total,) real-vs-padding node mask from per-frame counts. + + Compact-prefix layout: the first ``sum(n_node)`` nodes are real, the rest + are padding. jit-safe (no Python ``int`` cast on the traced sum). + """ + xp = array_api_compat.array_namespace(n_node) + idx = xp.arange(n_total, dtype=n_node.dtype, device=array_api_compat.device(n_node)) + return idx < xp.sum(n_node) diff --git a/deepmd/dpmodel/utils/neighbor_graph/segment.py b/deepmd/dpmodel/utils/neighbor_graph/segment.py new file mode 100644 index 0000000000..45d64af08c --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/segment.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Mask-aware, backend-dispatched segment reductions (the dpmodel scatter +primitive). Built on deepmd.dpmodel.array_api.xp_add_at so they work for +numpy / jax / torch. segment_index must be int64 (torch index_add requirement). +""" + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, + xp_add_at, +) + + +def segment_sum(data: Array, segment_ids: Array, num_segments: int) -> Array: + """out[s] = sum of data[i] over i with segment_ids[i] == s. Shape + ``(num_segments, *data.shape[1:])``; empty segments are zero. + """ + xp = array_api_compat.array_namespace(data) + out = xp.zeros( + (num_segments, *tuple(data.shape[1:])), + dtype=data.dtype, + device=array_api_compat.device(data), + ) + return xp_add_at(out, segment_ids, data) + + +def segment_mean(data: Array, segment_ids: Array, num_segments: int) -> Array: + """Per-segment mean; empty segments are zero (no division by zero).""" + xp = array_api_compat.array_namespace(data) + summed = segment_sum(data, segment_ids, num_segments) + ones = xp.ones( + (data.shape[0],), dtype=data.dtype, device=array_api_compat.device(data) + ) + counts = segment_sum(ones[:, None], segment_ids, num_segments) # (num_segments, 1) + safe = xp.where(counts == 0, xp.ones_like(counts), counts) + # broadcast counts over the trailing dims of summed + shape = (num_segments,) + (1,) * (summed.ndim - 1) + return summed / xp.reshape(safe, shape) diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py new file mode 100644 index 0000000000..fa84ef7ba4 --- /dev/null +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.utils.neighbor_graph import ( + edge_force_virial, +) + + +class TestEdgeForceVirial(unittest.TestCase): + def setUp(self) -> None: + # 1 frame, 2 nodes, 2 real edges: e0 = (src=1, dst=0), e1 = (src=0, dst=1) + self.edge_index = np.array([[1, 0], [0, 1]], dtype=np.int64) + self.edge_vec = np.array([[1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]) + self.edge_mask = np.array([True, True]) + self.g = np.array([[0.5, 0.0, 0.0], [0.3, 0.0, 0.0]]) # per-edge grad + self.n_node = np.array([2], dtype=np.int64) # one frame, 2 nodes + + def test_force_formula(self) -> None: + force, _, _ = edge_force_virial( + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.n_node + ) + # F_k = sum_{dst=k} g - sum_{src=k} g + # node 0: dst of e0 (+g0), src of e1 (-g1) => 0.5 - 0.3 = 0.2 + # node 1: dst of e1 (+g1), src of e0 (-g0) => 0.3 - 0.5 = -0.2 + np.testing.assert_allclose(force[:, 0], np.array([0.2, -0.2])) + + def test_atom_virial_full_to_src_sums_to_frame_virial(self) -> None: + _, av, vir = edge_force_virial( + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.n_node + ) + self.assertEqual(av.shape, (2, 3, 3)) + # all 2 nodes are in frame 0 => their atom-virials sum to that frame's virial + np.testing.assert_allclose(np.sum(av, axis=0), vir[0]) + # full-to-src: e0 virial on node 1 (src), e1 virial on node 0 (src) + w0 = -np.einsum("k,j->kj", self.g[0], self.edge_vec[0]) + w1 = -np.einsum("k,j->kj", self.g[1], self.edge_vec[1]) + np.testing.assert_allclose(av[1], w0) # src of e0 is node 1 + np.testing.assert_allclose(av[0], w1) # src of e1 is node 0 + + def test_empty_frame_no_nodes_or_edges(self) -> None: + # 3 frames with the MIDDLE one EMPTY (0 nodes/edges): n_node=[3,0,5]. + # frame 0 = {0,1,2}, frame 1 = EMPTY, frame 2 = {3,4,5,6,7}. Also a padded + # node axis (node_capacity 9 > sum 8). Verifies the zero-width-block frame + # assignment (searchsorted on duplicate cumsum boundaries [3,3,8] must skip + # the empty frame) and that the empty frame's virial is exactly zero. + n_node = np.array([3, 0, 5], dtype=np.int64) # sum = 8 + node_capacity = 9 # 1 padded node slot (8) + edge_index = np.array( + [ + [1, 4, 6], # src + [0, 3, 7], + ], # dst (frame 0: dst 0 ; frame 2: dst 3,7 ; frame 1: NONE) + dtype=np.int64, + ) + edge_vec = np.array([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]]) + edge_mask = np.array([True, True, True]) + g = np.array([[0.5, 0.0, 0.0], [0.0, 0.4, 0.0], [0.0, 0.0, 0.6]]) + force, av, vir = edge_force_virial( + g, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity + ) + self.assertEqual(vir.shape, (3, 3, 3)) + self.assertEqual(force.shape, (9, 3)) + # the empty middle frame contributes exactly zero + np.testing.assert_allclose(vir[1], 0.0) + w = [-np.einsum("k,j->kj", g[i], edge_vec[i]) for i in range(3)] + np.testing.assert_allclose(vir[0], w[0]) # frame 0: edge 0 + np.testing.assert_allclose( + vir[2], w[1] + w[2] + ) # frame 2: edges 1,2 (node 3..7) + # padded node slot (8) is unreferenced -> zero + np.testing.assert_allclose(force[8], 0.0) + np.testing.assert_allclose(av[8], 0.0) + # per-frame atom-virial closure across the empty frame: frame-2 nodes 3..7 + np.testing.assert_allclose(np.sum(av[3:8], axis=0), vir[2]) + + def test_all_edges_masked_gives_zero(self) -> None: + # ALL-EMPTY: nodes exist but there are ZERO real edges (isolated atoms, or + # rcut below all pair distances) -> only masked guard edges remain. Single- + # and multi-frame; every output must be exactly zero with correct shapes. + for n_node in ( + np.array([3], dtype=np.int64), # single frame + np.array([2, 3], dtype=np.int64), # multi-frame + ): + nf = int(n_node.shape[0]) + n = int(n_node.sum()) + # two masked guard edges at pad node 0 with nonzero g (must be ignored) + edge_index = np.array([[0, 0], [0, 0]], dtype=np.int64) + edge_vec = np.array([[9.0, 9.0, 9.0], [9.0, 9.0, 9.0]]) + edge_mask = np.array([False, False]) + g = np.array([[7.0, 7.0, 7.0], [7.0, 7.0, 7.0]]) + force, av, vir = edge_force_virial( + g, edge_vec, edge_index, edge_mask, n_node + ) + np.testing.assert_allclose(force, np.zeros((n, 3))) + np.testing.assert_allclose(av, np.zeros((n, 3, 3))) + np.testing.assert_allclose(vir, np.zeros((nf, 3, 3))) + + def test_ragged_multiframe_with_edge_and_node_padding(self) -> None: + # MOST GENERAL case: 2 frames with DIFFERENT node counts (3 and 5) AND + # different edge counts (2 and 3), masked guard EDGES, and a padded NODE + # axis (node_capacity 10 > sum(n_node)=8). + n_node = np.array( + [3, 5], dtype=np.int64 + ) # ragged: frame0={0,1,2}, frame1={3..7} + node_capacity = 10 # 2 padded node slots (8, 9) at the global tail + edge_index = np.array( + [ + [1, 2, 4, 5, 6, 0, 0], # src + [0, 1, 3, 4, 7, 0, 0], + ], # dst (frame0: dst 0,1 ; frame1: dst 3,4,7) + dtype=np.int64, + ) + edge_vec = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], # frame 0 (2 edges) + [0.0, 0.0, 1.0], + [2.0, 0.0, 0.0], + [0.0, 2.0, 0.0], # frame 1 (3 edges) + [9.0, 9.0, 9.0], + [9.0, 9.0, 9.0], # masked guard edges + ] + ) + edge_mask = np.array([True, True, True, True, True, False, False]) + g = np.array( + [ + [0.5, 0.0, 0.0], + [0.0, 0.3, 0.0], + [0.0, 0.0, 0.7], + [0.1, 0.0, 0.0], + [0.0, 0.2, 0.0], + [7.0, 7.0, 7.0], + [7.0, 7.0, 7.0], + ] + ) + force, av, vir = edge_force_virial( + g, edge_vec, edge_index, edge_mask, n_node, node_capacity=node_capacity + ) + # shapes: padded node axis + per-frame virial + self.assertEqual(force.shape, (10, 3)) + self.assertEqual(av.shape, (10, 3, 3)) + self.assertEqual(vir.shape, (2, 3, 3)) + # padded node slots (8, 9) are never referenced -> zero + np.testing.assert_allclose(force[8:], 0.0) + np.testing.assert_allclose(av[8:], 0.0) + # per-frame virial = sum of THAT frame's real edges only (ragged edge counts) + w = [-np.einsum("k,j->kj", g[i], edge_vec[i]) for i in range(5)] + np.testing.assert_allclose(vir[0], w[0] + w[1]) # frame 0: edges 0,1 + np.testing.assert_allclose(vir[1], w[2] + w[3] + w[4]) # frame 1: edges 2,3,4 + self.assertFalse(np.allclose(vir[0], vir[1])) + # per-frame atom-virial closure (ragged node blocks): frame nodes -> frame virial + np.testing.assert_allclose( + np.sum(av[0:3], axis=0), vir[0] + ) # frame 0 nodes 0,1,2 + np.testing.assert_allclose( + np.sum(av[3:8], axis=0), vir[1] + ) # frame 1 nodes 3..7 + # guard edges contributed nothing: result == running with real edges only + f2, a2, v2 = edge_force_virial( + g[:5], + edge_vec[:5], + edge_index[:, :5], + edge_mask[:5], + n_node, + node_capacity=node_capacity, + ) + np.testing.assert_allclose(force, f2) + np.testing.assert_allclose(av, a2) + np.testing.assert_allclose(vir, v2) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py new file mode 100644 index 0000000000..ef8066850b --- /dev/null +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.utils.neighbor_graph import ( + GraphLayout, + NeighborGraph, +) + + +class TestNeighborGraphDataclass(unittest.TestCase): + def test_construct_minimal(self) -> None: + ng = NeighborGraph( + n_node=np.array([2], dtype=np.int64), + edge_index=np.array([[1, 0], [0, 1]], dtype=np.int64), # (2, E) + edge_vec=np.zeros((2, 3), dtype=np.float64), + edge_mask=np.array([True, True]), + ) + self.assertEqual(ng.edge_index.shape, (2, 2)) + self.assertEqual(ng.edge_vec.shape, (2, 3)) + # optionals default to None + self.assertIsNone(ng.n_local) + self.assertIsNone(ng.angle_index) + self.assertIsNone(ng.angle_mask) + + def test_graphlayout_defaults(self) -> None: + lay = GraphLayout() + self.assertIsNone(lay.edge_capacity) + self.assertIsNone(lay.angle_capacity) + self.assertIsNone(lay.node_capacity) + self.assertIsNone(lay.frame_capacity) + self.assertEqual(lay.min_edges, 2) + + +from deepmd.dpmodel.utils.neighbor_graph import ( + node_validity_mask, +) + + +class TestNodeValidityMask(unittest.TestCase): + def test_no_padding_all_true(self) -> None: + n_node = np.array([2, 3], dtype=np.int64) # sum = 5 + mask = node_validity_mask(n_node, 5) + np.testing.assert_array_equal(mask, np.array([True] * 5)) + + def test_with_padding_prefix(self) -> None: + n_node = np.array([2, 3], dtype=np.int64) # 5 real + mask = node_validity_mask(n_node, 8) # N_max = 8 => 3 padding + np.testing.assert_array_equal(mask, np.array([True] * 5 + [False] * 3)) + + +from deepmd.dpmodel.utils.neighbor_graph import ( + pad_and_guard_edges, +) + + +class TestPadAndGuardEdges(unittest.TestCase): + def setUp(self) -> None: + self.edge_index = np.array([[1, 0, 2], [0, 1, 0]], dtype=np.int64) # E=3 + self.edge_vec = np.arange(9, dtype=np.float64).reshape(3, 3) + + def test_dynamic_appends_min_edges_guards(self) -> None: + # capacity=None (torch): append min_edges masked dummies at the tail + ei, ev, em = pad_and_guard_edges( + self.edge_index, self.edge_vec, capacity=None, min_edges=2 + ) + self.assertEqual(ei.shape, (2, 5)) # 3 real + 2 guard + self.assertEqual(ev.shape, (5, 3)) + np.testing.assert_array_equal(em, np.array([True, True, True, False, False])) + # real edges unchanged at the front + np.testing.assert_array_equal(ei[:, :3], self.edge_index) + np.testing.assert_allclose(ev[:3], self.edge_vec) + # guard edges are zero-vec, in-range index (pad_value=0) + np.testing.assert_allclose(ev[3:], 0.0) + np.testing.assert_array_equal(ei[:, 3:], 0) + + def test_static_capacity_pads_to_E_max(self) -> None: + ei, ev, em = pad_and_guard_edges( + self.edge_index, self.edge_vec, capacity=6, min_edges=2 + ) + self.assertEqual(ei.shape, (2, 6)) + np.testing.assert_array_equal( + em, np.array([True, True, True, False, False, False]) + ) + + def test_overflow_raises(self) -> None: + with self.assertRaises(ValueError): + pad_and_guard_edges(self.edge_index, self.edge_vec, capacity=2, min_edges=2) + + +class TestPublicExports(unittest.TestCase): + def test_importable_from_utils(self) -> None: + from deepmd.dpmodel.utils import ( + GraphLayout, + NeighborGraph, + build_neighbor_graph, + edge_force_virial, + from_dense_quartet, + segment_sum, + ) + + self.assertTrue(callable(segment_sum)) + self.assertTrue(callable(edge_force_virial)) + self.assertTrue(callable(build_neighbor_graph)) + self.assertTrue(callable(from_dense_quartet)) + self.assertIsNotNone(NeighborGraph) + self.assertIsNotNone(GraphLayout) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py new file mode 100644 index 0000000000..9ba25c0ccb --- /dev/null +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -0,0 +1,317 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the dpmodel NeighborGraph builder/converter. + +``build_neighbor_graph`` is the CARRY-ALL ``dense`` search backend: it builds a +graph DIRECTLY from coordinates and keeps EVERY neighbor within ``rcut`` (no +``sel`` truncation). We validate it against an INDEPENDENT brute-force all-pairs +oracle defined locally in this test file. + +``from_dense_quartet`` is the backward-compat CONVERTER: it adapts an existing +(``sel``-truncated) extended quartet and performs no search. +""" + +import itertools +import unittest + +import numpy as np + +from deepmd.dpmodel.utils.neighbor_graph import ( + GraphLayout, + build_neighbor_graph, + from_dense_quartet, +) +from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, +) + + +def brute_force_neighbor_sets(coord, box, rcut): + """Independent all-pairs oracle: per center i, the set of (local-owner j, + rounded edge_vec) within rcut. edge_vec = coord[j] + S@box - coord[i]. + """ + nloc = coord.shape[0] + if box is None: + shells = [np.zeros(3, dtype=np.int64)] + else: + h = np.min(np.abs(np.diag(box))) + n = int(np.ceil(rcut / h)) + shells = [ + np.array(s, dtype=np.int64) + for s in itertools.product(range(-n, n + 1), repeat=3) + ] + sets = [set() for _ in range(nloc)] + for s in shells: + sc = np.zeros(3) if box is None else s.astype(float) @ box + for i in range(nloc): + for j in range(nloc): + vec = coord[j] + sc - coord[i] + r = np.linalg.norm(vec) + if 1e-10 < r < rcut: + sets[i].add((j, tuple(np.round(vec, 6)))) + return sets + + +def graph_neighbor_sets(ng, nloc): + """Per dst-center, the set of (src local owner, rounded edge_vec); real edges only.""" + ei = ng.edge_index[:, ng.edge_mask] + ev = ng.edge_vec[ng.edge_mask] + sets = [set() for _ in range(nloc)] + for k in range(ei.shape[1]): + src, dst = int(ei[0, k]), int(ei[1, k]) + sets[dst].add((src, tuple(np.round(ev[k], 6)))) + return sets + + +def graph_neighbor_sets_frame(ng, frame, nloc): + """Per-frame neighbor sets (src/dst de-offset to local [0, nloc)); real edges only. + + Selects the edges whose dst lives in frame ``frame``'s node block + ``[frame*nloc, (frame+1)*nloc)`` and de-offsets indices, so the result is + directly comparable to a single-frame oracle. + """ + off = frame * nloc + ei = ng.edge_index[:, ng.edge_mask] + ev = ng.edge_vec[ng.edge_mask] + sets = [set() for _ in range(nloc)] + for k in range(ei.shape[1]): + src, dst = int(ei[0, k]), int(ei[1, k]) + if off <= dst < off + nloc: + sets[dst - off].add((src - off, tuple(np.round(ev[k], 6)))) + return sets + + +class TestNeighborGraphBuilder(unittest.TestCase): + def setUp(self) -> None: + self.rcut = 4.0 + # atom 2 at y=2.3 (not 2.0): avoids a degenerate pair sitting exactly at + # rcut under PBC (box 6, image distance 6-2=4==rcut), where strict-< vs + # <= cutoff conventions disagree. Real geometries never sit exactly at rcut. + 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) + + def test_nonperiodic_matches_brute_force(self) -> None: + ng = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + np.testing.assert_array_equal(ng.n_node, np.array([4], dtype=np.int64)) + self.assertEqual( + graph_neighbor_sets(ng, 4), + brute_force_neighbor_sets(self.coord[0], None, self.rcut), + ) + + def test_periodic_matches_brute_force(self) -> None: + box = np.eye(3, dtype=np.float64)[None] * 6.0 + ng = build_neighbor_graph(self.coord, self.atype, box, self.rcut) + self.assertEqual( + graph_neighbor_sets(ng, 4), + brute_force_neighbor_sets(self.coord[0], box[0], self.rcut), + ) + + def test_neighbor_only_across_periodic_boundary(self) -> None: + # DISCRIMINATING PBC case: a pair that is a neighbor ONLY across the + # boundary. atoms at x=0.5 and x=5.5 in a box of 6: direct distance 5.0 > + # rcut=4 (NOT a direct neighbor), but the minimum image is 1.0 < rcut. + # A build that ignored periodic images would find ZERO edges here. + box = np.eye(3, dtype=np.float64)[None] * 6.0 + coord = np.array([[0.5, 0.0, 0.0], [5.5, 0.0, 0.0]], dtype=np.float64).reshape( + 1, 2, 3 + ) + atype = np.array([[0, 0]], dtype=np.int64) + ng = build_neighbor_graph(coord, atype, box, self.rcut) + got = graph_neighbor_sets(ng, 2) # per-center list of neighbor sets + # each atom's ONLY neighbor is the other's periodic image, at +-1.0 + want = [{(1, (-1.0, 0.0, 0.0))}, {(0, (1.0, 0.0, 0.0))}] + self.assertEqual(got, want) + # the direct (non-image) separation of 5.0 must NOT appear as an edge + ev = ng.edge_vec[ng.edge_mask] + self.assertFalse(bool(np.any(np.linalg.norm(ev, axis=1) > 4.0))) + # independent brute-force oracle agrees on the cross-boundary environment + self.assertEqual(got, brute_force_neighbor_sets(coord[0], box[0], self.rcut)) + # and WITHOUT the box the same atoms are NOT neighbors (direct 5.0 > rcut) + ng_free = build_neighbor_graph(coord, atype, None, self.rcut) + self.assertEqual(int(ng_free.edge_mask.sum()), 0) + + def test_edge_vec_within_rcut(self) -> None: + ng = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + ev = ng.edge_vec[ng.edge_mask] + self.assertTrue(np.all(np.linalg.norm(ev, axis=1) < self.rcut)) + + def test_carry_all_keeps_more_than_truncated_quartet(self) -> None: + # THE carry-all contract: with a binding ``sel``, the legacy quartet + # converter drops real neighbors, but the dense search keeps them all. + box = np.eye(3, dtype=np.float64)[None] * 6.0 + # sel=1 per type -> heavily truncates under PBC (many images within rcut). + ext_coord, _ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + self.coord, self.atype, self.rcut, [1, 1], mixed_types=True, box=box + ) + ng_trunc = from_dense_quartet(ext_coord, nlist, mapping) + ng_all = build_neighbor_graph(self.coord, self.atype, box, self.rcut) + n_trunc = int(ng_trunc.edge_mask.sum()) + n_all = int(ng_all.edge_mask.sum()) + n_oracle = sum( + len(s) for s in brute_force_neighbor_sets(self.coord[0], box[0], self.rcut) + ) + # the truncated converter loses edges; the carry-all search recovers them all + self.assertLess(n_trunc, n_all) + self.assertEqual(n_all, n_oracle) + + def test_multiframe_per_frame_neighbor_sets(self) -> None: + # TWO DIFFERENT frames -> different per-frame EDGE counts. (Node counts are + # equal because build_neighbor_graph takes a rectangular (nf,nloc,3) coord; + # ragged node counts need a future ragged builder and are exercised on the + # flat primitives, e.g. test_edge_force_virial multi-frame.) + coord_b = np.array( + [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 1.0, 0.0], [1.5, 1.5, 0.0]], + dtype=np.float64, + ).reshape(1, 4, 3) + coord2 = np.concatenate([self.coord, coord_b], axis=0) # (2,4,3), DIFFERENT + atype2 = np.concatenate([self.atype, self.atype], axis=0) + ng = build_neighbor_graph(coord2, atype2, None, self.rcut) + np.testing.assert_array_equal(ng.n_node, np.array([4, 4], dtype=np.int64)) + # each frame's edges match THAT frame's own brute-force oracle + self.assertEqual( + graph_neighbor_sets_frame(ng, 0, 4), + brute_force_neighbor_sets(coord2[0], None, self.rcut), + ) + self.assertEqual( + graph_neighbor_sets_frame(ng, 1, 4), + brute_force_neighbor_sets(coord2[1], None, self.rcut), + ) + # the two frames are genuinely different environments (different edge sets) + self.assertNotEqual( + graph_neighbor_sets_frame(ng, 0, 4), + graph_neighbor_sets_frame(ng, 1, 4), + ) + # node-offset invariant: frame-0 edges in [0,4), frame-1 in [4,8) + ei = ng.edge_index[:, ng.edge_mask] + self.assertTrue(np.all(ei[:, ei[1] < 4] < 4)) + self.assertTrue(np.all(ei[:, ei[1] >= 4] >= 4)) + + def test_multiframe_periodic_per_frame(self) -> None: + box = np.eye(3, dtype=np.float64)[None] * 6.0 + coord2 = np.concatenate([self.coord, self.coord + 0.3], axis=0) # different + atype2 = np.concatenate([self.atype, self.atype], axis=0) + box2 = np.concatenate([box, box], axis=0) + ng = build_neighbor_graph(coord2, atype2, box2, self.rcut) + for f in (0, 1): + self.assertEqual( + graph_neighbor_sets_frame(ng, f, 4), + brute_force_neighbor_sets(coord2[f], box2[f], self.rcut), + ) + + def test_virtual_atoms_excluded(self) -> None: + # a virtual atom (type < 0) is excluded BOTH as a center (dst) and as a + # neighbor (src). atom 0 (origin) has in-range neighbors 1 (dist 1.0) and + # 2 (dist 2.3), so making it virtual actively exercises center-exclusion: + # without the virtual-center guard, edges 0<-1 and 0<-2 would appear. + atype = np.array([[-1, 1, 0, 1]], dtype=np.int64) # atom 0 virtual + ng = build_neighbor_graph(self.coord, atype, None, self.rcut) + ei = ng.edge_index[:, ng.edge_mask] + src, dst = ei[0], ei[1] + self.assertFalse(bool(np.any(dst == 0))) # never a center (center exclusion) + self.assertFalse(bool(np.any(src == 0))) # never a neighbor (neighbor excl.) + # the remaining real atoms still neighbor each other (we didn't nuke all edges) + self.assertGreater(int(ng.edge_mask.sum()), 0) + + def test_min_edges_guard_pads_sparse_frame(self) -> None: + # a single isolated atom yields ZERO real edges; the dynamic (capacity=None) + # layout must still emit the min_edges=2 guard edges, all masked out. + coord = np.zeros((1, 1, 3), dtype=np.float64) + atype = np.array([[0]], dtype=np.int64) + ng = build_neighbor_graph(coord, atype, None, self.rcut) # default layout + self.assertEqual(ng.edge_index.shape[1], 2) # min_edges guard edges + self.assertEqual(int(ng.edge_mask.sum()), 0) # none real + self.assertTrue(np.all(ng.edge_vec == 0.0)) + + def test_flat_coord_input_matches_rectangular(self) -> None: + # coord given flattened (nf, nloc*3) must match the (nf, nloc, 3) form. + coord_flat = self.coord.reshape(1, 4 * 3) + ng_flat = build_neighbor_graph(coord_flat, self.atype, None, self.rcut) + ng_rect = build_neighbor_graph(self.coord, self.atype, None, self.rcut) + self.assertEqual( + graph_neighbor_sets(ng_flat, 4), graph_neighbor_sets(ng_rect, 4) + ) + + def test_static_capacity_padding(self) -> None: + ng = build_neighbor_graph( + self.coord, + self.atype, + None, + self.rcut, + layout=GraphLayout(edge_capacity=64), + ) + self.assertEqual(ng.edge_index.shape[1], 64) + self.assertEqual(ng.edge_vec.shape[0], 64) + # exactly the real edges are marked, padded compactly at the tail + n_real = sum( + len(s) for s in brute_force_neighbor_sets(self.coord[0], None, self.rcut) + ) + self.assertEqual(int(ng.edge_mask.sum()), n_real) + self.assertTrue(bool(np.all(ng.edge_mask[:n_real]))) + self.assertFalse(bool(np.any(ng.edge_mask[n_real:]))) + # masked-out tail contributes no real edges + self.assertTrue(np.all(ng.edge_vec[~ng.edge_mask] == 0.0)) + + +class TestFromDenseQuartet(unittest.TestCase): + def test_adapter_on_handmade_quartet(self) -> None: + # 2 local atoms, no ghosts; each is the other's only neighbor. + extended_coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]]) # (1,2,3) + nlist = np.array([[[1, -1], [0, -1]]], dtype=np.int64) # (1,2,2) + mapping = np.array([[0, 1]], dtype=np.int64) # (1,2) local->self + ng = from_dense_quartet(extended_coord, nlist, mapping) + ei = ng.edge_index[:, ng.edge_mask] + ev = ng.edge_vec[ng.edge_mask] + got = { + (int(ei[0, k]), int(ei[1, k]), tuple(np.round(ev[k], 6))) + for k in range(ei.shape[1]) + } + want = { + (1, 0, (1.0, 0.0, 0.0)), # center 0, neighbor 1, vec = r1 - r0 + (0, 1, (-1.0, 0.0, 0.0)), # center 1, neighbor 0, vec = r0 - r1 + } + self.assertEqual(got, want) + + def test_adapter_multiframe_offsets(self) -> None: + # 2 frames, 2 local atoms each; each atom's only neighbor is the other. + # Frame 1 has a different separation so its edge_vec differs. + extended_coord = np.array( + [ + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], # frame 0 + [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], # frame 1 (different) + ] + ) + nlist = np.array( + [[[1, -1], [0, -1]], [[1, -1], [0, -1]]], dtype=np.int64 + ) # (2,2,2) + mapping = np.array([[0, 1], [0, 1]], dtype=np.int64) + ng = from_dense_quartet(extended_coord, nlist, mapping) + np.testing.assert_array_equal(ng.n_node, np.array([2, 2], dtype=np.int64)) + ei = ng.edge_index[:, ng.edge_mask] + ev = ng.edge_vec[ng.edge_mask] + per = {} + for k in range(ei.shape[1]): + per[(int(ei[0, k]), int(ei[1, k]))] = tuple(np.round(ev[k], 6)) + # frame 0 nodes {0,1} with sep 1.0; frame 1 nodes {2,3} with sep 2.0 + self.assertEqual(per[(1, 0)], (1.0, 0.0, 0.0)) + self.assertEqual(per[(0, 1)], (-1.0, 0.0, 0.0)) + self.assertEqual(per[(3, 2)], (2.0, 0.0, 0.0)) + self.assertEqual(per[(2, 3)], (-2.0, 0.0, 0.0)) + + def test_adapter_maps_ghost_to_local_owner(self) -> None: + # 1 local atom (0) + 1 ghost (1) which is a periodic image of atom 0. + extended_coord = np.array([[[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]]]) # (1,2,3) + nlist = np.array([[[1, -1]]], dtype=np.int64) # (1, nloc=1, nsel=2) + mapping = np.array([[0, 0]], dtype=np.int64) # ghost 1 -> owner 0 + ng = from_dense_quartet(extended_coord, nlist, mapping) + ei = ng.edge_index[:, ng.edge_mask] + ev = ng.edge_vec[ng.edge_mask] + self.assertEqual(ei.shape[1], 1) + # src = local owner of the ghost (0), dst = center (0); vec carries the shift + self.assertEqual((int(ei[0, 0]), int(ei[1, 0])), (0, 0)) + np.testing.assert_allclose(ev[0], np.array([3.0, 0.0, 0.0])) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/common/dpmodel/test_segment.py b/source/tests/common/dpmodel/test_segment.py new file mode 100644 index 0000000000..22046911e9 --- /dev/null +++ b/source/tests/common/dpmodel/test_segment.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.utils.neighbor_graph import ( + segment_mean, + segment_sum, +) + + +class TestSegment(unittest.TestCase): + def test_segment_sum_1d_values(self) -> None: + data = np.array([[1.0], [2.0], [3.0], [4.0], [5.0]]) + seg = np.array([0, 0, 1, 1, 2], dtype=np.int64) + out = segment_sum(data, seg, 3) + np.testing.assert_allclose(out, np.array([[3.0], [7.0], [5.0]])) + + def test_segment_sum_matrix_values(self) -> None: + # (E, 3, 3) per-edge tensors aggregate per segment + data = np.ones((4, 3, 3)) + seg = np.array([0, 0, 0, 1], dtype=np.int64) + out = segment_sum(data, seg, 2) + self.assertEqual(out.shape, (2, 3, 3)) + np.testing.assert_allclose(out[0], 3.0 * np.ones((3, 3))) + np.testing.assert_allclose(out[1], np.ones((3, 3))) + + def test_segment_sum_empty_segment_is_zero(self) -> None: + data = np.array([[1.0], [2.0]]) + seg = np.array([0, 2], dtype=np.int64) # segment 1 gets nothing + out = segment_sum(data, seg, 3) + np.testing.assert_allclose(out, np.array([[1.0], [0.0], [2.0]])) + + def test_segment_mean(self) -> None: + data = np.array([[2.0], [4.0], [9.0]]) + seg = np.array([0, 0, 1], dtype=np.int64) + out = segment_mean(data, seg, 2) + np.testing.assert_allclose(out, np.array([[3.0], [9.0]])) + + def test_segment_mean_empty_segment_no_nan(self) -> None: + data = np.array([[2.0], [4.0]]) + seg = np.array([0, 0], dtype=np.int64) + out = segment_mean(data, seg, 2) # segment 1 empty -> 0, not nan + np.testing.assert_allclose(out, np.array([[3.0], [0.0]]))