feat(dpmodel): NeighborGraph 3-body angle machinery (PR-E)

deepmd/dpmodel/utils/neighbor_graph/__init__.py

@@ -9,6 +9,14 @@
 See the design discussion wanghan-iapcm/deepmd-kit#4.
 """
 
+from .angles import (
+    angle_padding_fraction,
+    angle_to_edge_sum,
+    angle_to_node_sum,
+    attach_angles,
+    build_angle_index,
+    graph_angle_cos,
+)
 from .ase_builder import (
     build_neighbor_graph_ase,
 )
@@ -30,6 +38,7 @@
     NeighborGraph,
     frame_id_from_n_node,
     node_validity_mask,
+    pad_and_guard_angles,
     pad_and_guard_edges,
 )
 from .pairs import (
@@ -45,15 +54,22 @@
 __all__ = [
     "GraphLayout",
     "NeighborGraph",
+    "angle_padding_fraction",
+    "angle_to_edge_sum",
+    "angle_to_node_sum",
+    "attach_angles",
+    "build_angle_index",
     "build_neighbor_graph",
     "build_neighbor_graph_ase",
     "center_edge_pairs",
     "edge_env_mat",
     "edge_force_virial",
     "frame_id_from_n_node",
     "from_dense_quartet",
+    "graph_angle_cos",
     "neighbor_graph_from_ijs",
     "node_validity_mask",
+    "pad_and_guard_angles",
     "pad_and_guard_edges",
     "segment_max",
     "segment_mean",

deepmd/dpmodel/utils/neighbor_graph/angles.py

@@ -0,0 +1,254 @@
+# 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
+
+import dataclasses
+
+from .graph import (
+    GraphLayout,
+    NeighborGraph,
+    pad_and_guard_angles,
+)
+from .pairs import (
+    center_edge_pairs,
+)
+from .segment import (
+    segment_sum,
+)
+
+
+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
+
+
+def attach_angles(
+    graph: NeighborGraph,
+    a_rcut: float,
+    *,
+    ordered: bool = False,
+    include_self: bool = False,
+    layout: GraphLayout | None = None,
+) -> NeighborGraph:
+    """Attach angle_index/angle_mask to an existing edge-only NeighborGraph.
+
+    Parameters
+    ----------
+    graph : NeighborGraph
+        Input graph (edge fields must be populated).
+    a_rcut : float
+        Angle cutoff radius. Only edges with norm < a_rcut participate.
+    ordered : bool, optional
+        If True, include both (a, b) and (b, a) angle pairs.
+    include_self : bool, optional
+        If True, include self-angle pairs (a, a).
+    layout : GraphLayout or None, optional
+        If provided, uses layout.angle_capacity and layout.node_capacity.
+
+    Returns
+    -------
+    NeighborGraph
+        A new NeighborGraph with angle_index and angle_mask populated;
+        all edge/node fields are unchanged.
+    """
+    xp = array_api_compat.array_namespace(graph.edge_index)
+    if layout is not None and layout.node_capacity is not None:
+        n_total = layout.node_capacity
+    else:
+        n_total = int(xp.sum(graph.n_node))
+    ai, am = build_angle_index(
+        graph.edge_index,
+        graph.edge_vec,
+        graph.edge_mask,
+        n_total,
+        a_rcut,
+        ordered=ordered,
+        include_self=include_self,
+        layout=layout,
+    )
+    return dataclasses.replace(graph, angle_index=ai, angle_mask=am)
+
+
+def graph_angle_cos(angle_index: Array, edge_vec: Array, eps: float = 1e-6) -> Array:
+    """Per-angle cosine, mirroring dpa3 ``cosine_ij`` (repflows.py:632-644).
+
+    Parameters
+    ----------
+    angle_index : Array
+        Shape (2, A) index pairs into edge list. ``angle_index[0, a]`` is
+        edge_a and ``angle_index[1, a]`` is edge_b for angle ``a``.
+    edge_vec : Array
+        Shape (E, 3) edge vectors (r_src - r_dst, i.e. neighbor - center).
+    eps : float, optional
+        Numerical stabiliser: norm denominators use ``||v|| + eps`` and the
+        dot product is scaled by ``(1 - eps)``.  Mirrors the dpa3 dense
+        channel exactly (repflows.py:643-649).
+
+    Returns
+    -------
+    Array
+        Shape (A,) cosine values, one per angle slot (valid and padding).
+        Padding slots carry arbitrary values; mask with angle_mask before use.
+    """
+    xp = array_api_compat.array_namespace(edge_vec)
+    va = xp.take(edge_vec, angle_index[0, :], axis=0)  # (A, 3)
+    vb = xp.take(edge_vec, angle_index[1, :], axis=0)  # (A, 3)
+    na = va / (xp.linalg.vector_norm(va, axis=-1, keepdims=True) + eps)
+    nb = vb / (xp.linalg.vector_norm(vb, axis=-1, keepdims=True) + eps)
+    return xp.sum(na * nb, axis=-1) * (1.0 - eps)
+
+
+def angle_to_edge_sum(data: Array, angle_index: Array, num_edges: int) -> Array:
+    """Aggregate per-angle data to the angle's query edge (edge_a).
+
+    Parameters
+    ----------
+    data : Array
+        Shape (A,) or (A, ...) per-angle data to aggregate.
+    angle_index : Array
+        Shape (2, A) angle index pairs into edges.
+    num_edges : int
+        Total number of edges (E).
+
+    Returns
+    -------
+    Array
+        Shape (E,) or (E, ...) aggregated per-edge data.
+    """
+    return segment_sum(data, angle_index[0, :], num_edges)
+
+
+def angle_to_node_sum(
+    data: Array, angle_index: Array, edge_index: Array, num_nodes: int
+) -> Array:
+    """Aggregate per-angle data to the shared center (dst of edge_a).
+
+    Parameters
+    ----------
+    data : Array
+        Shape (A,) or (A, ...) per-angle data to aggregate.
+    angle_index : Array
+        Shape (2, A) angle index pairs into edges.
+    edge_index : Array
+        Shape (2, E) edge indices [src, dst].
+    num_nodes : int
+        Total number of nodes (N).
+
+    Returns
+    -------
+    Array
+        Shape (N,) or (N, ...) aggregated per-node data.
+    """
+    xp = array_api_compat.array_namespace(data)
+    center = xp.take(edge_index[1, :], angle_index[0, :], axis=0)
+    return segment_sum(data, center, num_nodes)
+
+
+def angle_padding_fraction(graph: NeighborGraph) -> float:
+    """Return the fraction of angle slots that are padding (guard entries).
+
+    Parameters
+    ----------
+    graph : NeighborGraph
+        A graph with ``angle_mask`` set (i.e., after :func:`attach_angles`
+        with a static ``GraphLayout.angle_capacity``).
+
+    Returns
+    -------
+    float
+        ``1 - A_real / A_max`` where ``A_real`` is the count of valid angles
+        and ``A_max`` is ``angle_mask.shape[0]``.  Returns ``0.0`` when the
+        mask is empty.
+    """
+    if graph.angle_mask is None:
+        return 0.0
+    xp = array_api_compat.array_namespace(graph.angle_mask)
+    total = graph.angle_mask.shape[0]
+    if total == 0:
+        return 0.0
+    real = int(xp.sum(xp.astype(graph.angle_mask, xp.int64)))
+    return 1.0 - real / total

deepmd/dpmodel/utils/neighbor_graph/graph.py

@@ -123,6 +123,57 @@ def pad_and_guard_edges(
     return ei, ev, edge_mask
 
 
+def pad_and_guard_angles(
+    angle_index: Array,
+    angle_capacity: int | None = None,
+    min_angles: int = 2,
+    pad_value: int = 0,
+) -> tuple[Array, Array]:
+    """Append padding/guard angles as a contiguous suffix and build angle_mask.
+
+    Real angles (``angle_index``) stay at the front (compact layout).
+    Dummy angles point at edge ``pad_value`` (in-range).
+
+    Parameters
+    ----------
+    angle_index
+        (2, A_real) ``[edge_a, edge_b]`` edge endpoints of the real angles.
+    angle_capacity
+        Target angle-axis length ``A_max``. ``None`` (torch dynamic) appends
+        exactly ``min_angles`` masked dummy angles so the axis has a known lower
+        bound and shape-stable guards for export; an int (jax static) pads to
+        ``A_max = angle_capacity`` and raises ``ValueError`` on overflow.
+    min_angles
+        Number of dummy angles appended when ``angle_capacity is None``.
+    pad_value
+        Edge index the dummy angles point at (must be in range).
+
+    Returns
+    -------
+    angle_index
+        (2, target) padded angle endpoints.
+    angle_mask
+        (target,) boolean mask, ``True`` for the real-angle prefix.
+    """
+    xp = array_api_compat.array_namespace(angle_index)
+    dev = array_api_compat.device(angle_index)
+    a_real = angle_index.shape[1]
+    if angle_capacity is None:
+        target = a_real + min_angles
+    else:
+        if a_real > angle_capacity:
+            raise ValueError(
+                f"angle overflow: {a_real} real angles > angle_capacity {angle_capacity}"
+            )
+        target = angle_capacity
+    n_pad = target - a_real
+    pad_idx = xp.full((2, n_pad), pad_value, dtype=angle_index.dtype, device=dev)
+    ai = xp.concat([angle_index, pad_idx], axis=1)
+    arange = xp.arange(target, dtype=angle_index.dtype, device=dev)
+    angle_mask = arange < a_real
+    return ai, angle_mask
+
+
 def frame_id_from_n_node(n_node: Array, n_total: int | None = None) -> Array:
     """Node->frame map for a flat node axis: ``repeat(arange(nf), n_node)``.