feat(dpmodel): build_angle_index (edge pairs within a_rcut)

Author: Han Wang

deepmd/dpmodel/utils/neighbor_graph/__init__.py

@@ -9,6 +9,9 @@
 See the design discussion wanghan-iapcm/deepmd-kit#4.
 """
 
+from .angles import (
+    build_angle_index,
+)
 from .ase_builder import (
     build_neighbor_graph_ase,
 )
@@ -46,6 +49,7 @@
 __all__ = [
     "GraphLayout",
     "NeighborGraph",
+    "build_angle_index",
     "build_neighbor_graph",
     "build_neighbor_graph_ase",
     "center_edge_pairs",

deepmd/dpmodel/utils/neighbor_graph/angles.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""3-body angle graph: pairs of edges sharing a center within a_rcut.
+
+Angles reference EDGES (angle_index into [0,E)); edge_vec stays the only
+geometry leaf. a_sel is normalization-only (not a truncation). Reuses PR-D's
+center_edge_pairs; a_rcut filters the participating edges.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import array_api_compat
+
+if TYPE_CHECKING:
+    from deepmd.dpmodel.array_api import Array
+
+from .graph import GraphLayout, pad_and_guard_angles
+from .pairs import center_edge_pairs
+
+
+def build_angle_index(
+    edge_index: Array,
+    edge_vec: Array,
+    edge_mask: Array,
+    n_total: int,
+    a_rcut: float,
+    *,
+    ordered: bool = False,
+    include_self: bool = False,
+    layout: GraphLayout | None = None,
+) -> tuple[Array, Array]:
+    """Build angle index for 3-body terms.
+
+    Parameters
+    ----------
+    edge_index : Array
+        Shape (2, E) [src, dst] SoA edge indices.
+    edge_vec : Array
+        Shape (E, 3) edge vectors (neighbor - center).
+    edge_mask : Array
+        Shape (E,) boolean validity mask for edges.
+    n_total : int
+        Total number of nodes.
+    a_rcut : float
+        Angle cutoff. Only edges with norm < a_rcut participate in angles.
+    ordered : bool, optional
+        If True, include both (a, b) and (b, a) pairs (ordered pairs).
+    include_self : bool, optional
+        If True, include self-angle pairs (a, a).
+    layout : GraphLayout or None, optional
+        If provided, uses layout.angle_capacity as static padding capacity.
+
+    Returns
+    -------
+    angle_index : Array
+        Shape (2, A) index pairs into the edge list.
+    angle_mask : Array
+        Shape (A,) boolean mask for valid angles.
+    """
+    xp = array_api_compat.array_namespace(edge_index)
+    # a_rcut edge gate: only edges within a_rcut may participate in an angle
+    dist = xp.linalg.vector_norm(edge_vec, axis=-1)  # (E,)
+    a_edge_mask = xp.astype(edge_mask, xp.bool) & (dist < a_rcut)
+    # compact eager form only (static_nnei not exposed until angle export is
+    # needed, PR-G). dst = edge_index[1, :] per the [src, dst] SoA convention.
+    q_e, k_e, pair_mask = center_edge_pairs(
+        edge_index[1, :],
+        a_edge_mask,
+        n_total,
+        include_self=include_self,
+        ordered=ordered,
+    )
+    # compact form returns all-True pair_mask, but NEVER discard it: the
+    # shape-static form keeps filtered pairs and invalidates them only here.
+    angle_index = xp.stack([q_e, k_e], axis=0)  # (2, A_real)
+    cap = layout.angle_capacity if layout is not None else None
+    ai, am = pad_and_guard_angles(angle_index, cap, min_angles=2)
+    # fold pair_mask into the real-angle prefix of the padded mask
+    pm_padded = xp.concat(
+        [
+            pair_mask,
+            xp.zeros(
+                (am.shape[0] - pair_mask.shape[0],),
+                dtype=xp.bool,
+                device=array_api_compat.device(pair_mask),
+            ),
+        ],
+        axis=0,
+    )
+    return ai, am & pm_padded

source/tests/common/dpmodel/test_angle_builder.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-from deepmd.dpmodel.utils.neighbor_graph import pad_and_guard_angles
+from deepmd.dpmodel.utils.neighbor_graph import build_angle_index, pad_and_guard_angles
 
 
 def test_pad_angles_dynamic_appends_min_guard():
@@ -22,3 +22,101 @@ def test_pad_angles_overflow_raises():
     ai = np.zeros((2, 6), dtype=np.int64)
     with pytest.raises(ValueError):
         pad_and_guard_angles(ai, angle_capacity=4)
+
+
+def _angle_oracle(dst, evnorm, mask, a_rcut, ordered, include_self):
+    """All (edge_a, edge_b) sharing a center, both edges within a_rcut."""
+    out = set()
+    for a in range(len(dst)):
+        if not mask[a] or evnorm[a] >= a_rcut:
+            continue
+        for b in range(len(dst)):
+            if not mask[b] or evnorm[b] >= a_rcut or dst[a] != dst[b]:
+                continue
+            if not include_self and a == b:
+                continue
+            if not ordered and b < a:
+                continue
+            out.add((a, b))
+    return out
+
+
+def test_build_angle_index_matches_oracle():
+    # 4 edges, all dst=0; norms [0.5, 0.9, 2.5, 0.7], a_rcut=1.0
+    # DEFAULT = unordered, no-self => within a_rcut {0,1,3} give {(0,1),(0,3),(1,3)}
+    edge_index = np.array([[1, 2, 3, 1], [0, 0, 0, 0]], dtype=np.int64)
+    edge_vec = np.array([[0.5, 0, 0], [0.9, 0, 0], [2.5, 0, 0], [0.7, 0, 0]])
+    edge_mask = np.array([True, True, True, True])
+    ai, am = build_angle_index(edge_index, edge_vec, edge_mask, 4, a_rcut=1.0)
+    got = {(int(ai[0, p]), int(ai[1, p])) for p in range(ai.shape[1]) if am[p]}
+    evnorm = np.linalg.norm(edge_vec, axis=-1)
+    assert got == _angle_oracle(
+        [0, 0, 0, 0], evnorm, edge_mask, 1.0, ordered=False, include_self=False
+    )
+    assert got == {
+        (0, 1),
+        (0, 3),
+        (1, 3),
+    }  # unordered, no-self, edge 2 dropped (norm>a_rcut)
+    assert all(
+        2 not in (int(ai[0, p]), int(ai[1, p])) for p in range(ai.shape[1]) if am[p]
+    )
+    assert all(
+        int(ai[0, p]) != int(ai[1, p]) for p in range(ai.shape[1]) if am[p]
+    )  # no self
+
+
+def test_build_angle_index_ordered_include_self():
+    # ordered + include_self: (0,0),(0,1),(0,3),(1,0),(1,1),(1,3),(3,0),(3,1),(3,3)
+    edge_index = np.array([[1, 2, 3, 1], [0, 0, 0, 0]], dtype=np.int64)
+    edge_vec = np.array([[0.5, 0, 0], [0.9, 0, 0], [2.5, 0, 0], [0.7, 0, 0]])
+    edge_mask = np.array([True, True, True, True])
+    ai, am = build_angle_index(
+        edge_index, edge_vec, edge_mask, 4, a_rcut=1.0, ordered=True, include_self=True
+    )
+    got = {(int(ai[0, p]), int(ai[1, p])) for p in range(ai.shape[1]) if am[p]}
+    evnorm = np.linalg.norm(edge_vec, axis=-1)
+    assert got == _angle_oracle(
+        [0, 0, 0, 0], evnorm, edge_mask, 1.0, ordered=True, include_self=True
+    )
+
+
+def test_build_angle_index_masked_edge():
+    # edge 1 masked out — should not appear in any angle
+    edge_index = np.array([[1, 2, 3, 1], [0, 0, 0, 0]], dtype=np.int64)
+    edge_vec = np.array([[0.5, 0, 0], [0.9, 0, 0], [0.3, 0, 0], [0.7, 0, 0]])
+    edge_mask = np.array([True, False, True, True])
+    ai, am = build_angle_index(edge_index, edge_vec, edge_mask, 4, a_rcut=1.0)
+    got = {(int(ai[0, p]), int(ai[1, p])) for p in range(ai.shape[1]) if am[p]}
+    evnorm = np.linalg.norm(edge_vec, axis=-1)
+    assert got == _angle_oracle(
+        [0, 0, 0, 0], evnorm, edge_mask, 1.0, ordered=False, include_self=False
+    )
+    assert all(
+        1 not in (int(ai[0, p]), int(ai[1, p])) for p in range(ai.shape[1]) if am[p]
+    )
+
+
+def test_build_angle_index_torch_namespace():
+    # Step 4b: torch-namespace smoke test (function-level import for TID253)
+    import torch
+
+    edge_index = np.array([[1, 2, 3, 1], [0, 0, 0, 0]], dtype=np.int64)
+    edge_vec = np.array([[0.5, 0, 0], [0.9, 0, 0], [2.5, 0, 0], [0.7, 0, 0]])
+    edge_mask = np.array([True, True, True, True])
+
+    ai_np, am_np = build_angle_index(edge_index, edge_vec, edge_mask, 4, a_rcut=1.0)
+    got_np = {
+        (int(ai_np[0, p]), int(ai_np[1, p])) for p in range(ai_np.shape[1]) if am_np[p]
+    }
+
+    t_edge_index = torch.from_numpy(edge_index)
+    t_edge_vec = torch.from_numpy(edge_vec)
+    t_edge_mask = torch.from_numpy(edge_mask)
+    ai_t, am_t = build_angle_index(t_edge_index, t_edge_vec, t_edge_mask, 4, a_rcut=1.0)
+    got_t = {
+        (int(ai_t[0, p].item()), int(ai_t[1, p].item()))
+        for p in range(ai_t.shape[1])
+        if am_t[p].item()
+    }
+    assert got_t == got_np