From 44f44acb2b9a56f55609ec7f0594d9b6fa01cb92 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:00:30 +0800 Subject: [PATCH 01/24] feat(dpmodel): add NeighborGraph/GraphLayout contract dataclasses --- deepmd/dpmodel/utils/neighbor_graph.py | 59 +++++++++++++++++++ .../common/dpmodel/test_neighbor_graph.py | 33 +++++++++++ 2 files changed, 92 insertions(+) create mode 100644 deepmd/dpmodel/utils/neighbor_graph.py create mode 100644 source/tests/common/dpmodel/test_neighbor_graph.py diff --git a/deepmd/dpmodel/utils/neighbor_graph.py b/deepmd/dpmodel/utils/neighbor_graph.py new file mode 100644 index 0000000000..fbc403e432 --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph.py @@ -0,0 +1,59 @@ +# 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 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 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..12407b087a --- /dev/null +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -0,0 +1,33 @@ +# 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 7526989aead60642b0780bc38fb43bee45aa361e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:01:36 +0800 Subject: [PATCH 02/24] feat(dpmodel): derive node validity mask from n_node (compact-prefix) --- deepmd/dpmodel/utils/neighbor_graph.py | 15 +++++++++++++++ .../common/dpmodel/test_neighbor_graph.py | 19 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/deepmd/dpmodel/utils/neighbor_graph.py b/deepmd/dpmodel/utils/neighbor_graph.py index fbc403e432..ac2766363e 100644 --- a/deepmd/dpmodel/utils/neighbor_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph.py @@ -16,6 +16,8 @@ dataclass, ) +import array_api_compat + from deepmd.dpmodel.array_api import ( Array, ) @@ -57,3 +59,16 @@ class GraphLayout: node_capacity: int | None = None frame_capacity: int | None = None min_edges: int = 2 + + +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/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index 12407b087a..004a818108 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -31,3 +31,22 @@ def test_graphlayout_defaults(self) -> None: 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 2b681a253e56c8c0ec6214ddedff1d33a0431314 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:02:44 +0800 Subject: [PATCH 03/24] feat(dpmodel): add mask-aware segment_sum/segment_mean reduction toolkit --- deepmd/dpmodel/utils/segment.py | 39 ++++++++++++++++++ source/tests/common/dpmodel/test_segment.py | 44 +++++++++++++++++++++ 2 files changed, 83 insertions(+) create mode 100644 deepmd/dpmodel/utils/segment.py create mode 100644 source/tests/common/dpmodel/test_segment.py diff --git a/deepmd/dpmodel/utils/segment.py b/deepmd/dpmodel/utils/segment.py new file mode 100644 index 0000000000..45d64af08c --- /dev/null +++ b/deepmd/dpmodel/utils/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_segment.py b/source/tests/common/dpmodel/test_segment.py new file mode 100644 index 0000000000..d73e7fa043 --- /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.segment 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]])) From f62cf5f30831fc7128c0095eca9450e391f7e1ee Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:04:00 +0800 Subject: [PATCH 04/24] feat(dpmodel): pad_and_guard_edges (compact layout, capacity, edge_mask) --- deepmd/dpmodel/utils/neighbor_graph.py | 37 +++++++++++++++++ .../common/dpmodel/test_neighbor_graph.py | 41 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/deepmd/dpmodel/utils/neighbor_graph.py b/deepmd/dpmodel/utils/neighbor_graph.py index ac2766363e..30ab57b425 100644 --- a/deepmd/dpmodel/utils/neighbor_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph.py @@ -61,6 +61,43 @@ class GraphLayout: 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. diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index 004a818108..785ca8aca1 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -50,3 +50,44 @@ def test_with_padding_prefix(self) -> None: 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 + ) From 608d379da423e5b6f051bc5467edc2b53116abe5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:05:36 +0800 Subject: [PATCH 05/24] feat(dpmodel): NumpyNeighborList reference builder (coords -> NeighborGraph) --- deepmd/dpmodel/utils/numpy_neighbor_list.py | 98 +++++++++++++++++++ .../dpmodel/test_numpy_neighbor_list.py | 91 +++++++++++++++++ 2 files changed, 189 insertions(+) create mode 100644 deepmd/dpmodel/utils/numpy_neighbor_list.py create mode 100644 source/tests/common/dpmodel/test_numpy_neighbor_list.py diff --git a/deepmd/dpmodel/utils/numpy_neighbor_list.py b/deepmd/dpmodel/utils/numpy_neighbor_list.py new file mode 100644 index 0000000000..3367df859b --- /dev/null +++ b/deepmd/dpmodel/utils/numpy_neighbor_list.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Numpy reference neighbor-list builder producing a NeighborGraph. + +All-pairs over periodic-image shells (over-enumeration filtered by rcut) - +correct but not optimized; the production O(N) builders (vesin / nvalchemiops) +live in the pt/pt_expt backends (later plan). Reference/test use only. +Limitation: shell count from the box diagonal (orthorhombic-ish test boxes). +""" + +import itertools + +import numpy as np + +from deepmd.dpmodel.utils.neighbor_graph import ( + GraphLayout, + NeighborGraph, + pad_and_guard_edges, +) + + +def _frame_edges(pos: np.ndarray, box: np.ndarray | None, rcut: float): + """Return (src_j, dst_i, edge_vec) for one frame. src=neighbor, dst=center.""" + nloc = pos.shape[0] + if box is None: + shells = [np.zeros(3, dtype=np.int64)] + else: + h = float(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) + ] + src, dst, vec = [], [], [] + rcut2 = rcut * rcut + for s in shells: + sc = np.zeros(3) if box is None else s.astype(np.float64) @ box + shifted = pos + sc # (nloc, 3) positions of image-shifted neighbors + for i in range(nloc): + d = shifted - pos[i] # (nloc, 3) + r2 = np.sum(d * d, axis=1) + for j in range(nloc): + if 1e-20 < r2[j] < rcut2: + src.append(j) + dst.append(i) + vec.append(d[j]) + if len(src) == 0: + return ( + np.zeros((0,), dtype=np.int64), + np.zeros((0,), dtype=np.int64), + np.zeros((0, 3), dtype=np.float64), + ) + return ( + np.array(src, dtype=np.int64), + np.array(dst, dtype=np.int64), + np.stack(vec, axis=0).astype(np.float64), + ) + + +class NumpyNeighborList: + """Reference builder: coords -> NeighborGraph.""" + + def build( + self, + coord: np.ndarray, + box: np.ndarray | None, + rcut: float, + layout: GraphLayout | None = None, + ) -> NeighborGraph: + if layout is None: + layout = GraphLayout() + coord = np.asarray(coord, dtype=np.float64) + nf, nloc = coord.shape[0], coord.shape[1] + n_node = np.full((nf,), nloc, dtype=np.int64) + src_all, dst_all, vec_all = [], [], [] + for f in range(nf): + bx = None if box is None else np.asarray(box, dtype=np.float64)[f] + src, dst, vec = _frame_edges(coord[f], bx, rcut) + offset = f * nloc + src_all.append(src + offset) + dst_all.append(dst + offset) + vec_all.append(vec) + src_cat = np.concatenate(src_all) if src_all else np.zeros((0,), np.int64) + dst_cat = np.concatenate(dst_all) if dst_all else np.zeros((0,), np.int64) + edge_index = np.stack([src_cat, dst_cat], axis=0).astype(np.int64) # (2, E) + edge_vec = ( + np.concatenate(vec_all, axis=0) + if vec_all + else np.zeros((0, 3), np.float64) + ) + edge_index, edge_vec, edge_mask = pad_and_guard_edges( + edge_index, edge_vec, layout.edge_capacity, layout.min_edges + ) + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + ) diff --git a/source/tests/common/dpmodel/test_numpy_neighbor_list.py b/source/tests/common/dpmodel/test_numpy_neighbor_list.py new file mode 100644 index 0000000000..63974f134c --- /dev/null +++ b/source/tests/common/dpmodel/test_numpy_neighbor_list.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np + +from deepmd.dpmodel.utils.neighbor_graph import ( + GraphLayout, +) +from deepmd.dpmodel.utils.numpy_neighbor_list import ( + NumpyNeighborList, +) + + +def brute_force_neighbor_sets(coord, box, rcut): + """Reference: per center i, the multiset of (src j, rounded edge_vec).""" + 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 + + +class TestNumpyNeighborList(unittest.TestCase): + def setUp(self) -> None: + self.rcut = 4.0 + self.coord = np.array( + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [3.5, 0.0, 0.0]], + dtype=np.float64, + ).reshape(1, 4, 3) + + def _graph_neighbor_sets(self, ng, nloc): + ei = ng.edge_index[:, ng.edge_mask] # real edges only + 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 test_nonperiodic_matches_brute_force(self) -> None: + ng = NumpyNeighborList().build(self.coord, box=None, rcut=self.rcut) + np.testing.assert_array_equal(ng.n_node, np.array([4], dtype=np.int64)) + got = self._graph_neighbor_sets(ng, 4) + want = brute_force_neighbor_sets(self.coord[0], None, self.rcut) + self.assertEqual(got, want) + + def test_periodic_matches_brute_force(self) -> None: + box = np.eye(3, dtype=np.float64)[None] * 6.0 # (1,3,3) cubic L=6 + ng = NumpyNeighborList().build(self.coord, box=box, rcut=self.rcut) + got = self._graph_neighbor_sets(ng, 4) + want = brute_force_neighbor_sets(self.coord[0], box[0], self.rcut) + self.assertEqual(got, want) + + def test_edge_vec_within_rcut(self) -> None: + ng = NumpyNeighborList().build(self.coord, box=None, rcut=self.rcut) + ev = ng.edge_vec[ng.edge_mask] + self.assertTrue(np.all(np.linalg.norm(ev, axis=1) < self.rcut)) + + def test_multiframe_offsets_nodes(self) -> None: + coord2 = np.concatenate([self.coord, self.coord], axis=0) # nf=2, nloc=4 + ng = NumpyNeighborList().build(coord2, box=None, rcut=self.rcut) + np.testing.assert_array_equal(ng.n_node, np.array([4, 4], dtype=np.int64)) + ei = ng.edge_index[:, ng.edge_mask] + # every frame-1 edge endpoint is in [4, 8); every frame-0 edge in [0, 4) + f0 = ei[:, ei[1] < 4] + f1 = ei[:, ei[1] >= 4] + self.assertTrue(np.all(f0 < 4)) + self.assertTrue(np.all(f1 >= 4)) + + def test_static_capacity_padding(self) -> None: + ng = NumpyNeighborList().build( + self.coord, box=None, rcut=self.rcut, layout=GraphLayout(edge_capacity=64) + ) + self.assertEqual(ng.edge_index.shape[1], 64) + self.assertEqual(ng.edge_vec.shape[0], 64) + self.assertEqual(int(ng.edge_mask.sum()), int(ng.edge_mask[:].sum())) + # masked-out tail contributes no real edges + self.assertTrue(np.all(ng.edge_vec[~ng.edge_mask] == 0.0)) From 0fb91f74e365e0e3ba57fe95f1e0b2a9d4c11984 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:07:13 +0800 Subject: [PATCH 06/24] feat(dpmodel): edge force/virial assembly (full-to-src atom virial) --- deepmd/dpmodel/utils/edge_derivatives.py | 44 +++++++++++++ .../common/dpmodel/test_edge_force_virial.py | 62 +++++++++++++++++++ 2 files changed, 106 insertions(+) create mode 100644 deepmd/dpmodel/utils/edge_derivatives.py create mode 100644 source/tests/common/dpmodel/test_edge_force_virial.py diff --git a/deepmd/dpmodel/utils/edge_derivatives.py b/deepmd/dpmodel/utils/edge_derivatives.py new file mode 100644 index 0000000000..215de4fdee --- /dev/null +++ b/deepmd/dpmodel/utils/edge_derivatives.py @@ -0,0 +1,44 @@ +# 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) + global virial = sum_e w_e +Padding/guard edges (edge_mask == 0) are zeroed before any scatter. +""" + +import array_api_compat + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.utils.segment import ( + segment_sum, +) + + +def edge_force_virial( + g_e: Array, + edge_vec: Array, + edge_index: Array, + edge_mask: Array, + n_node_total: int, +) -> tuple[Array, Array, Array]: + """Returns (force (N,3), atom_virial (N,3,3), global_virial (3,3)).""" + xp = array_api_compat.array_namespace(g_e) + g = g_e * edge_mask[:, None] # zero padding/guard contributions + src = edge_index[0] + dst = edge_index[1] + # force + force = segment_sum(g, dst, n_node_total) - segment_sum(g, src, n_node_total) + # 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_node_total) # (N, 3, 3) + global_virial = xp.sum(w_edge, axis=0) # (3, 3) + return force, atom_virial, global_virial 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..5eff325808 --- /dev/null +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.utils.edge_derivatives import ( + edge_force_virial, +) + + +class TestEdgeForceVirial(unittest.TestCase): + def setUp(self) -> None: + # 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 = 2 + + def test_force_formula(self) -> None: + force, _, _ = edge_force_virial( + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N + ) + # 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_global_virial_is_sum_of_edge_outer(self) -> None: + _, _, gv = edge_force_virial( + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N + ) + # W = -sum_e g_e (x) edge_vec_e + want = -( + np.einsum("k,j->kj", self.g[0], self.edge_vec[0]) + + np.einsum("k,j->kj", self.g[1], self.edge_vec[1]) + ) + np.testing.assert_allclose(gv, want) + + def test_atom_virial_full_to_src_sums_to_global(self) -> None: + _, av, gv = edge_force_virial( + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N + ) + self.assertEqual(av.shape, (2, 3, 3)) + np.testing.assert_allclose(np.sum(av, axis=0), gv) + # 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_padding_edges_contribute_nothing(self) -> None: + # append a masked guard edge pointing at node 0 with nonzero g (should be ignored) + ei = np.concatenate([self.edge_index, np.array([[0], [0]], dtype=np.int64)], axis=1) + ev = np.concatenate([self.edge_vec, np.array([[9.0, 9.0, 9.0]])], axis=0) + em = np.array([True, True, False]) + g = np.concatenate([self.g, np.array([[7.0, 7.0, 7.0]])], axis=0) + f1, a1, v1 = edge_force_virial(g, ev, ei, em, self.N) + f0, a0, v0 = edge_force_virial(self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N) + np.testing.assert_allclose(f1, f0) + np.testing.assert_allclose(a1, a0) + np.testing.assert_allclose(v1, v0) From 9c6cc80fed835a8dde7461267823b04f4524c165 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:11:29 +0800 Subject: [PATCH 07/24] feat(dpmodel): export NeighborGraph foundation symbols --- deepmd/dpmodel/utils/__init__.py | 24 +++++++++++++++++++ deepmd/dpmodel/utils/neighbor_graph.py | 14 ++++++----- deepmd/dpmodel/utils/numpy_neighbor_list.py | 8 +++---- .../common/dpmodel/test_neighbor_graph.py | 17 +++++++++++++ 4 files changed, 53 insertions(+), 10 deletions(-) diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index a9af7a50e5..c8c0442f35 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -2,6 +2,9 @@ from .default_neighbor_list import ( DefaultNeighborList, ) +from .edge_derivatives import ( + edge_force_virial, +) from .env_mat import ( EnvMat, ) @@ -18,9 +21,18 @@ is_lmdb, make_neighbor_stat_data, ) +from .neighbor_graph import ( + GraphLayout, + NeighborGraph, + node_validity_mask, + pad_and_guard_edges, +) from .neighbor_list import ( NeighborList, ) +from .numpy_neighbor_list import ( + NumpyNeighborList, +) from .network import ( EmbeddingNet, FittingNet, @@ -46,6 +58,10 @@ phys2inter, to_face_distance, ) +from .segment import ( + segment_mean, + segment_sum, +) from .serialization import ( load_dp_model, save_dp_model, @@ -60,6 +76,7 @@ __all__ = [ "AtomExcludeMask", "DefaultNeighborList", + "GraphLayout", "DistributedSameNlocBatchSampler", "EmbeddingNet", "EnvMat", @@ -69,12 +86,15 @@ "LmdbTestDataNlocView", "NativeLayer", "NativeNet", + "NeighborGraph", "NeighborList", "NetworkCollection", + "NumpyNeighborList", "PairExcludeMask", "SameNlocBatchSampler", "aggregate", "build_multiple_neighbor_list", + "edge_force_virial", "build_neighbor_list", "compute_total_numb_batch", "extend_coord_with_ghosts", @@ -88,11 +108,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.py b/deepmd/dpmodel/utils/neighbor_graph.py index 30ab57b425..232145bda0 100644 --- a/deepmd/dpmodel/utils/neighbor_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph.py @@ -15,12 +15,16 @@ from dataclasses import ( dataclass, ) +from typing import ( + TYPE_CHECKING, +) import array_api_compat -from deepmd.dpmodel.array_api import ( - Array, -) +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import ( + Array, + ) @dataclass @@ -105,7 +109,5 @@ def node_validity_mask(n_node: Array, n_total: int) -> Array: 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) - ) + 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/numpy_neighbor_list.py b/deepmd/dpmodel/utils/numpy_neighbor_list.py index 3367df859b..be1b6a9b57 100644 --- a/deepmd/dpmodel/utils/numpy_neighbor_list.py +++ b/deepmd/dpmodel/utils/numpy_neighbor_list.py @@ -18,7 +18,9 @@ ) -def _frame_edges(pos: np.ndarray, box: np.ndarray | None, rcut: float): +def _frame_edges( + pos: np.ndarray, box: np.ndarray | None, rcut: float +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Return (src_j, dst_i, edge_vec) for one frame. src=neighbor, dst=center.""" nloc = pos.shape[0] if box is None: @@ -83,9 +85,7 @@ def build( dst_cat = np.concatenate(dst_all) if dst_all else np.zeros((0,), np.int64) edge_index = np.stack([src_cat, dst_cat], axis=0).astype(np.int64) # (2, E) edge_vec = ( - np.concatenate(vec_all, axis=0) - if vec_all - else np.zeros((0, 3), np.float64) + np.concatenate(vec_all, axis=0) if vec_all else np.zeros((0, 3), np.float64) ) edge_index, edge_vec, edge_mask = pad_and_guard_edges( edge_index, edge_vec, layout.edge_capacity, layout.min_edges diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index 785ca8aca1..6183cc84cb 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -91,3 +91,20 @@ def test_overflow_raises(self) -> None: 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, + NumpyNeighborList, + edge_force_virial, + segment_sum, + ) + + self.assertTrue(callable(segment_sum)) + self.assertTrue(callable(edge_force_virial)) + self.assertIsNotNone(NeighborGraph) + self.assertIsNotNone(GraphLayout) + self.assertIsNotNone(NumpyNeighborList) From 341a4081e47575de14d2425a783ba7cc421b298a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 16:18:54 +0800 Subject: [PATCH 08/24] fix(dpmodel): sort __all__ (RUF022) + cast edge mask to g dtype Addresses two-stage review: RUF022 isort ordering of __all__ (CI blocker) and array-API-pure mask multiply (edge_force_virial) per CLAUDE.md guideline. --- deepmd/dpmodel/utils/__init__.py | 4 ++-- deepmd/dpmodel/utils/edge_derivatives.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index c8c0442f35..6c47dab613 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -76,11 +76,11 @@ __all__ = [ "AtomExcludeMask", "DefaultNeighborList", - "GraphLayout", "DistributedSameNlocBatchSampler", "EmbeddingNet", "EnvMat", "FittingNet", + "GraphLayout", "LmdbDataReader", "LmdbTestData", "LmdbTestDataNlocView", @@ -94,9 +94,9 @@ "SameNlocBatchSampler", "aggregate", "build_multiple_neighbor_list", - "edge_force_virial", "build_neighbor_list", "compute_total_numb_batch", + "edge_force_virial", "extend_coord_with_ghosts", "get_graph_index", "get_multiple_nlist_key", diff --git a/deepmd/dpmodel/utils/edge_derivatives.py b/deepmd/dpmodel/utils/edge_derivatives.py index 215de4fdee..bb94217c2c 100644 --- a/deepmd/dpmodel/utils/edge_derivatives.py +++ b/deepmd/dpmodel/utils/edge_derivatives.py @@ -31,7 +31,9 @@ def edge_force_virial( ) -> tuple[Array, Array, Array]: """Returns (force (N,3), atom_virial (N,3,3), global_virial (3,3)).""" xp = array_api_compat.array_namespace(g_e) - g = g_e * edge_mask[:, None] # zero padding/guard contributions + # 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 From a7e438249a18a6d3baf0ab31e8e726feacf9c1fa Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 24 Jun 2026 08:28:25 +0000 Subject: [PATCH 09/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/utils/__init__.py | 6 +++--- source/tests/common/dpmodel/test_edge_force_virial.py | 8 ++++++-- source/tests/common/dpmodel/test_neighbor_graph.py | 8 ++------ source/tests/common/dpmodel/test_numpy_neighbor_list.py | 5 ++++- 4 files changed, 15 insertions(+), 12 deletions(-) diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 6c47dab613..c3f11a2796 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -30,9 +30,6 @@ from .neighbor_list import ( NeighborList, ) -from .numpy_neighbor_list import ( - NumpyNeighborList, -) from .network import ( EmbeddingNet, FittingNet, @@ -52,6 +49,9 @@ get_multiple_nlist_key, nlist_distinguish_types, ) +from .numpy_neighbor_list import ( + NumpyNeighborList, +) from .region import ( inter2phys, normalize_coord, diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index 5eff325808..92d36437dc 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -51,12 +51,16 @@ def test_atom_virial_full_to_src_sums_to_global(self) -> None: def test_padding_edges_contribute_nothing(self) -> None: # append a masked guard edge pointing at node 0 with nonzero g (should be ignored) - ei = np.concatenate([self.edge_index, np.array([[0], [0]], dtype=np.int64)], axis=1) + ei = np.concatenate( + [self.edge_index, np.array([[0], [0]], dtype=np.int64)], axis=1 + ) ev = np.concatenate([self.edge_vec, np.array([[9.0, 9.0, 9.0]])], axis=0) em = np.array([True, True, False]) g = np.concatenate([self.g, np.array([[7.0, 7.0, 7.0]])], axis=0) f1, a1, v1 = edge_force_virial(g, ev, ei, em, self.N) - f0, a0, v0 = edge_force_virial(self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N) + f0, a0, v0 = edge_force_virial( + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N + ) np.testing.assert_allclose(f1, f0) np.testing.assert_allclose(a1, a0) np.testing.assert_allclose(v1, v0) diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index 6183cc84cb..20a27ef677 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -47,9 +47,7 @@ def test_no_padding_all_true(self) -> None: 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) - ) + np.testing.assert_array_equal(mask, np.array([True] * 5 + [False] * 3)) from deepmd.dpmodel.utils.neighbor_graph import ( @@ -88,9 +86,7 @@ def test_static_capacity_pads_to_E_max(self) -> None: 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 - ) + pad_and_guard_edges(self.edge_index, self.edge_vec, capacity=2, min_edges=2) class TestPublicExports(unittest.TestCase): diff --git a/source/tests/common/dpmodel/test_numpy_neighbor_list.py b/source/tests/common/dpmodel/test_numpy_neighbor_list.py index 63974f134c..ed1ed0c2b2 100644 --- a/source/tests/common/dpmodel/test_numpy_neighbor_list.py +++ b/source/tests/common/dpmodel/test_numpy_neighbor_list.py @@ -20,7 +20,10 @@ def brute_force_neighbor_sets(coord, box, rcut): 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)] + 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 From 52ca488a4bb777ba90626fe7a172094656f2006e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 19:16:35 +0800 Subject: [PATCH 10/24] refactor(dpmodel): build NeighborGraph from existing nlist; drop NumpyNeighborList Add neighbor_graph_from_extended (quartet -> ghost-free NeighborGraph, src = mapping[neighbor] local owner) and build_neighbor_graph (reuses the tested extend_input_and_build_neighbor_list). Drop the all-pairs NumpyNeighborList from the library; the independent brute-force oracle now lives in the test file (test_neighbor_graph_builder.py) and cross-validates the adapter. --- deepmd/dpmodel/utils/__init__.py | 5 +- deepmd/dpmodel/utils/neighbor_graph.py | 104 ++++++++++++ deepmd/dpmodel/utils/numpy_neighbor_list.py | 98 ----------- .../common/dpmodel/test_neighbor_graph.py | 6 +- .../dpmodel/test_neighbor_graph_builder.py | 157 ++++++++++++++++++ .../dpmodel/test_numpy_neighbor_list.py | 94 ----------- 6 files changed, 269 insertions(+), 195 deletions(-) delete mode 100644 deepmd/dpmodel/utils/numpy_neighbor_list.py create mode 100644 source/tests/common/dpmodel/test_neighbor_graph_builder.py delete mode 100644 source/tests/common/dpmodel/test_numpy_neighbor_list.py diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index c3f11a2796..26858cf369 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -24,6 +24,8 @@ from .neighbor_graph import ( GraphLayout, NeighborGraph, + build_neighbor_graph, + neighbor_graph_from_extended, node_validity_mask, pad_and_guard_edges, ) @@ -89,11 +91,11 @@ "NeighborGraph", "NeighborList", "NetworkCollection", - "NumpyNeighborList", "PairExcludeMask", "SameNlocBatchSampler", "aggregate", "build_multiple_neighbor_list", + "build_neighbor_graph", "build_neighbor_list", "compute_total_numb_batch", "edge_force_virial", @@ -107,6 +109,7 @@ "make_fitting_network", "make_multilayer_network", "make_neighbor_stat_data", + "neighbor_graph_from_extended", "nlist_distinguish_types", "node_validity_mask", "normalize_coord", diff --git a/deepmd/dpmodel/utils/neighbor_graph.py b/deepmd/dpmodel/utils/neighbor_graph.py index 232145bda0..526599d0a7 100644 --- a/deepmd/dpmodel/utils/neighbor_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph.py @@ -111,3 +111,107 @@ def node_validity_mask(n_node: Array, n_total: int) -> Array: 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) + + +def neighbor_graph_from_extended( + extended_coord: Array, + nlist: Array, + mapping: Array, + layout: GraphLayout | None = None, +) -> NeighborGraph: + """Convert the legacy extended quartet into a ghost-free NeighborGraph. + + This is the dpmodel/array-API adapter that REUSES deepmd's existing, tested, + general-cell neighbor list (``build_neighbor_list`` / ``extend_coord_with_ghosts``) + instead of re-deriving neighbors. 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 = nlist.shape[0] + nloc = nlist.shape[1] + nsel = nlist.shape[2] + n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) + src_parts: list[Array] = [] + dst_parts: list[Array] = [] + vec_parts: list[Array] = [] + center_full = xp.broadcast_to( + xp.reshape(xp.arange(nloc, dtype=xp.int64, device=dev), (nloc, 1)), + (nloc, nsel), + ) + center_flat = xp.reshape(center_full, (nloc * nsel,)) + for ff in range(nf): + nl_flat = xp.reshape(nlist[ff], (nloc * nsel,)) + keep = xp.reshape(xp.nonzero(nl_flat >= 0)[0], (-1,)) + j_ext = xp.take(nl_flat, keep, axis=0) # extended neighbor indices + dst = xp.take(center_flat, keep, axis=0) # local center indices + src = xp.take(mapping[ff], j_ext, axis=0) # local owner of neighbor + vec = xp.take(extended_coord[ff], j_ext, axis=0) - xp.take( + extended_coord[ff], dst, axis=0 + ) + offset = ff * nloc + src_parts.append(src + offset) + dst_parts.append(dst + offset) + vec_parts.append(vec) + edge_index = xp.astype( + xp.stack([xp.concat(src_parts), xp.concat(dst_parts)], axis=0), xp.int64 + ) + edge_vec = xp.concat(vec_parts, axis=0) + edge_index, edge_vec, edge_mask = pad_and_guard_edges( + edge_index, edge_vec, layout.edge_capacity, layout.min_edges + ) + 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, + sel: int | list[int], + mixed_types: bool = True, + layout: GraphLayout | None = None, +) -> NeighborGraph: + """Build a NeighborGraph by reusing the tested dense nlist (dpmodel default). + + Calls ``extend_input_and_build_neighbor_list`` (general-cell, tested) then + :func:`neighbor_graph_from_extended`. With ``sel`` large enough that no real + neighbor is truncated, the result is exactly the in-``rcut`` environment (the + ``sel``-as-normalization regime; see memory/spec_unified_edge_nlist.md). + """ + from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, + ) + + extended_coord, _extended_atype, mapping, nlist = ( + extend_input_and_build_neighbor_list( + coord, atype, rcut, sel, mixed_types=mixed_types, box=box + ) + ) + return neighbor_graph_from_extended(extended_coord, nlist, mapping, layout) diff --git a/deepmd/dpmodel/utils/numpy_neighbor_list.py b/deepmd/dpmodel/utils/numpy_neighbor_list.py deleted file mode 100644 index be1b6a9b57..0000000000 --- a/deepmd/dpmodel/utils/numpy_neighbor_list.py +++ /dev/null @@ -1,98 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Numpy reference neighbor-list builder producing a NeighborGraph. - -All-pairs over periodic-image shells (over-enumeration filtered by rcut) - -correct but not optimized; the production O(N) builders (vesin / nvalchemiops) -live in the pt/pt_expt backends (later plan). Reference/test use only. -Limitation: shell count from the box diagonal (orthorhombic-ish test boxes). -""" - -import itertools - -import numpy as np - -from deepmd.dpmodel.utils.neighbor_graph import ( - GraphLayout, - NeighborGraph, - pad_and_guard_edges, -) - - -def _frame_edges( - pos: np.ndarray, box: np.ndarray | None, rcut: float -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Return (src_j, dst_i, edge_vec) for one frame. src=neighbor, dst=center.""" - nloc = pos.shape[0] - if box is None: - shells = [np.zeros(3, dtype=np.int64)] - else: - h = float(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) - ] - src, dst, vec = [], [], [] - rcut2 = rcut * rcut - for s in shells: - sc = np.zeros(3) if box is None else s.astype(np.float64) @ box - shifted = pos + sc # (nloc, 3) positions of image-shifted neighbors - for i in range(nloc): - d = shifted - pos[i] # (nloc, 3) - r2 = np.sum(d * d, axis=1) - for j in range(nloc): - if 1e-20 < r2[j] < rcut2: - src.append(j) - dst.append(i) - vec.append(d[j]) - if len(src) == 0: - return ( - np.zeros((0,), dtype=np.int64), - np.zeros((0,), dtype=np.int64), - np.zeros((0, 3), dtype=np.float64), - ) - return ( - np.array(src, dtype=np.int64), - np.array(dst, dtype=np.int64), - np.stack(vec, axis=0).astype(np.float64), - ) - - -class NumpyNeighborList: - """Reference builder: coords -> NeighborGraph.""" - - def build( - self, - coord: np.ndarray, - box: np.ndarray | None, - rcut: float, - layout: GraphLayout | None = None, - ) -> NeighborGraph: - if layout is None: - layout = GraphLayout() - coord = np.asarray(coord, dtype=np.float64) - nf, nloc = coord.shape[0], coord.shape[1] - n_node = np.full((nf,), nloc, dtype=np.int64) - src_all, dst_all, vec_all = [], [], [] - for f in range(nf): - bx = None if box is None else np.asarray(box, dtype=np.float64)[f] - src, dst, vec = _frame_edges(coord[f], bx, rcut) - offset = f * nloc - src_all.append(src + offset) - dst_all.append(dst + offset) - vec_all.append(vec) - src_cat = np.concatenate(src_all) if src_all else np.zeros((0,), np.int64) - dst_cat = np.concatenate(dst_all) if dst_all else np.zeros((0,), np.int64) - edge_index = np.stack([src_cat, dst_cat], axis=0).astype(np.int64) # (2, E) - edge_vec = ( - np.concatenate(vec_all, axis=0) if vec_all else np.zeros((0, 3), np.float64) - ) - edge_index, edge_vec, edge_mask = pad_and_guard_edges( - edge_index, edge_vec, layout.edge_capacity, layout.min_edges - ) - return NeighborGraph( - n_node=n_node, - edge_index=edge_index, - edge_vec=edge_vec, - edge_mask=edge_mask, - ) diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index 20a27ef677..20730b0048 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -94,13 +94,15 @@ def test_importable_from_utils(self) -> None: from deepmd.dpmodel.utils import ( GraphLayout, NeighborGraph, - NumpyNeighborList, + build_neighbor_graph, edge_force_virial, + neighbor_graph_from_extended, segment_sum, ) self.assertTrue(callable(segment_sum)) self.assertTrue(callable(edge_force_virial)) + self.assertTrue(callable(build_neighbor_graph)) + self.assertTrue(callable(neighbor_graph_from_extended)) self.assertIsNotNone(NeighborGraph) self.assertIsNotNone(GraphLayout) - self.assertIsNotNone(NumpyNeighborList) 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..f25804723f --- /dev/null +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the dpmodel default NeighborGraph builder. + +``build_neighbor_graph`` reuses deepmd's tested extended neighbor list +(``extend_input_and_build_neighbor_list``) and converts it to a NeighborGraph +via ``neighbor_graph_from_extended``. We validate it against an INDEPENDENT +brute-force all-pairs oracle defined locally in this test file (kept here, not +in the library, because the production builder reuses the already-tested +extended nlist). +""" + +import itertools +import unittest + +import numpy as np + +from deepmd.dpmodel.utils.neighbor_graph import ( + GraphLayout, + build_neighbor_graph, + neighbor_graph_from_extended, +) + + +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 + + +class TestNeighborGraphBuilder(unittest.TestCase): + def setUp(self) -> None: + self.rcut = 4.0 + self.sel = [50, 50] # large -> no truncation (sel-as-normalization regime) + # 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, self.sel, mixed_types=True + ) + 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.sel, mixed_types=True + ) + self.assertEqual( + graph_neighbor_sets(ng, 4), + brute_force_neighbor_sets(self.coord[0], box[0], self.rcut), + ) + + def test_edge_vec_within_rcut(self) -> None: + ng = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, self.sel, mixed_types=True + ) + ev = ng.edge_vec[ng.edge_mask] + self.assertTrue(np.all(np.linalg.norm(ev, axis=1) < self.rcut)) + + def test_multiframe_offsets_nodes(self) -> None: + coord2 = np.concatenate([self.coord, self.coord], axis=0) + atype2 = np.concatenate([self.atype, self.atype], axis=0) + ng = build_neighbor_graph( + coord2, atype2, None, self.rcut, self.sel, mixed_types=True + ) + np.testing.assert_array_equal(ng.n_node, np.array([4, 4], dtype=np.int64)) + 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_static_capacity_padding(self) -> None: + ng = build_neighbor_graph( + self.coord, + self.atype, + None, + self.rcut, + self.sel, + mixed_types=True, + layout=GraphLayout(edge_capacity=64), + ) + self.assertEqual(ng.edge_index.shape[1], 64) + self.assertEqual(ng.edge_vec.shape[0], 64) + self.assertTrue(np.all(ng.edge_vec[~ng.edge_mask] == 0.0)) + + +class TestNeighborGraphFromExtended(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 = neighbor_graph_from_extended(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_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 = neighbor_graph_from_extended(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])) diff --git a/source/tests/common/dpmodel/test_numpy_neighbor_list.py b/source/tests/common/dpmodel/test_numpy_neighbor_list.py deleted file mode 100644 index ed1ed0c2b2..0000000000 --- a/source/tests/common/dpmodel/test_numpy_neighbor_list.py +++ /dev/null @@ -1,94 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import itertools -import unittest - -import numpy as np - -from deepmd.dpmodel.utils.neighbor_graph import ( - GraphLayout, -) -from deepmd.dpmodel.utils.numpy_neighbor_list import ( - NumpyNeighborList, -) - - -def brute_force_neighbor_sets(coord, box, rcut): - """Reference: per center i, the multiset of (src j, rounded edge_vec).""" - 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 - - -class TestNumpyNeighborList(unittest.TestCase): - def setUp(self) -> None: - self.rcut = 4.0 - self.coord = np.array( - [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [3.5, 0.0, 0.0]], - dtype=np.float64, - ).reshape(1, 4, 3) - - def _graph_neighbor_sets(self, ng, nloc): - ei = ng.edge_index[:, ng.edge_mask] # real edges only - 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 test_nonperiodic_matches_brute_force(self) -> None: - ng = NumpyNeighborList().build(self.coord, box=None, rcut=self.rcut) - np.testing.assert_array_equal(ng.n_node, np.array([4], dtype=np.int64)) - got = self._graph_neighbor_sets(ng, 4) - want = brute_force_neighbor_sets(self.coord[0], None, self.rcut) - self.assertEqual(got, want) - - def test_periodic_matches_brute_force(self) -> None: - box = np.eye(3, dtype=np.float64)[None] * 6.0 # (1,3,3) cubic L=6 - ng = NumpyNeighborList().build(self.coord, box=box, rcut=self.rcut) - got = self._graph_neighbor_sets(ng, 4) - want = brute_force_neighbor_sets(self.coord[0], box[0], self.rcut) - self.assertEqual(got, want) - - def test_edge_vec_within_rcut(self) -> None: - ng = NumpyNeighborList().build(self.coord, box=None, rcut=self.rcut) - ev = ng.edge_vec[ng.edge_mask] - self.assertTrue(np.all(np.linalg.norm(ev, axis=1) < self.rcut)) - - def test_multiframe_offsets_nodes(self) -> None: - coord2 = np.concatenate([self.coord, self.coord], axis=0) # nf=2, nloc=4 - ng = NumpyNeighborList().build(coord2, box=None, rcut=self.rcut) - np.testing.assert_array_equal(ng.n_node, np.array([4, 4], dtype=np.int64)) - ei = ng.edge_index[:, ng.edge_mask] - # every frame-1 edge endpoint is in [4, 8); every frame-0 edge in [0, 4) - f0 = ei[:, ei[1] < 4] - f1 = ei[:, ei[1] >= 4] - self.assertTrue(np.all(f0 < 4)) - self.assertTrue(np.all(f1 >= 4)) - - def test_static_capacity_padding(self) -> None: - ng = NumpyNeighborList().build( - self.coord, box=None, rcut=self.rcut, layout=GraphLayout(edge_capacity=64) - ) - self.assertEqual(ng.edge_index.shape[1], 64) - self.assertEqual(ng.edge_vec.shape[0], 64) - self.assertEqual(int(ng.edge_mask.sum()), int(ng.edge_mask[:].sum())) - # masked-out tail contributes no real edges - self.assertTrue(np.all(ng.edge_vec[~ng.edge_mask] == 0.0)) From 28a36275499812dceccdcca2a5698da3cdd51300 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 19:20:11 +0800 Subject: [PATCH 11/24] fix(dpmodel): drop stale numpy_neighbor_list import (rebase leftover) --- deepmd/dpmodel/utils/__init__.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 26858cf369..5c3821dc2b 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -51,9 +51,6 @@ get_multiple_nlist_key, nlist_distinguish_types, ) -from .numpy_neighbor_list import ( - NumpyNeighborList, -) from .region import ( inter2phys, normalize_coord, From 807787788e352a94548d6dda7cdf4f0432ac6e0e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 19:26:37 +0800 Subject: [PATCH 12/24] refactor(dpmodel): group NeighborGraph into a subpackage Move the flat utils modules (neighbor_graph.py, segment.py, edge_derivatives.py) into deepmd/dpmodel/utils/neighbor_graph/ as a cohesive subpackage: graph.py (contract + node-validity + edge padding), builder.py (neighbor_graph_from_extended, build_neighbor_graph), segment.py, derivatives.py. Public API re-exported from the package __init__ and from utils/__init__, so import paths (deepmd.dpmodel.utils.*) are unchanged. --- deepmd/dpmodel/utils/__init__.py | 10 +- .../dpmodel/utils/neighbor_graph/__init__.py | 39 ++++++ .../builder.py} | 106 ++-------------- .../derivatives.py} | 3 +- deepmd/dpmodel/utils/neighbor_graph/graph.py | 113 ++++++++++++++++++ .../utils/{ => neighbor_graph}/segment.py | 0 .../common/dpmodel/test_edge_force_virial.py | 2 +- source/tests/common/dpmodel/test_segment.py | 2 +- 8 files changed, 170 insertions(+), 105 deletions(-) create mode 100644 deepmd/dpmodel/utils/neighbor_graph/__init__.py rename deepmd/dpmodel/utils/{neighbor_graph.py => neighbor_graph/builder.py} (53%) rename deepmd/dpmodel/utils/{edge_derivatives.py => neighbor_graph/derivatives.py} (97%) create mode 100644 deepmd/dpmodel/utils/neighbor_graph/graph.py rename deepmd/dpmodel/utils/{ => neighbor_graph}/segment.py (100%) diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 5c3821dc2b..6eaa326129 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -2,9 +2,6 @@ from .default_neighbor_list import ( DefaultNeighborList, ) -from .edge_derivatives import ( - edge_force_virial, -) from .env_mat import ( EnvMat, ) @@ -25,9 +22,12 @@ GraphLayout, NeighborGraph, build_neighbor_graph, + edge_force_virial, neighbor_graph_from_extended, node_validity_mask, pad_and_guard_edges, + segment_mean, + segment_sum, ) from .neighbor_list import ( NeighborList, @@ -57,10 +57,6 @@ phys2inter, to_face_distance, ) -from .segment import ( - segment_mean, - segment_sum, -) from .serialization import ( load_dp_model, save_dp_model, diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py new file mode 100644 index 0000000000..ee4394322f --- /dev/null +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -0,0 +1,39 @@ +# 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`` (reuse the tested dense nlist -> ``NeighborGraph``), +``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, + neighbor_graph_from_extended, +) +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", + "neighbor_graph_from_extended", + "node_validity_mask", + "pad_and_guard_edges", + "segment_mean", + "segment_sum", +] diff --git a/deepmd/dpmodel/utils/neighbor_graph.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py similarity index 53% rename from deepmd/dpmodel/utils/neighbor_graph.py rename to deepmd/dpmodel/utils/neighbor_graph/builder.py index 526599d0a7..04d63f68df 100644 --- a/deepmd/dpmodel/utils/neighbor_graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -1,118 +1,34 @@ # 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. +"""Builders that produce a :class:`NeighborGraph`. -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. +``neighbor_graph_from_extended`` converts the legacy extended quartet +(extended_coord, nlist, mapping) into a ghost-free NeighborGraph; +``build_neighbor_graph`` is the dpmodel default that reuses deepmd's tested +``extend_input_and_build_neighbor_list`` and then calls the adapter. """ from __future__ import ( annotations, ) -from dataclasses import ( - dataclass, -) 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, ) -@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) - - def neighbor_graph_from_extended( extended_coord: Array, nlist: Array, diff --git a/deepmd/dpmodel/utils/edge_derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py similarity index 97% rename from deepmd/dpmodel/utils/edge_derivatives.py rename to deepmd/dpmodel/utils/neighbor_graph/derivatives.py index bb94217c2c..369c719861 100644 --- a/deepmd/dpmodel/utils/edge_derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -17,7 +17,8 @@ from deepmd.dpmodel.array_api import ( Array, ) -from deepmd.dpmodel.utils.segment import ( + +from .segment import ( segment_sum, ) 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/segment.py b/deepmd/dpmodel/utils/neighbor_graph/segment.py similarity index 100% rename from deepmd/dpmodel/utils/segment.py rename to deepmd/dpmodel/utils/neighbor_graph/segment.py diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index 92d36437dc..c271ee0e92 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -3,7 +3,7 @@ import numpy as np -from deepmd.dpmodel.utils.edge_derivatives import ( +from deepmd.dpmodel.utils.neighbor_graph import ( edge_force_virial, ) diff --git a/source/tests/common/dpmodel/test_segment.py b/source/tests/common/dpmodel/test_segment.py index d73e7fa043..22046911e9 100644 --- a/source/tests/common/dpmodel/test_segment.py +++ b/source/tests/common/dpmodel/test_segment.py @@ -3,7 +3,7 @@ import numpy as np -from deepmd.dpmodel.utils.segment import ( +from deepmd.dpmodel.utils.neighbor_graph import ( segment_mean, segment_sum, ) From 2d9fd985683aa588218033b536f87f42c4b309b2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 24 Jun 2026 23:37:18 +0800 Subject: [PATCH 13/24] fix(dpmodel): honor int sel in build_neighbor_graph; strengthen padding test Address CodeRabbit review on #5581: - build_neighbor_graph normalizes an int sel to list form before calling extend_input_and_build_neighbor_list (annotated list[int]), so the public int | list[int] contract is honored. No behavior change -- the underlying build_neighbor_list already accepts int -- but the call is now type-consistent. Add a test that int sel reproduces the list-sel real-edge environment. - test_static_capacity_padding: replace the (already-removed) tautological mask check's intent with explicit assertions on the exact real-edge count and the compact-prefix padding (real edges first, masked tail after). --- .../dpmodel/utils/neighbor_graph/builder.py | 6 ++++- .../dpmodel/test_neighbor_graph_builder.py | 22 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index 04d63f68df..ae16886755 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -125,9 +125,13 @@ def build_neighbor_graph( extend_input_and_build_neighbor_list, ) + # ``extend_input_and_build_neighbor_list`` is annotated ``sel: list[int]``; + # normalize the integer form so the public ``int | list[int]`` contract is + # honored (the underlying ``build_neighbor_list`` accepts both). + sel_list = [sel] if isinstance(sel, int) else sel extended_coord, _extended_atype, mapping, nlist = ( extend_input_and_build_neighbor_list( - coord, atype, rcut, sel, mixed_types=mixed_types, box=box + coord, atype, rcut, sel_list, mixed_types=mixed_types, box=box ) ) return neighbor_graph_from_extended(extended_coord, nlist, mapping, layout) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index f25804723f..e0223f2ef2 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -109,6 +109,20 @@ def test_multiframe_offsets_nodes(self) -> None: self.assertTrue(np.all(ei[:, ei[1] < 4] < 4)) self.assertTrue(np.all(ei[:, ei[1] >= 4] >= 4)) + def test_int_sel_matches_list_sel(self) -> None: + # an integer ``sel`` (normalized to list form) must yield the same + # real-edge environment as the equivalent large list ``sel``. + nloc = self.coord.shape[1] + ng_int = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, 64, mixed_types=True + ) + ng_list = build_neighbor_graph( + self.coord, self.atype, None, self.rcut, self.sel, mixed_types=True + ) + self.assertEqual( + graph_neighbor_sets(ng_int, nloc), graph_neighbor_sets(ng_list, nloc) + ) + def test_static_capacity_padding(self) -> None: ng = build_neighbor_graph( self.coord, @@ -121,6 +135,14 @@ def test_static_capacity_padding(self) -> None: ) 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)) From a2c5da8574021771c9dc7a8737b63630bc3927c2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 00:15:20 +0800 Subject: [PATCH 14/24] fix(dpmodel): per-frame virial in edge_force_virial; clarify legacy sel contract Address PR #5581 review (iProzd + chatgpt-codex): - edge_force_virial now returns PER-FRAME virial (nf,3,3) instead of a single (3,3) that collapsed all frames; takes n_node (per-frame counts), assigns each edge to its dst's frame via searchsorted(cumsum(n_node)). Adds a multi-frame regression test (the bug was invisible to the single-frame UT). - builder.py: document the LEGACY-COMPATIBLE / sel-truncating contract of neighbor_graph_from_extended / build_neighbor_graph explicitly; the carry-all graph-native builder is separate (from_ijs via ASE/vesin). --- .../dpmodel/utils/neighbor_graph/builder.py | 34 +++++++-- .../utils/neighbor_graph/derivatives.py | 42 +++++++++-- .../common/dpmodel/test_edge_force_virial.py | 73 +++++++++++++++---- 3 files changed, 119 insertions(+), 30 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index ae16886755..8907735315 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -1,10 +1,21 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Builders that produce a :class:`NeighborGraph`. -``neighbor_graph_from_extended`` converts the legacy extended quartet -(extended_coord, nlist, mapping) into a ghost-free NeighborGraph; -``build_neighbor_graph`` is the dpmodel default that reuses deepmd's tested -``extend_input_and_build_neighbor_list`` and then calls the adapter. +CONTRACT — these are the **legacy-compatible** builders: the graph they return +carries exactly the neighbors of the dense nlist they are built from, i.e. it +INHERITS that nlist's ``sel`` selection/truncation. They are the backward-compat +adapter (and the test oracle), NOT the graph-native "all neighbors within +``rcut``" builder. The complete-environment (carry-all, ``sel``-as-normalization) +contract is provided by a SEPARATE builder (``from_ijs`` fed by ASE/vesin; see +memory/spec_unified_edge_nlist.md decision #17). Keeping the two contracts in +distinct functions avoids the footgun of a consumer assuming it sees all +neighbors while a builder silently truncated them. + +- ``neighbor_graph_from_extended``: convert an existing extended quartet + (extended_coord, nlist, mapping) -> ghost-free graph (inherits the nlist's sel). +- ``build_neighbor_graph``: reuse deepmd's tested + ``extend_input_and_build_neighbor_list`` (which TRUNCATES to ``sel``) then the + adapter -> legacy-compatible graph. """ from __future__ import ( @@ -114,12 +125,19 @@ def build_neighbor_graph( mixed_types: bool = True, layout: GraphLayout | None = None, ) -> NeighborGraph: - """Build a NeighborGraph by reusing the tested dense nlist (dpmodel default). + """Build a LEGACY-COMPATIBLE NeighborGraph by reusing the tested dense nlist. Calls ``extend_input_and_build_neighbor_list`` (general-cell, tested) then - :func:`neighbor_graph_from_extended`. With ``sel`` large enough that no real - neighbor is truncated, the result is exactly the in-``rcut`` environment (the - ``sel``-as-normalization regime; see memory/spec_unified_edge_nlist.md). + :func:`neighbor_graph_from_extended`. + + CONTRACT: the returned graph contains the neighbors selected by ``sel`` and + **inherits the legacy ``sel`` truncation** — it does NOT carry all neighbors + within ``rcut`` when ``sel`` binds. It coincides with the complete in-``rcut`` + environment ONLY when ``sel`` is large enough that no real neighbor is dropped + (the ``sel``-as-normalization regime). For the carry-all graph-native contract + use the dedicated carry-all builder (``from_ijs`` via ASE/vesin), NOT this + function. See memory/spec_unified_edge_nlist.md (decision #17). This builder is + the backward-compat adapter and the test oracle. """ from deepmd.dpmodel.utils.nlist import ( extend_input_and_build_neighbor_list, diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index 369c719861..a1c6e35b64 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -8,7 +8,8 @@ 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) - global virial = sum_e w_e + 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. """ @@ -28,20 +29,47 @@ def edge_force_virial( edge_vec: Array, edge_index: Array, edge_mask: Array, - n_node_total: int, + n_node: Array, ) -> tuple[Array, Array, Array]: - """Returns (force (N,3), atom_virial (N,3,3), global_virial (3,3)).""" + """Assemble per-node force/atom-virial and PER-FRAME virial from ``g_e``. + + Parameters + ---------- + n_node + (nf,) per-frame node counts. The flat node axis is ``N = sum(n_node)`` + (compact, frame-major); ``nf = n_node.shape[0]``. + + 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_total = int(xp.sum(n_node)) # flat node count N (static shape) + 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 - force = segment_sum(g, dst, n_node_total) - segment_sum(g, src, n_node_total) + force = segment_sum(g, dst, n_total) - segment_sum(g, src, n_total) # 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_node_total) # (N, 3, 3) - global_virial = xp.sum(w_edge, axis=0) # (3, 3) - return force, atom_virial, global_virial + atom_virial = segment_sum(w_edge, src, n_total) # (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 + # the node axis is 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/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index c271ee0e92..64bd8968d8 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -10,39 +10,41 @@ class TestEdgeForceVirial(unittest.TestCase): def setUp(self) -> None: - # 2 nodes, 2 real edges: e0 = (src=1, dst=0), e1 = (src=0, dst=1) + # 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 = 2 + 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 + 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_global_virial_is_sum_of_edge_outer(self) -> None: - _, _, gv = edge_force_virial( - self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N + def test_virial_is_per_frame(self) -> None: + _, _, vir = edge_force_virial( + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.n_node ) - # W = -sum_e g_e (x) edge_vec_e + # single frame => shape (1, 3, 3); value = -sum_e g_e (x) edge_vec_e + self.assertEqual(vir.shape, (1, 3, 3)) want = -( np.einsum("k,j->kj", self.g[0], self.edge_vec[0]) + np.einsum("k,j->kj", self.g[1], self.edge_vec[1]) ) - np.testing.assert_allclose(gv, want) + np.testing.assert_allclose(vir[0], want) - def test_atom_virial_full_to_src_sums_to_global(self) -> None: - _, av, gv = edge_force_virial( - self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N + 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)) - np.testing.assert_allclose(np.sum(av, axis=0), gv) + # 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]) @@ -50,17 +52,58 @@ def test_atom_virial_full_to_src_sums_to_global(self) -> None: np.testing.assert_allclose(av[0], w1) # src of e1 is node 0 def test_padding_edges_contribute_nothing(self) -> None: - # append a masked guard edge pointing at node 0 with nonzero g (should be ignored) + # append a masked guard edge pointing at node 0 with nonzero g (ignored) ei = np.concatenate( [self.edge_index, np.array([[0], [0]], dtype=np.int64)], axis=1 ) ev = np.concatenate([self.edge_vec, np.array([[9.0, 9.0, 9.0]])], axis=0) em = np.array([True, True, False]) g = np.concatenate([self.g, np.array([[7.0, 7.0, 7.0]])], axis=0) - f1, a1, v1 = edge_force_virial(g, ev, ei, em, self.N) + f1, a1, v1 = edge_force_virial(g, ev, ei, em, self.n_node) f0, a0, v0 = edge_force_virial( - self.g, self.edge_vec, self.edge_index, self.edge_mask, self.N + self.g, self.edge_vec, self.edge_index, self.edge_mask, self.n_node ) np.testing.assert_allclose(f1, f0) np.testing.assert_allclose(a1, a0) np.testing.assert_allclose(v1, v0) + + def test_multiframe_virials_not_collapsed(self) -> None: + # 2 frames, 2 nodes each: nodes {0,1} in frame 0, {2,3} in frame 1. + # Frame edges chosen so the two frames have DISTINCT virials; the bug + # (summing all edges into one (3,3)) would merge them. + n_node = np.array([2, 2], dtype=np.int64) # N = 4 + # frame 0 edges: e0=(src=1,dst=0), e1=(src=0,dst=1); frame 1: e2=(3,2), e3=(2,3) + edge_index = np.array([[1, 0, 3, 2], [0, 1, 2, 3]], dtype=np.int64) + edge_vec = np.array( + [ + [1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0], + [0.0, 2.0, 0.0], # frame 1 uses a different direction/scale + [0.0, -2.0, 0.0], + ] + ) + edge_mask = np.array([True, True, True, True]) + g = np.array( + [ + [0.5, 0.0, 0.0], + [0.3, 0.0, 0.0], + [0.0, 0.7, 0.0], + [0.0, 0.1, 0.0], + ] + ) + force, av, vir = edge_force_virial(g, edge_vec, edge_index, edge_mask, n_node) + self.assertEqual(vir.shape, (2, 3, 3)) # per-frame, NOT collapsed + w = [-np.einsum("k,j->kj", g[i], edge_vec[i]) for i in range(4)] + np.testing.assert_allclose(vir[0], w[0] + w[1]) # frame 0 edges only + np.testing.assert_allclose(vir[1], w[2] + w[3]) # frame 1 edges only + # the two frames are genuinely different (would be equal-ish if merged-then-split wrong) + self.assertFalse(np.allclose(vir[0], vir[1])) + # per-frame atom-virial closure: frame-f nodes' atom-virials sum to vir[f] + np.testing.assert_allclose(np.sum(av[0:2], axis=0), vir[0]) + np.testing.assert_allclose(np.sum(av[2:4], axis=0), vir[1]) + # force is per-node (flat across frames), unaffected + self.assertEqual(force.shape, (4, 3)) + + +if __name__ == "__main__": + unittest.main() From 04d49fb16a0f151e561217a382c4f2fd6e918754 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 00:20:06 +0800 Subject: [PATCH 15/24] test(dpmodel): general multi-frame coverage (ragged nodes/edges + padding) Address review follow-up: the per-frame-virial bug slipped through because the only test was single-frame. Generalize coverage: - edge_force_virial: add node_capacity for a padded node axis; new test exercises the MOST GENERAL layout (ragged n_node=[3,5], uneven per-frame edges, masked guard edges, padded node axis) + per-frame virial/atom-virial closure. - builder: strengthen multi-frame test to TWO DIFFERENT frames checked per-frame vs the brute-force oracle (non-periodic + periodic); add multi-frame adapter test. (Builder is rectangular-node by construction; ragged-node building is a future ragged builder.) --- .../utils/neighbor_graph/derivatives.py | 26 ++++-- .../common/dpmodel/test_edge_force_virial.py | 73 ++++++++++++++++ .../dpmodel/test_neighbor_graph_builder.py | 85 ++++++++++++++++++- 3 files changed, 175 insertions(+), 9 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py index a1c6e35b64..1c0bafc234 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/derivatives.py +++ b/deepmd/dpmodel/utils/neighbor_graph/derivatives.py @@ -30,14 +30,25 @@ def edge_force_virial( 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 node counts. The flat node axis is ``N = sum(n_node)`` - (compact, frame-major); ``nf = n_node.shape[0]``. + (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 ------- @@ -51,22 +62,23 @@ def edge_force_virial( frame via the frame of their ``dst`` node. """ xp = array_api_compat.array_namespace(g_e) - n_total = int(xp.sum(n_node)) # flat node count N (static shape) + 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 - force = segment_sum(g, dst, n_total) - segment_sum(g, src, n_total) + # 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_total) # (N, 3, 3) + 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 - # the node axis is compact frame-major (frame f owns a contiguous block). + # 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 diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index 64bd8968d8..fa35d4497d 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -67,6 +67,79 @@ def test_padding_edges_contribute_nothing(self) -> None: np.testing.assert_allclose(a1, a0) np.testing.assert_allclose(v1, v0) + 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) + def test_multiframe_virials_not_collapsed(self) -> None: # 2 frames, 2 nodes each: nodes {0,1} in frame 0, {2,3} in frame 1. # Frame edges chosen so the two frames have DISTINCT virials; the bug diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index e0223f2ef2..63ed682e14 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -58,6 +58,24 @@ def graph_neighbor_sets(ng, nloc): 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 @@ -98,17 +116,54 @@ def test_edge_vec_within_rcut(self) -> None: ev = ng.edge_vec[ng.edge_mask] self.assertTrue(np.all(np.linalg.norm(ev, axis=1) < self.rcut)) - def test_multiframe_offsets_nodes(self) -> None: - coord2 = np.concatenate([self.coord, self.coord], axis=0) + 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, self.sel, mixed_types=True ) 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, self.sel, mixed_types=True + ) + 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_int_sel_matches_list_sel(self) -> None: # an integer ``sel`` (normalized to list form) must yield the same # real-edge environment as the equivalent large list ``sel``. @@ -165,6 +220,32 @@ def test_adapter_on_handmade_quartet(self) -> None: } 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 = neighbor_graph_from_extended(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) From 675a781588a10b359bf94ad1476e145291383f98 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 00:22:44 +0800 Subject: [PATCH 16/24] test(dpmodel): drop redundant single-frame/square-multiframe edge_force_virial tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ragged+padded general test subsumes the square multi-frame and the single-frame per-frame/edge-padding cases. Keep two minimal convention-pinning tests the general test cannot assert: force dst-src sign, and full-to-src attribution (decision #5 — the general closure is attribution-agnostic). --- .../common/dpmodel/test_edge_force_virial.py | 65 ------------------- 1 file changed, 65 deletions(-) diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index fa35d4497d..5e54901d64 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -26,18 +26,6 @@ def test_force_formula(self) -> None: # 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_virial_is_per_frame(self) -> None: - _, _, vir = edge_force_virial( - self.g, self.edge_vec, self.edge_index, self.edge_mask, self.n_node - ) - # single frame => shape (1, 3, 3); value = -sum_e g_e (x) edge_vec_e - self.assertEqual(vir.shape, (1, 3, 3)) - want = -( - np.einsum("k,j->kj", self.g[0], self.edge_vec[0]) - + np.einsum("k,j->kj", self.g[1], self.edge_vec[1]) - ) - np.testing.assert_allclose(vir[0], want) - 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 @@ -51,22 +39,6 @@ def test_atom_virial_full_to_src_sums_to_frame_virial(self) -> None: 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_padding_edges_contribute_nothing(self) -> None: - # append a masked guard edge pointing at node 0 with nonzero g (ignored) - ei = np.concatenate( - [self.edge_index, np.array([[0], [0]], dtype=np.int64)], axis=1 - ) - ev = np.concatenate([self.edge_vec, np.array([[9.0, 9.0, 9.0]])], axis=0) - em = np.array([True, True, False]) - g = np.concatenate([self.g, np.array([[7.0, 7.0, 7.0]])], axis=0) - f1, a1, v1 = edge_force_virial(g, ev, ei, em, self.n_node) - f0, a0, v0 = edge_force_virial( - self.g, self.edge_vec, self.edge_index, self.edge_mask, self.n_node - ) - np.testing.assert_allclose(f1, f0) - np.testing.assert_allclose(a1, a0) - np.testing.assert_allclose(v1, v0) - 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 @@ -140,43 +112,6 @@ def test_ragged_multiframe_with_edge_and_node_padding(self) -> None: np.testing.assert_allclose(av, a2) np.testing.assert_allclose(vir, v2) - def test_multiframe_virials_not_collapsed(self) -> None: - # 2 frames, 2 nodes each: nodes {0,1} in frame 0, {2,3} in frame 1. - # Frame edges chosen so the two frames have DISTINCT virials; the bug - # (summing all edges into one (3,3)) would merge them. - n_node = np.array([2, 2], dtype=np.int64) # N = 4 - # frame 0 edges: e0=(src=1,dst=0), e1=(src=0,dst=1); frame 1: e2=(3,2), e3=(2,3) - edge_index = np.array([[1, 0, 3, 2], [0, 1, 2, 3]], dtype=np.int64) - edge_vec = np.array( - [ - [1.0, 0.0, 0.0], - [-1.0, 0.0, 0.0], - [0.0, 2.0, 0.0], # frame 1 uses a different direction/scale - [0.0, -2.0, 0.0], - ] - ) - edge_mask = np.array([True, True, True, True]) - g = np.array( - [ - [0.5, 0.0, 0.0], - [0.3, 0.0, 0.0], - [0.0, 0.7, 0.0], - [0.0, 0.1, 0.0], - ] - ) - force, av, vir = edge_force_virial(g, edge_vec, edge_index, edge_mask, n_node) - self.assertEqual(vir.shape, (2, 3, 3)) # per-frame, NOT collapsed - w = [-np.einsum("k,j->kj", g[i], edge_vec[i]) for i in range(4)] - np.testing.assert_allclose(vir[0], w[0] + w[1]) # frame 0 edges only - np.testing.assert_allclose(vir[1], w[2] + w[3]) # frame 1 edges only - # the two frames are genuinely different (would be equal-ish if merged-then-split wrong) - self.assertFalse(np.allclose(vir[0], vir[1])) - # per-frame atom-virial closure: frame-f nodes' atom-virials sum to vir[f] - np.testing.assert_allclose(np.sum(av[0:2], axis=0), vir[0]) - np.testing.assert_allclose(np.sum(av[2:4], axis=0), vir[1]) - # force is per-node (flat across frames), unaffected - self.assertEqual(force.shape, (4, 3)) - if __name__ == "__main__": unittest.main() From a64135548ab76fb3add24bd22be48accc2deb8b2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 00:27:04 +0800 Subject: [PATCH 17/24] test(dpmodel): cover empty frame (n_node has a 0) in edge_force_virial n_node=[3,0,5] with an empty middle frame: verifies searchsorted frame assignment on duplicate cumsum boundaries [3,3,8] skips the zero-width block (node 3 -> frame 2) and the empty frame's virial is exactly zero. Combined with a padded node axis (node_capacity 9 > sum 8) to also exercise node padding at the primitive level (the only node_capacity consumer in this PR; a model-level caller arrives with the graph lower in PR-A/jax in PR-F). --- .../common/dpmodel/test_edge_force_virial.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index 5e54901d64..d34ac434a5 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -39,6 +39,42 @@ def test_atom_virial_full_to_src_sums_to_frame_virial(self) -> None: 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_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 From dd76bd601550d8e4a72922d3bddcc447929dc1b6 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 00:31:56 +0800 Subject: [PATCH 18/24] test(dpmodel): cover all-empty graph (zero real edges) in edge_force_virial Single- and multi-frame systems with nodes but NO real edges (only masked guard edges): force/atom_virial/per-frame virial must all be exactly zero with correct shapes. Covers isolated-atom / rcut-below-all-distances geometries. --- .../common/dpmodel/test_edge_force_virial.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/source/tests/common/dpmodel/test_edge_force_virial.py b/source/tests/common/dpmodel/test_edge_force_virial.py index d34ac434a5..fa84ef7ba4 100644 --- a/source/tests/common/dpmodel/test_edge_force_virial.py +++ b/source/tests/common/dpmodel/test_edge_force_virial.py @@ -75,6 +75,28 @@ def test_empty_frame_no_nodes_or_edges(self) -> None: # 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 From 014a9d5fe6a43cac6e6113e8e93c12763a7bef51 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 08:47:50 +0800 Subject: [PATCH 19/24] feat(dpmodel): make build_neighbor_graph a carry-all search; rename converter to from_dense_quartet build_neighbor_graph now SEARCHES neighbors directly from coordinates and keeps every neighbor within rcut (sel-free), instead of routing through the sel-truncating dense nlist. The legacy quartet adapter is renamed neighbor_graph_from_extended -> from_dense_quartet and documented as the backward-compat converter (inherits sel truncation), separating the 'compute from geometry' dispatcher from the 'adapt an existing list' converters (spec decision #17). --- deepmd/dpmodel/utils/__init__.py | 4 +- .../dpmodel/utils/neighbor_graph/__init__.py | 11 +- .../dpmodel/utils/neighbor_graph/builder.py | 193 +++++++++++++----- .../common/dpmodel/test_neighbor_graph.py | 4 +- .../dpmodel/test_neighbor_graph_builder.py | 93 +++++---- 5 files changed, 202 insertions(+), 103 deletions(-) diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 6eaa326129..0179543dd4 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -23,7 +23,7 @@ NeighborGraph, build_neighbor_graph, edge_force_virial, - neighbor_graph_from_extended, + from_dense_quartet, node_validity_mask, pad_and_guard_edges, segment_mean, @@ -93,6 +93,7 @@ "compute_total_numb_batch", "edge_force_virial", "extend_coord_with_ghosts", + "from_dense_quartet", "get_graph_index", "get_multiple_nlist_key", "inter2phys", @@ -102,7 +103,6 @@ "make_fitting_network", "make_multilayer_network", "make_neighbor_stat_data", - "neighbor_graph_from_extended", "nlist_distinguish_types", "node_validity_mask", "normalize_coord", diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index ee4394322f..08b165f861 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -3,14 +3,15 @@ The unified edge/graph neighbor-list contract and its supporting machinery: ``graph`` (the ``NeighborGraph``/``GraphLayout`` contract + derived node-validity -+ edge padding), ``builder`` (reuse the tested dense nlist -> ``NeighborGraph``), -``segment`` (mask-aware segment-reduction toolkit), and ``derivatives`` (edge -force/virial assembly). See memory/spec_unified_edge_nlist.md. ++ 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, - neighbor_graph_from_extended, + from_dense_quartet, ) from .derivatives import ( edge_force_virial, @@ -31,7 +32,7 @@ "NeighborGraph", "build_neighbor_graph", "edge_force_virial", - "neighbor_graph_from_extended", + "from_dense_quartet", "node_validity_mask", "pad_and_guard_edges", "segment_mean", diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index 8907735315..d680385818 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -1,21 +1,27 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Builders that produce a :class:`NeighborGraph`. - -CONTRACT — these are the **legacy-compatible** builders: the graph they return -carries exactly the neighbors of the dense nlist they are built from, i.e. it -INHERITS that nlist's ``sel`` selection/truncation. They are the backward-compat -adapter (and the test oracle), NOT the graph-native "all neighbors within -``rcut``" builder. The complete-environment (carry-all, ``sel``-as-normalization) -contract is provided by a SEPARATE builder (``from_ijs`` fed by ASE/vesin; see -memory/spec_unified_edge_nlist.md decision #17). Keeping the two contracts in -distinct functions avoids the footgun of a consumer assuming it sees all -neighbors while a builder silently truncated them. - -- ``neighbor_graph_from_extended``: convert an existing extended quartet - (extended_coord, nlist, mapping) -> ghost-free graph (inherits the nlist's sel). -- ``build_neighbor_graph``: reuse deepmd's tested - ``extend_input_and_build_neighbor_list`` (which TRUNCATES to ``sel``) then the - adapter -> legacy-compatible graph. +"""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?". """ from __future__ import ( @@ -40,22 +46,25 @@ ) -def neighbor_graph_from_extended( +def from_dense_quartet( extended_coord: Array, nlist: Array, mapping: Array, layout: GraphLayout | None = None, ) -> NeighborGraph: - """Convert the legacy extended quartet into a ghost-free NeighborGraph. + """Convert a legacy extended quartet into a ghost-free NeighborGraph (CONVERTER). - This is the dpmodel/array-API adapter that REUSES deepmd's existing, tested, - general-cell neighbor list (``build_neighbor_list`` / ``extend_coord_with_ghosts``) - instead of re-deriving neighbors. 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``. + 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 @@ -121,35 +130,117 @@ def build_neighbor_graph( atype: Array, box: Array | None, rcut: float, - sel: int | list[int], - mixed_types: bool = True, layout: GraphLayout | None = None, ) -> NeighborGraph: - """Build a LEGACY-COMPATIBLE NeighborGraph by reusing the tested dense nlist. - - Calls ``extend_input_and_build_neighbor_list`` (general-cell, tested) then - :func:`neighbor_graph_from_extended`. - - CONTRACT: the returned graph contains the neighbors selected by ``sel`` and - **inherits the legacy ``sel`` truncation** — it does NOT carry all neighbors - within ``rcut`` when ``sel`` binds. It coincides with the complete in-``rcut`` - environment ONLY when ``sel`` is large enough that no real neighbor is dropped - (the ``sel``-as-normalization regime). For the carry-all graph-native contract - use the dedicated carry-all builder (``from_ijs`` via ASE/vesin), NOT this - function. See memory/spec_unified_edge_nlist.md (decision #17). This builder is - the backward-compat adapter and the test oracle. + """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. 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_input_and_build_neighbor_list, + extend_coord_with_ghosts, + ) + from deepmd.dpmodel.utils.region import ( + normalize_coord, ) - # ``extend_input_and_build_neighbor_list`` is annotated ``sel: list[int]``; - # normalize the integer form so the public ``int | list[int]`` contract is - # honored (the underlying ``build_neighbor_list`` accepts both). - sel_list = [sel] if isinstance(sel, int) else sel - extended_coord, _extended_atype, mapping, nlist = ( - extend_input_and_build_neighbor_list( - coord, atype, rcut, sel_list, mixed_types=mixed_types, box=box + 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] + n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) + arange_nloc = xp.arange(nloc, dtype=xp.int64, device=dev) + arange_nall = xp.arange(nall, dtype=xp.int64, device=dev) + # (nloc, nall) static index grids reused per frame + ii_flat = xp.reshape( + xp.broadcast_to(arange_nloc[:, None], (nloc, nall)), (nloc * nall,) + ) + jj_flat = xp.reshape( + xp.broadcast_to(arange_nall[None, :], (nloc, nall)), (nloc * nall,) + ) + src_parts: list[Array] = [] + dst_parts: list[Array] = [] + vec_parts: list[Array] = [] + for ff in range(nf): + ec = extended_coord[ff] # (nall, 3) + centers = ec[:nloc, :] # (nloc, 3) + diff = ( + ec[None, :, :] - centers[:, None, :] + ) # (nloc, nall, 3) = ext[j]-center[i] + dist = xp.linalg.vector_norm(diff, axis=-1) # (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. Uncapped -- no sel truncation. + not_self = jj_flat != ii_flat # (nloc*nall,) + vir_nei = xp.reshape( + xp.broadcast_to((extended_atype[ff] < 0)[None, :], (nloc, nall)), + (nloc * nall,), + ) + vir_cen = xp.reshape( + xp.broadcast_to((atype[ff] < 0)[:, None], (nloc, nall)), + (nloc * nall,), ) + within = xp.reshape(dist <= rcut, (nloc * nall,)) + keep_mask = ( + within & not_self & xp.logical_not(vir_nei) & xp.logical_not(vir_cen) + ) + keep = xp.reshape(xp.nonzero(keep_mask)[0], (-1,)) + dst = xp.take(ii_flat, keep, axis=0) # local center + j_ext = xp.take(jj_flat, keep, axis=0) # extended neighbor index + src = xp.take(mapping[ff], j_ext, axis=0) # local owner of neighbor + vec = xp.take(xp.reshape(diff, (nloc * nall, 3)), keep, axis=0) + offset = ff * nloc + src_parts.append(src + offset) + dst_parts.append(dst + offset) + vec_parts.append(vec) + edge_index = xp.astype( + xp.stack([xp.concat(src_parts), xp.concat(dst_parts)], axis=0), xp.int64 + ) + edge_vec = xp.concat(vec_parts, axis=0) + edge_index, edge_vec, edge_mask = pad_and_guard_edges( + edge_index, edge_vec, layout.edge_capacity, layout.min_edges + ) + return NeighborGraph( + n_node=n_node, + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, ) - return neighbor_graph_from_extended(extended_coord, nlist, mapping, layout) diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index 20730b0048..ef8066850b 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -96,13 +96,13 @@ def test_importable_from_utils(self) -> None: NeighborGraph, build_neighbor_graph, edge_force_virial, - neighbor_graph_from_extended, + 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(neighbor_graph_from_extended)) + 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 index 63ed682e14..08c6aba6e1 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Tests for the dpmodel default NeighborGraph builder. +"""Tests for the dpmodel NeighborGraph builder/converter. -``build_neighbor_graph`` reuses deepmd's tested extended neighbor list -(``extend_input_and_build_neighbor_list``) and converts it to a NeighborGraph -via ``neighbor_graph_from_extended``. We validate it against an INDEPENDENT -brute-force all-pairs oracle defined locally in this test file (kept here, not -in the library, because the production builder reuses the already-tested -extended nlist). +``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 @@ -17,7 +18,10 @@ from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, build_neighbor_graph, - neighbor_graph_from_extended, + from_dense_quartet, +) +from deepmd.dpmodel.utils.nlist import ( + extend_input_and_build_neighbor_list, ) @@ -79,7 +83,6 @@ def graph_neighbor_sets_frame(ng, frame, nloc): class TestNeighborGraphBuilder(unittest.TestCase): def setUp(self) -> None: self.rcut = 4.0 - self.sel = [50, 50] # large -> no truncation (sel-as-normalization regime) # 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. @@ -90,9 +93,7 @@ def setUp(self) -> None: 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, self.sel, mixed_types=True - ) + 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), @@ -101,21 +102,36 @@ def test_nonperiodic_matches_brute_force(self) -> None: 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.sel, mixed_types=True - ) + 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_edge_vec_within_rcut(self) -> None: - ng = build_neighbor_graph( - self.coord, self.atype, None, self.rcut, self.sel, mixed_types=True - ) + 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; @@ -127,9 +143,7 @@ def test_multiframe_per_frame_neighbor_sets(self) -> None: ).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, self.sel, mixed_types=True - ) + 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( @@ -155,28 +169,19 @@ def test_multiframe_periodic_per_frame(self) -> None: 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, self.sel, mixed_types=True - ) + 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_int_sel_matches_list_sel(self) -> None: - # an integer ``sel`` (normalized to list form) must yield the same - # real-edge environment as the equivalent large list ``sel``. - nloc = self.coord.shape[1] - ng_int = build_neighbor_graph( - self.coord, self.atype, None, self.rcut, 64, mixed_types=True - ) - ng_list = build_neighbor_graph( - self.coord, self.atype, None, self.rcut, self.sel, mixed_types=True - ) - self.assertEqual( - graph_neighbor_sets(ng_int, nloc), graph_neighbor_sets(ng_list, nloc) - ) + def test_virtual_atoms_excluded(self) -> None: + # a virtual atom (type < 0) is neither a center nor a neighbor. + atype = np.array([[0, 1, -1, 1]], dtype=np.int64) # atom 2 virtual + ng = build_neighbor_graph(self.coord, atype, None, self.rcut) + ei = ng.edge_index[:, ng.edge_mask] + self.assertFalse(bool(np.any(ei == 2))) # node 2 never appears as src or dst def test_static_capacity_padding(self) -> None: ng = build_neighbor_graph( @@ -184,8 +189,6 @@ def test_static_capacity_padding(self) -> None: self.atype, None, self.rcut, - self.sel, - mixed_types=True, layout=GraphLayout(edge_capacity=64), ) self.assertEqual(ng.edge_index.shape[1], 64) @@ -201,13 +204,13 @@ def test_static_capacity_padding(self) -> None: self.assertTrue(np.all(ng.edge_vec[~ng.edge_mask] == 0.0)) -class TestNeighborGraphFromExtended(unittest.TestCase): +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 = neighbor_graph_from_extended(extended_coord, nlist, mapping) + ng = from_dense_quartet(extended_coord, nlist, mapping) ei = ng.edge_index[:, ng.edge_mask] ev = ng.edge_vec[ng.edge_mask] got = { @@ -233,7 +236,7 @@ def test_adapter_multiframe_offsets(self) -> None: [[[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 = neighbor_graph_from_extended(extended_coord, nlist, mapping) + 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] @@ -251,10 +254,14 @@ def test_adapter_maps_ghost_to_local_owner(self) -> None: 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 = neighbor_graph_from_extended(extended_coord, nlist, mapping) + 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() From 83f12c23817233efd9d509e45a8a93953c72e40e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 08:50:17 +0800 Subject: [PATCH 20/24] test(dpmodel): cover min_edges guard, flat coord input, virtual-center exclusion in build_neighbor_graph --- .../dpmodel/test_neighbor_graph_builder.py | 32 +++++++++++++++++-- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 08c6aba6e1..68b16cb16a 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -177,11 +177,37 @@ def test_multiframe_periodic_per_frame(self) -> None: ) def test_virtual_atoms_excluded(self) -> None: - # a virtual atom (type < 0) is neither a center nor a neighbor. - atype = np.array([[0, 1, -1, 1]], dtype=np.int64) # atom 2 virtual + # 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] - self.assertFalse(bool(np.any(ei == 2))) # node 2 never appears as src or dst + 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( From 6a3c5c4a431ad2240c66648e4302a18160d810f0 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 09:13:04 +0800 Subject: [PATCH 21/24] refactor(dpmodel): vectorize build_neighbor_graph and from_dense_quartet over frames Replace the per-frame Python loop in both the carry-all search and the quartet converter with flat (frame, center, neighbor) index grids and cross-frame gathers via (frame * nall + idx) flat indices. No behavior change (edges stay frame-major); tests unchanged and green. --- .../dpmodel/utils/neighbor_graph/builder.py | 150 +++++++++--------- 1 file changed, 72 insertions(+), 78 deletions(-) diff --git a/deepmd/dpmodel/utils/neighbor_graph/builder.py b/deepmd/dpmodel/utils/neighbor_graph/builder.py index d680385818..9a10d3f805 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/builder.py +++ b/deepmd/dpmodel/utils/neighbor_graph/builder.py @@ -22,6 +22,11 @@ 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 ( @@ -85,38 +90,39 @@ def from_dense_quartet( layout = GraphLayout() xp = array_api_compat.array_namespace(extended_coord, nlist, mapping) dev = array_api_compat.device(extended_coord) - nf = nlist.shape[0] - nloc = nlist.shape[1] - nsel = nlist.shape[2] - n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) - src_parts: list[Array] = [] - dst_parts: list[Array] = [] - vec_parts: list[Array] = [] - center_full = xp.broadcast_to( - xp.reshape(xp.arange(nloc, dtype=xp.int64, device=dev), (nloc, 1)), - (nloc, nsel), - ) - center_flat = xp.reshape(center_full, (nloc * nsel,)) - for ff in range(nf): - nl_flat = xp.reshape(nlist[ff], (nloc * nsel,)) - keep = xp.reshape(xp.nonzero(nl_flat >= 0)[0], (-1,)) - j_ext = xp.take(nl_flat, keep, axis=0) # extended neighbor indices - dst = xp.take(center_flat, keep, axis=0) # local center indices - src = xp.take(mapping[ff], j_ext, axis=0) # local owner of neighbor - vec = xp.take(extended_coord[ff], j_ext, axis=0) - xp.take( - extended_coord[ff], dst, axis=0 - ) - offset = ff * nloc - src_parts.append(src + offset) - dst_parts.append(dst + offset) - vec_parts.append(vec) + 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([xp.concat(src_parts), xp.concat(dst_parts)], axis=0), xp.int64 + xp.stack([ff_k * nloc + src_local, ff_k * nloc + dst_local], axis=0), xp.int64 ) - edge_vec = xp.concat(vec_parts, axis=0) 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, @@ -144,10 +150,10 @@ def build_neighbor_graph( 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. 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. + ``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 ---------- @@ -186,58 +192,46 @@ def build_neighbor_graph( ) extended_coord = xp.reshape(extended_coord, (nf, -1, 3)) nall = extended_coord.shape[1] - n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev) - arange_nloc = xp.arange(nloc, dtype=xp.int64, device=dev) - arange_nall = xp.arange(nall, dtype=xp.int64, device=dev) - # (nloc, nall) static index grids reused per frame - ii_flat = xp.reshape( - xp.broadcast_to(arange_nloc[:, None], (nloc, nall)), (nloc * nall,) - ) - jj_flat = xp.reshape( - xp.broadcast_to(arange_nall[None, :], (nloc, nall)), (nloc * nall,) - ) - src_parts: list[Array] = [] - dst_parts: list[Array] = [] - vec_parts: list[Array] = [] - for ff in range(nf): - ec = extended_coord[ff] # (nall, 3) - centers = ec[:nloc, :] # (nloc, 3) - diff = ( - ec[None, :, :] - centers[:, None, :] - ) # (nloc, nall, 3) = ext[j]-center[i] - dist = xp.linalg.vector_norm(diff, axis=-1) # (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. Uncapped -- no sel truncation. - not_self = jj_flat != ii_flat # (nloc*nall,) - vir_nei = xp.reshape( - xp.broadcast_to((extended_atype[ff] < 0)[None, :], (nloc, nall)), - (nloc * nall,), - ) - vir_cen = xp.reshape( - xp.broadcast_to((atype[ff] < 0)[:, None], (nloc, nall)), - (nloc * nall,), - ) - within = xp.reshape(dist <= rcut, (nloc * nall,)) - keep_mask = ( - within & not_self & xp.logical_not(vir_nei) & xp.logical_not(vir_cen) - ) - keep = xp.reshape(xp.nonzero(keep_mask)[0], (-1,)) - dst = xp.take(ii_flat, keep, axis=0) # local center - j_ext = xp.take(jj_flat, keep, axis=0) # extended neighbor index - src = xp.take(mapping[ff], j_ext, axis=0) # local owner of neighbor - vec = xp.take(xp.reshape(diff, (nloc * nall, 3)), keep, axis=0) - offset = ff * nloc - src_parts.append(src + offset) - dst_parts.append(dst + offset) - vec_parts.append(vec) + # 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([xp.concat(src_parts), xp.concat(dst_parts)], axis=0), xp.int64 + xp.stack([ff_k * nloc + src_local, ff_k * nloc + dst_local], axis=0), xp.int64 ) - edge_vec = xp.concat(vec_parts, axis=0) 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, From 1a97db9952454e0e15c5c3a58ffe96b501c80df6 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 09:15:24 +0800 Subject: [PATCH 22/24] test(dpmodel): add discriminating cross-periodic-boundary neighbor case A pair that is a neighbor ONLY across the boundary (direct distance > rcut, minimum image < rcut): build_neighbor_graph must find the image edge with the right vector, find nothing without the box, and agree with the brute-force oracle. Existing periodic test had the pair as a direct neighbor too, so it did not isolate periodic-image handling. --- .../dpmodel/test_neighbor_graph_builder.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 68b16cb16a..9ba25c0ccb 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -108,6 +108,30 @@ def test_periodic_matches_brute_force(self) -> None: 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] From 0eed16e23006f23beb764e7ba01828171f460897 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 09:17:38 +0800 Subject: [PATCH 23/24] test(dpmodel): add torch-namespace smoke test for graph builders Run build_neighbor_graph and from_dense_quartet on torch.from_numpy inputs (periodic + non-periodic) and assert the neighbor environment matches the numpy result and the output stays in the torch namespace. Guards the vectorized broadcasting/gather against numpy-only ops that break torch/jax. --- .../dpmodel/test_neighbor_graph_builder.py | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index 9ba25c0ccb..a49eb9dfbf 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -12,8 +12,12 @@ import itertools import unittest +from types import ( + SimpleNamespace, +) import numpy as np +import pytest from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, @@ -25,6 +29,22 @@ ) +def _detach_graph(ng): + """View a NeighborGraph's arrays as numpy (accepts torch or numpy tensors), + so the numpy-only ``graph_neighbor_sets`` helper can consume a torch result. + """ + + def to_np(x): + return x.detach().cpu().numpy() if hasattr(x, "detach") else np.asarray(x) + + return SimpleNamespace( + edge_index=to_np(ng.edge_index), + edge_vec=to_np(ng.edge_vec), + edge_mask=to_np(ng.edge_mask), + n_node=to_np(ng.n_node), + ) + + 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]. @@ -253,6 +273,39 @@ def test_static_capacity_padding(self) -> None: # masked-out tail contributes no real edges self.assertTrue(np.all(ng.edge_vec[~ng.edge_mask] == 0.0)) + def test_torch_namespace_smoke(self) -> None: + # array-API smoke test: both builders must run under the torch namespace + # (torch.from_numpy inputs) and produce the SAME neighbor environment as + # numpy. Catches numpy-only ops (bool*float, wrong index dtype, missing + # device=) that pass under numpy but break torch/jax (CLAUDE.md guidance). + torch = pytest.importorskip("torch") + box = np.eye(3, dtype=np.float64)[None] * 6.0 + for b in (None, box): # non-periodic AND periodic + ref = graph_neighbor_sets( + build_neighbor_graph(self.coord, self.atype, b, self.rcut), 4 + ) + t_box = None if b is None else torch.from_numpy(b) + ng_pt = build_neighbor_graph( + torch.from_numpy(self.coord), + torch.from_numpy(self.atype), + t_box, + self.rcut, + ) + # namespace preserved (output stays a torch tensor, not silently numpy) + self.assertIn("torch", type(ng_pt.edge_index).__module__) + self.assertEqual(graph_neighbor_sets(_detach_graph(ng_pt), 4), ref) + # the converter must travel the torch namespace too + ext_coord, _ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( + self.coord, self.atype, self.rcut, [50, 50], mixed_types=True, box=box + ) + ref_c = graph_neighbor_sets(from_dense_quartet(ext_coord, nlist, mapping), 4) + ng_pt_c = from_dense_quartet( + torch.from_numpy(ext_coord), + torch.from_numpy(nlist), + torch.from_numpy(mapping), + ) + self.assertEqual(graph_neighbor_sets(_detach_graph(ng_pt_c), 4), ref_c) + class TestFromDenseQuartet(unittest.TestCase): def test_adapter_on_handmade_quartet(self) -> None: From f32efc862230bb77c192c29ee6f077316b62a911 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 25 Jun 2026 09:30:35 +0800 Subject: [PATCH 24/24] test(dpmodel): drop the torch-namespace smoke test Remove the shallow 'smoke' test; torch-namespace behavior of the graph builders is covered properly by PR-A's cross-backend consistency harness (full graph forward under pt_expt/torch), not a standalone smoke check. --- .../dpmodel/test_neighbor_graph_builder.py | 53 ------------------- 1 file changed, 53 deletions(-) diff --git a/source/tests/common/dpmodel/test_neighbor_graph_builder.py b/source/tests/common/dpmodel/test_neighbor_graph_builder.py index a49eb9dfbf..9ba25c0ccb 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph_builder.py +++ b/source/tests/common/dpmodel/test_neighbor_graph_builder.py @@ -12,12 +12,8 @@ import itertools import unittest -from types import ( - SimpleNamespace, -) import numpy as np -import pytest from deepmd.dpmodel.utils.neighbor_graph import ( GraphLayout, @@ -29,22 +25,6 @@ ) -def _detach_graph(ng): - """View a NeighborGraph's arrays as numpy (accepts torch or numpy tensors), - so the numpy-only ``graph_neighbor_sets`` helper can consume a torch result. - """ - - def to_np(x): - return x.detach().cpu().numpy() if hasattr(x, "detach") else np.asarray(x) - - return SimpleNamespace( - edge_index=to_np(ng.edge_index), - edge_vec=to_np(ng.edge_vec), - edge_mask=to_np(ng.edge_mask), - n_node=to_np(ng.n_node), - ) - - 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]. @@ -273,39 +253,6 @@ def test_static_capacity_padding(self) -> None: # masked-out tail contributes no real edges self.assertTrue(np.all(ng.edge_vec[~ng.edge_mask] == 0.0)) - def test_torch_namespace_smoke(self) -> None: - # array-API smoke test: both builders must run under the torch namespace - # (torch.from_numpy inputs) and produce the SAME neighbor environment as - # numpy. Catches numpy-only ops (bool*float, wrong index dtype, missing - # device=) that pass under numpy but break torch/jax (CLAUDE.md guidance). - torch = pytest.importorskip("torch") - box = np.eye(3, dtype=np.float64)[None] * 6.0 - for b in (None, box): # non-periodic AND periodic - ref = graph_neighbor_sets( - build_neighbor_graph(self.coord, self.atype, b, self.rcut), 4 - ) - t_box = None if b is None else torch.from_numpy(b) - ng_pt = build_neighbor_graph( - torch.from_numpy(self.coord), - torch.from_numpy(self.atype), - t_box, - self.rcut, - ) - # namespace preserved (output stays a torch tensor, not silently numpy) - self.assertIn("torch", type(ng_pt.edge_index).__module__) - self.assertEqual(graph_neighbor_sets(_detach_graph(ng_pt), 4), ref) - # the converter must travel the torch namespace too - ext_coord, _ext_atype, mapping, nlist = extend_input_and_build_neighbor_list( - self.coord, self.atype, self.rcut, [50, 50], mixed_types=True, box=box - ) - ref_c = graph_neighbor_sets(from_dense_quartet(ext_coord, nlist, mapping), 4) - ng_pt_c = from_dense_quartet( - torch.from_numpy(ext_coord), - torch.from_numpy(nlist), - torch.from_numpy(mapping), - ) - self.assertEqual(graph_neighbor_sets(_detach_graph(ng_pt_c), 4), ref_c) - class TestFromDenseQuartet(unittest.TestCase): def test_adapter_on_handmade_quartet(self) -> None: