From 93ff993ab79610b5511d0949cd8289cea1f38094 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 11:29:24 +0800
Subject: [PATCH 01/12] feat(dpmodel): pad_and_guard_angles angle-axis padder

---
 .../dpmodel/utils/neighbor_graph/__init__.py  |  2 +
 deepmd/dpmodel/utils/neighbor_graph/graph.py  | 51 +++++++++++++++++++
 .../common/dpmodel/test_angle_builder.py      | 24 +++++++++
 3 files changed, 77 insertions(+)
 create mode 100644 source/tests/common/dpmodel/test_angle_builder.py

diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index 24fb090309..9df8ad5b24 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
@@ -30,6 +30,7 @@
     NeighborGraph,
     frame_id_from_n_node,
     node_validity_mask,
+    pad_and_guard_angles,
     pad_and_guard_edges,
 )
 from .pairs import (
@@ -54,6 +55,7 @@
     "from_dense_quartet",
     "neighbor_graph_from_ijs",
     "node_validity_mask",
+    "pad_and_guard_angles",
     "pad_and_guard_edges",
     "segment_max",
     "segment_mean",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py
index 0ce10efdf6..67f7d53845 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/graph.py
+++ b/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)``.
 
diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
new file mode 100644
index 0000000000..ec8418ee77
--- /dev/null
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -0,0 +1,24 @@
+import numpy as np
+import pytest
+
+from deepmd.dpmodel.utils.neighbor_graph import pad_and_guard_angles
+
+
+def test_pad_angles_dynamic_appends_min_guard():
+    ai = np.array([[0, 1], [1, 0]], dtype=np.int64)  # 2 real angles
+    out_ai, out_mask = pad_and_guard_angles(ai, angle_capacity=None, min_angles=2)
+    assert out_ai.shape == (2, 4)  # 2 real + 2 guard
+    np.testing.assert_array_equal(out_mask, [True, True, False, False])
+
+
+def test_pad_angles_static_capacity():
+    ai = np.array([[0, 1], [1, 0]], dtype=np.int64)
+    out_ai, out_mask = pad_and_guard_angles(ai, angle_capacity=5)
+    assert out_ai.shape == (2, 5)
+    assert int(out_mask.sum()) == 2
+
+
+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)

From fa7576d092f3adb9a7344b55db4eed92b4d6e972 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 11:38:09 +0800
Subject: [PATCH 02/12] feat(dpmodel): build_angle_index (edge pairs within
 a_rcut)

---
 .../dpmodel/utils/neighbor_graph/__init__.py  |   4 +
 deepmd/dpmodel/utils/neighbor_graph/angles.py |  91 ++++++++++++++++
 .../common/dpmodel/test_angle_builder.py      | 100 +++++++++++++++++-
 3 files changed, 194 insertions(+), 1 deletion(-)
 create mode 100644 deepmd/dpmodel/utils/neighbor_graph/angles.py

diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index 9df8ad5b24..2e796ace4f 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/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",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/angles.py b/deepmd/dpmodel/utils/neighbor_graph/angles.py
new file mode 100644
index 0000000000..9bf16e8615
--- /dev/null
+++ b/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
diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index ec8418ee77..d8e7206837 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/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

From 0467e43564859a367bc50846413e4e8ceff5ca10 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 11:43:17 +0800
Subject: [PATCH 03/12] test(dpmodel): build_angle_index multi-center,
 static-layout, ordered-no-self coverage

---
 .../common/dpmodel/test_angle_builder.py      | 85 ++++++++++++++++++-
 1 file changed, 84 insertions(+), 1 deletion(-)

diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index d8e7206837..ef6be7d947 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -1,7 +1,11 @@
 import numpy as np
 import pytest
 
-from deepmd.dpmodel.utils.neighbor_graph import build_angle_index, pad_and_guard_angles
+from deepmd.dpmodel.utils.neighbor_graph import (
+    GraphLayout,
+    build_angle_index,
+    pad_and_guard_angles,
+)
 
 
 def test_pad_angles_dynamic_appends_min_guard():
@@ -120,3 +124,82 @@ def test_build_angle_index_torch_namespace():
         if am_t[p].item()
     }
     assert got_t == got_np
+
+
+def test_build_angle_index_multi_center():
+    # Edges with MIXED centers: dst=[0,1,0,1,2]; exercises the dst[a]!=dst[b] exclusion
+    # src=[1,2,3,4,5], dst=[0,1,0,1,2], all norms within a_rcut
+    # Expected angles per center:
+    #  dst=0: edges {0,2} => {(0,2),(2,0)}
+    #  dst=1: edges {1,3} => {(1,3),(3,1)}
+    #  dst=2: edge {4} => no pairs
+    edge_index = np.array([[1, 2, 3, 4, 5], [0, 1, 0, 1, 2]], dtype=np.int64)
+    edge_vec = np.array(
+        [[0.3, 0, 0], [0.4, 0, 0], [0.5, 0, 0], [0.6, 0, 0], [0.7, 0, 0]]
+    )
+    edge_mask = np.array([True, True, True, True, True])
+    ai, am = build_angle_index(edge_index, edge_vec, edge_mask, 6, 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)
+    expected = _angle_oracle(
+        [0, 1, 0, 1, 2], evnorm, edge_mask, 1.0, ordered=False, include_self=False
+    )
+    assert got == expected
+    # Verify no cross-center angles
+    assert all(
+        edge_index[1, int(ai[0, p])] == edge_index[1, int(ai[1, p])]
+        for p in range(ai.shape[1])
+        if am[p]
+    )
+
+
+def test_build_angle_index_static_layout():
+    # Test with static layout.angle_capacity; shape must be (2, capacity)
+    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])
+    layout = GraphLayout(edge_capacity=100, angle_capacity=10)
+    ai, am = build_angle_index(
+        edge_index, edge_vec, edge_mask, 4, a_rcut=1.0, layout=layout
+    )
+    # Check static shape
+    assert ai.shape == (2, 10)
+    assert am.shape == (10,)
+    # Check real angles match the dynamic result
+    got_static = {(int(ai[0, p]), int(ai[1, p])) for p in range(ai.shape[1]) if am[p]}
+    ai_dyn, am_dyn = build_angle_index(edge_index, edge_vec, edge_mask, 4, a_rcut=1.0)
+    got_dynamic = {
+        (int(ai_dyn[0, p]), int(ai_dyn[1, p]))
+        for p in range(ai_dyn.shape[1])
+        if am_dyn[p]
+    }
+    assert got_static == got_dynamic
+    # Check mask counts match
+    assert int(am.sum()) == int(am_dyn.sum())
+
+
+def test_build_angle_index_ordered_no_self():
+    # Test ordered=True, include_self=False; should be symmetric pairs excluding diagonals
+    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=False
+    )
+    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)
+    expected = _angle_oracle(
+        [0, 0, 0, 0], evnorm, edge_mask, 1.0, ordered=True, include_self=False
+    )
+    assert got == expected
+    # Verify no self-angles and ordered includes both directions
+    assert all(
+        int(ai[0, p]) != int(ai[1, p]) for p in range(ai.shape[1]) if am[p]
+    )  # no self
+    for pair in got:
+        a, b = pair
+        # For ordered=True, include_self=False, we expect both (a,b) and (b,a)
+        # if they are different and both within a_rcut and same center
+        rev_pair = (b, a)
+        if a != b and a not in [2] and b not in [2]:  # edge 2 is outside a_rcut
+            assert rev_pair in got  # symmetric pairs should both exist

From cae4c5ae3097b6293660236fb89ec98f06db3e45 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 11:48:22 +0800
Subject: [PATCH 04/12] feat(dpmodel): attach_angles post-hoc angle attachment

---
 .../dpmodel/utils/neighbor_graph/__init__.py  |  2 +
 deepmd/dpmodel/utils/neighbor_graph/angles.py | 51 ++++++++++-
 .../common/dpmodel/test_angle_builder.py      | 84 +++++++++++++++++++
 3 files changed, 136 insertions(+), 1 deletion(-)

diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index 2e796ace4f..03b83282e3 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
@@ -10,6 +10,7 @@
 """
 
 from .angles import (
+    attach_angles,
     build_angle_index,
 )
 from .ase_builder import (
@@ -49,6 +50,7 @@
 __all__ = [
     "GraphLayout",
     "NeighborGraph",
+    "attach_angles",
     "build_angle_index",
     "build_neighbor_graph",
     "build_neighbor_graph_ase",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/angles.py b/deepmd/dpmodel/utils/neighbor_graph/angles.py
index 9bf16e8615..5f6c7d9c72 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/angles.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/angles.py
@@ -15,7 +15,9 @@
 if TYPE_CHECKING:
     from deepmd.dpmodel.array_api import Array
 
-from .graph import GraphLayout, pad_and_guard_angles
+import dataclasses
+
+from .graph import GraphLayout, NeighborGraph, pad_and_guard_angles
 from .pairs import center_edge_pairs
 
 
@@ -89,3 +91,50 @@ def build_angle_index(
         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)
diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index ef6be7d947..97c834a7ae 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -3,7 +3,9 @@
 
 from deepmd.dpmodel.utils.neighbor_graph import (
     GraphLayout,
+    attach_angles,
     build_angle_index,
+    build_neighbor_graph,
     pad_and_guard_angles,
 )
 
@@ -203,3 +205,85 @@ def test_build_angle_index_ordered_no_self():
         rev_pair = (b, a)
         if a != b and a not in [2] and b not in [2]:  # edge 2 is outside a_rcut
             assert rev_pair in got  # symmetric pairs should both exist
+
+
+# ---------------------------------------------------------------------------
+# Task 3: attach_angles tests
+# ---------------------------------------------------------------------------
+
+
+def test_attach_angles_sets_fields_and_preserves_edges():
+    """attach_angles populates angle_index/mask; edge fields are unchanged."""
+    coord = np.array([[[0.0, 0, 0], [0.8, 0, 0], [0, 0.8, 0]]])
+    atype = np.array([[0, 0, 0]])  # (nf, nloc)
+    ng = build_neighbor_graph(coord, atype, None, 2.0)
+    # default carry-all builder leaves angles None
+    assert ng.angle_index is None
+    assert ng.angle_mask is None
+    ng2 = attach_angles(ng, a_rcut=1.5)
+    assert ng2.angle_index is not None and ng2.angle_mask is not None
+    # edge fields must be identical (by value and shape)
+    np.testing.assert_array_equal(np.asarray(ng2.edge_index), np.asarray(ng.edge_index))
+    np.testing.assert_array_equal(np.asarray(ng2.edge_mask), np.asarray(ng.edge_mask))
+    np.testing.assert_array_equal(np.asarray(ng2.edge_vec), np.asarray(ng.edge_vec))
+
+
+def test_attach_angles_angle_shape_consistent():
+    """angle_index has shape (2, A) and angle_mask has shape (A,)."""
+    coord = np.array([[[0.0, 0, 0], [0.5, 0, 0], [0, 0.5, 0]]])
+    atype = np.array([[0, 0, 0]])  # (nf, nloc)
+    ng = build_neighbor_graph(coord, atype, None, 2.0)
+    ng2 = attach_angles(ng, a_rcut=1.5)
+    assert ng2.angle_index.shape[0] == 2
+    assert ng2.angle_mask.shape[0] == ng2.angle_index.shape[1]
+
+
+def test_attach_angles_valid_angles_reference_valid_edges():
+    """All valid angle pairs (q_e, k_e) must index edges that are within a_rcut."""
+    coord = np.array([[[0.0, 0, 0], [0.6, 0, 0], [0, 0.6, 0]]])
+    atype = np.array([[0, 0, 0]])  # (nf, nloc)
+    ng = build_neighbor_graph(coord, atype, None, 2.0)
+    ng2 = attach_angles(ng, a_rcut=1.0)
+    ei = np.asarray(ng2.edge_index)
+    ev = np.asarray(ng2.edge_vec)
+    em = np.asarray(ng2.edge_mask)
+    ai = np.asarray(ng2.angle_index)
+    am = np.asarray(ng2.angle_mask)
+    for p in range(am.shape[0]):
+        if not am[p]:
+            continue
+        q, k = int(ai[0, p]), int(ai[1, p])
+        # both referenced edges must be valid and within a_rcut
+        assert em[q] and em[k]
+        assert np.linalg.norm(ev[q]) < 1.0
+        assert np.linalg.norm(ev[k]) < 1.0
+        # both referenced edges must share the same center (dst)
+        assert ei[1, q] == ei[1, k]
+
+
+def test_attach_angles_with_layout():
+    """Static layout.angle_capacity is respected."""
+    coord = np.array([[[0.0, 0, 0], [0.6, 0, 0], [0, 0.6, 0]]])
+    atype = np.array([[0, 0, 0]])  # (nf, nloc)
+    ng = build_neighbor_graph(coord, atype, None, 2.0)
+    layout = GraphLayout(edge_capacity=100, angle_capacity=20)
+    ng2 = attach_angles(ng, a_rcut=1.5, layout=layout)
+    assert ng2.angle_index.shape == (2, 20)
+    assert ng2.angle_mask.shape == (20,)
+
+
+def test_attach_angles_ordered_include_self():
+    """ordered=True, include_self=True produces a superset of default pairs."""
+    coord = np.array([[[0.0, 0, 0], [0.5, 0, 0], [0, 0.5, 0]]])
+    atype = np.array([[0, 0, 0]])  # (nf, nloc)
+    ng = build_neighbor_graph(coord, atype, None, 2.0)
+    ng_default = attach_angles(ng, a_rcut=1.5)
+    ng_full = attach_angles(ng, a_rcut=1.5, ordered=True, include_self=True)
+    ai_def = np.asarray(ng_default.angle_index)
+    am_def = np.asarray(ng_default.angle_mask)
+    ai_full = np.asarray(ng_full.angle_index)
+    am_full = np.asarray(ng_full.angle_mask)
+    n_default = int(am_def.sum())
+    n_full = int(am_full.sum())
+    # ordered+include_self must produce at least as many angles as default
+    assert n_full >= n_default

From 1bd1234fa5719b8c1f7ee86335bc0bf068cba28b Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 11:51:20 +0800
Subject: [PATCH 05/12] test(dpmodel): attach_angles layout.node_capacity
 branch

---
 .../common/dpmodel/test_angle_builder.py      | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index 97c834a7ae..f09a3349b5 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -287,3 +287,27 @@ def test_attach_angles_ordered_include_self():
     n_full = int(am_full.sum())
     # ordered+include_self must produce at least as many angles as default
     assert n_full >= n_default
+
+
+def test_attach_angles_with_layout_node_capacity():
+    """layout.node_capacity branch: node_capacity is used as n_total."""
+    coord = np.array([[[0.0, 0, 0], [0.6, 0, 0], [0, 0.6, 0]]])
+    atype = np.array([[0, 0, 0]])  # (nf, nloc)
+    ng = build_neighbor_graph(coord, atype, None, 2.0)
+    layout = GraphLayout(edge_capacity=100, angle_capacity=20, node_capacity=10)
+    ng2 = attach_angles(ng, a_rcut=1.5, layout=layout)
+    assert ng2.angle_index is not None
+    assert ng2.angle_mask.shape == (20,)
+    # same real-angle set as the dynamic path (node_capacity only oversizes n_total)
+    ng3 = attach_angles(ng, a_rcut=1.5)
+    got = {
+        (int(ng2.angle_index[0, p]), int(ng2.angle_index[1, p]))
+        for p in range(ng2.angle_index.shape[1])
+        if ng2.angle_mask[p]
+    }
+    ref = {
+        (int(ng3.angle_index[0, p]), int(ng3.angle_index[1, p]))
+        for p in range(ng3.angle_index.shape[1])
+        if ng3.angle_mask[p]
+    }
+    assert got == ref

From 76475182b8996df0f2913a7a5e3a57e629502bc5 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 11:55:38 +0800
Subject: [PATCH 06/12] feat(dpmodel): angle->edge/node segment aggregation

---
 .../dpmodel/utils/neighbor_graph/__init__.py  |  4 ++
 deepmd/dpmodel/utils/neighbor_graph/angles.py | 47 +++++++++++++++++++
 .../common/dpmodel/test_angle_builder.py      | 47 +++++++++++++++++++
 3 files changed, 98 insertions(+)

diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index 03b83282e3..f741e0e3f6 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
@@ -10,6 +10,8 @@
 """
 
 from .angles import (
+    angle_to_edge_sum,
+    angle_to_node_sum,
     attach_angles,
     build_angle_index,
 )
@@ -50,6 +52,8 @@
 __all__ = [
     "GraphLayout",
     "NeighborGraph",
+    "angle_to_edge_sum",
+    "angle_to_node_sum",
     "attach_angles",
     "build_angle_index",
     "build_neighbor_graph",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/angles.py b/deepmd/dpmodel/utils/neighbor_graph/angles.py
index 5f6c7d9c72..7f922d453f 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/angles.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/angles.py
@@ -19,6 +19,7 @@
 
 from .graph import GraphLayout, NeighborGraph, pad_and_guard_angles
 from .pairs import center_edge_pairs
+from .segment import segment_sum
 
 
 def build_angle_index(
@@ -138,3 +139,49 @@ def attach_angles(
         layout=layout,
     )
     return dataclasses.replace(graph, angle_index=ai, angle_mask=am)
+
+
+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)
diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index f09a3349b5..7efae7e5ce 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -3,6 +3,8 @@
 
 from deepmd.dpmodel.utils.neighbor_graph import (
     GraphLayout,
+    angle_to_edge_sum,
+    angle_to_node_sum,
     attach_angles,
     build_angle_index,
     build_neighbor_graph,
@@ -311,3 +313,48 @@ def test_attach_angles_with_layout_node_capacity():
         if ng3.angle_mask[p]
     }
     assert got == ref
+
+
+# ---------------------------------------------------------------------------
+# Task 4: angle aggregation (angle_to_edge_sum / angle_to_node_sum)
+# ---------------------------------------------------------------------------
+
+
+def test_angle_aggregation():
+    """Test angle->edge and angle->node aggregation."""
+    # edges: dst=[0,0]; angles: (a=0,b=0),(a=0,b=1),(a=1,b=0)
+    edge_index = np.array([[5, 5], [0, 0]], dtype=np.int64)
+    angle_index = np.array([[0, 0, 1], [0, 1, 0]], dtype=np.int64)
+    data = np.array([1.0, 2.0, 4.0])  # per-angle
+    # angle->edge (group by edge_a): edge0 gets angles 0,1 => 3; edge1 gets angle2 => 4
+    e = angle_to_edge_sum(data, angle_index, 2)
+    np.testing.assert_allclose(e, [3.0, 4.0])
+    # angle->node (center of edge_a): all 3 angles share center 0 => 7
+    n = angle_to_node_sum(data, angle_index, edge_index, 1)
+    np.testing.assert_allclose(n, [7.0])
+
+
+def test_angle_aggregation_torch_namespace():
+    """Step 4b: torch-namespace smoke test for angle aggregation."""
+    import torch
+
+    # edges: dst=[0,0]; angles: (a=0,b=0),(a=0,b=1),(a=1,b=0)
+    edge_index = np.array([[5, 5], [0, 0]], dtype=np.int64)
+    angle_index = np.array([[0, 0, 1], [0, 1, 0]], dtype=np.int64)
+    data = np.array([1.0, 2.0, 4.0])  # per-angle
+
+    # numpy reference
+    e_np = angle_to_edge_sum(data, angle_index, 2)
+    n_np = angle_to_node_sum(data, angle_index, edge_index, 1)
+
+    # torch version
+    t_edge_index = torch.from_numpy(edge_index)
+    t_angle_index = torch.from_numpy(angle_index)
+    t_data = torch.from_numpy(data)
+
+    e_t = angle_to_edge_sum(t_data, t_angle_index, 2)
+    n_t = angle_to_node_sum(t_data, t_angle_index, t_edge_index, 1)
+
+    # compare
+    np.testing.assert_allclose(np.asarray(e_t), e_np)
+    np.testing.assert_allclose(np.asarray(n_t), n_np)

From 6ecafb96e9439831bcac5206ef6a4250cf51097c Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 12:04:16 +0800
Subject: [PATCH 07/12] feat(dpmodel): graph_angle_cos with dpa3 dense-parity

---
 .../dpmodel/utils/neighbor_graph/__init__.py  |   2 +
 deepmd/dpmodel/utils/neighbor_graph/angles.py |  29 ++
 .../dpmodel/test_graph_angle_cos_parity.py    | 314 ++++++++++++++++++
 3 files changed, 345 insertions(+)
 create mode 100644 source/tests/common/dpmodel/test_graph_angle_cos_parity.py

diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index f741e0e3f6..3f9f6def24 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
@@ -14,6 +14,7 @@
     angle_to_node_sum,
     attach_angles,
     build_angle_index,
+    graph_angle_cos,
 )
 from .ase_builder import (
     build_neighbor_graph_ase,
@@ -63,6 +64,7 @@
     "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",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/angles.py b/deepmd/dpmodel/utils/neighbor_graph/angles.py
index 7f922d453f..ebef42d72b 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/angles.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/angles.py
@@ -141,6 +141,35 @@ def attach_angles(
     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).
 
diff --git a/source/tests/common/dpmodel/test_graph_angle_cos_parity.py b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
new file mode 100644
index 0000000000..06ad8b1c80
--- /dev/null
+++ b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Parity tests for graph_angle_cos vs dpa3 dense cosine_ij and se_t dot form."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from deepmd.dpmodel.utils.neighbor_graph import (
+    attach_angles,
+    build_neighbor_graph,
+    graph_angle_cos,
+)
+
+
+# ---------------------------------------------------------------------------
+# Step 1/4: explicit geometry sanity
+# ---------------------------------------------------------------------------
+
+
+def test_graph_angle_cos_matches_normalized_dot():
+    """Center at origin, neighbors along x and y: cos=0 (perpendicular).
+
+    Use rcut=1.1 so neighbors (distance 1.0 from center) are included but
+    the neighbor-to-neighbor distance (sqrt(2)~1.41) is NOT, making atom 0
+    the ONLY center with multiple neighbors => exactly ONE unordered angle.
+    """
+    coord = np.array([[[0.0, 0, 0], [1.0, 0, 0], [0.0, 1.0, 0]]])
+    atype = np.array([[0, 0, 0]])
+    ng = attach_angles(build_neighbor_graph(coord, atype, None, 1.1), a_rcut=1.1)
+    cos = graph_angle_cos(ng.angle_index, ng.edge_vec)
+    real = np.asarray(ng.angle_mask)
+    vals = np.asarray(cos)[real]
+    # only atom 0 has 2 neighbors within rcut => exactly ONE unordered angle
+    assert int(real.sum()) == 1
+    # perpendicular vectors: dot = 0, (1-eps) scaling => val ~0.0 to atol 1e-6
+    np.testing.assert_allclose(vals, [0.0], atol=1e-6)
+
+
+def test_graph_angle_cos_parallel():
+    """Two neighbors in the SAME direction as seen from center: cos ~(1-eps).
+
+    Use rcut=1.5 so only center→neighbor1 (dist=1) and center→neighbor2 (dist=2)
+    but atom1→atom2 (dist=1) is NOT seen from atom2 as a second neighbor.
+    Actually place center at 1.0, so neighbors at 0.0 and 2.0 are collinear.
+    rcut=1.1 => center (at 1.0) sees atoms at 0.0 (dist=1.0) and 2.0 (dist=1.0),
+    but those two atoms don't see each other (dist=2.0 > 1.1).
+    """
+    coord = np.array([[[1.0, 0, 0], [0.0, 0, 0], [2.0, 0, 0]]])
+    atype = np.array([[0, 0, 0]])
+    ng = attach_angles(build_neighbor_graph(coord, atype, None, 1.1), a_rcut=1.1)
+    cos = graph_angle_cos(ng.angle_index, ng.edge_vec)
+    real = np.asarray(ng.angle_mask)
+    vals = np.asarray(cos)[real]
+    assert int(real.sum()) == 1
+    # edge_vec = neighbor - center; both neighbors are in opposite x-directions
+    # from the center: edge to atom1=(0,0,0) is (-1,0,0); edge to atom2=(2,0,0) is (+1,0,0)
+    # For unit-norm vectors (norm=1.0):
+    #   na = va / (1 + eps),  nb = vb / (1 + eps)
+    #   dot(na, nb) = -1 / (1+eps)^2
+    #   cos = -1 / (1+eps)^2 * (1 - eps)
+    eps = 1e-6
+    expected = -1.0 / (1 + eps) ** 2 * (1 - eps)
+    np.testing.assert_allclose(vals, [expected], rtol=1e-6)
+
+
+def test_graph_angle_cos_no_self_angles():
+    """No graph angle should have both edge slots pointing to the same edge."""
+    coord = np.array([[[0.0, 0, 0], [1.0, 0, 0], [0.0, 1.0, 0], [0.0, 0.0, 1.0]]])
+    atype = np.array([[0, 0, 0, 0]])
+    ng = attach_angles(build_neighbor_graph(coord, atype, None, 3.0), a_rcut=3.0)
+    ai = np.asarray(ng.angle_index)
+    am = np.asarray(ng.angle_mask)
+    # For every VALID angle, the two edge indices must differ
+    real_angles = [(int(ai[0, p]), int(ai[1, p])) for p in range(am.shape[0]) if am[p]]
+    for a, b in real_angles:
+        assert a != b, f"Self-angle found: edge {a}"
+
+
+# ---------------------------------------------------------------------------
+# dpa3 dense parity
+# ---------------------------------------------------------------------------
+
+
+def _dense_cosine_ij(coord_ext, nlist, a_rcut, a_sel):
+    """Faithful numpy transcription of repflows.py:597-649.
+
+    Returns cosine_ij of shape (nf, nloc, a_sel, a_sel).
+    a_diff convention: coord_r - coord_l = neighbor - center.
+    """
+    nf, nloc, nnei = nlist.shape
+    # coord_ext: (nf, nall, 3)
+    # a_nlist: truncate to a_sel and mask beyond a_rcut
+    diff_full = np.zeros((nf, nloc, nnei, 3), dtype=np.float64)
+    for f in range(nf):
+        for i in range(nloc):
+            for k in range(nnei):
+                j = nlist[f, i, k]
+                if j >= 0:
+                    diff_full[f, i, k] = coord_ext[f, j] - coord_ext[f, i]
+    # a_rcut gate — clip a_sel to actual nnei columns available
+    nnei = nlist.shape[2]
+    eff_a_sel = min(a_sel, nnei)
+    a_dist_mask = (np.linalg.norm(diff_full, axis=-1) < a_rcut)[:, :, :eff_a_sel]
+    a_nlist = nlist[:, :, :eff_a_sel].copy()
+    a_nlist = np.where(a_dist_mask, a_nlist, np.full_like(a_nlist, -1))
+    # a_diff: shape (nf, nloc, eff_a_sel, 3)
+    a_diff = np.zeros((nf, nloc, eff_a_sel, 3), dtype=np.float64)
+    for f in range(nf):
+        for i in range(nloc):
+            for k in range(eff_a_sel):
+                j = a_nlist[f, i, k]
+                if j >= 0:
+                    a_diff[f, i, k] = coord_ext[f, j] - coord_ext[f, i]
+    # normalized_diff_i: (nf, nloc, eff_a_sel, 3)
+    norm = np.linalg.norm(a_diff, axis=-1, keepdims=True)
+    normalized_diff_i = a_diff / (norm + 1e-6)
+    # cosine_ij: (nf, nloc, eff_a_sel, eff_a_sel)
+    cosine_ij = np.matmul(normalized_diff_i, np.swapaxes(normalized_diff_i, -2, -1))
+    cosine_ij = cosine_ij * (1 - 1e-6)
+    return cosine_ij, a_nlist
+
+
+def test_graph_angle_cos_parity_vs_dpa3_dense():
+    """Graph unordered/no-self cos values must match dense OFF-DIAGONAL cosine_ij.
+
+    Uses a small 4-atom system with a large a_sel (non-binding) so that the
+    graph and dense see the same neighbor set.
+
+    - Tolerance: rtol=1e-12, atol=1e-12 (CPU fp64 same-math, identical eps).
+    - The graph is UNORDERED (no duplicates) and NO-SELF.
+    - Dense diagonal (j==k, cos≈1) must NOT appear in graph angle set.
+    - Dense off-diagonal (j!=k) are collected as unordered {j,k} pairs and
+      matched against graph angles by (center_node, unordered edge-src pair).
+    """
+    rng = np.random.default_rng(42)
+    # 4 atoms in a box; no PBC => set box to None
+    # Use a single frame, 4 atoms
+    nf, nloc = 1, 4
+    coord = rng.uniform(-1, 1, (nf, nloc, 3))
+    atype = np.zeros((nf, nloc), dtype=np.int32)
+    rcut = 3.0
+    a_rcut = 2.5
+    # Choose a_sel equal to nloc-1 (max neighbors) => non-binding
+    a_sel = nloc - 1  # =3; each atom has at most 3 neighbors => non-binding
+
+    # --- graph side ---
+    ng = attach_angles(build_neighbor_graph(coord, atype, None, rcut), a_rcut=a_rcut)
+    cos_graph = np.asarray(graph_angle_cos(ng.angle_index, ng.edge_vec))
+    am = np.asarray(ng.angle_mask)
+    ai = np.asarray(ng.angle_index)
+    ei = np.asarray(ng.edge_index)  # (2, E): [src, dst]
+    ev = np.asarray(ng.edge_vec)
+
+    # No self-angles
+    for p in range(am.shape[0]):
+        if am[p]:
+            assert int(ai[0, p]) != int(ai[1, p]), "Self-angle found in graph"
+
+    # --- dense side ---
+    # Build a dense nlist from the same coord (no PBC, full nlist)
+    # We construct nlist by brute force
+    coord3 = coord[0]  # (nloc, 3)
+    # For the dense side, coord_ext = coord (nloc=nall, no ghosts)
+    coord_ext = coord  # (1, nloc, 3)
+    # Build dense nlist: shape (1, nloc, nnei_max)
+    # For 4 atoms, each atom has at most nloc-1=3 neighbors
+    nnei = nloc - 1  # max possible neighbors (no self)
+    dense_nlist = np.full((nf, nloc, nnei), -1, dtype=np.int64)
+    for i in range(nloc):
+        k = 0
+        for j in range(nloc):
+            d = np.linalg.norm(coord3[j] - coord3[i])
+            if d < rcut and i != j:
+                dense_nlist[0, i, k] = j
+                k += 1
+
+    cosine_ij, a_nlist = _dense_cosine_ij(coord_ext, dense_nlist, a_rcut, a_sel)
+    # cosine_ij: (1, nloc, a_sel, a_sel)
+
+    # --- match graph angles to dense off-diagonal ---
+    # graph edge_index: src=edge_index[0], dst=edge_index[1] (center=dst)
+    # For each valid graph angle p: edge_a=ai[0,p], edge_b=ai[1,p]
+    # center = ei[1, edge_a] = ei[1, edge_b] (shared center)
+    # neighbor_a = ei[0, edge_a], neighbor_b = ei[0, edge_b]
+
+    # Build dense lookup: (center, na, nb) -> cos for off-diagonal
+    dense_cos_lookup = {}  # (center, unordered frozenset(na, nb)) -> cos
+    for i in range(nloc):
+        for j_idx in range(a_sel):
+            na = int(a_nlist[0, i, j_idx])
+            if na < 0:
+                continue
+            for k_idx in range(a_sel):
+                nb = int(a_nlist[0, i, k_idx])
+                if nb < 0:
+                    continue
+                if j_idx == k_idx:  # skip diagonal (self-angles = edge channel)
+                    continue
+                cos_val = float(cosine_ij[0, i, j_idx, k_idx])
+                key = (i, frozenset([na, nb]))
+                # unordered: both (j,k) and (k,j) map to same pair
+                # store the value for frozenset key (they should be equal by symmetry
+                # of cos, but we verify)
+                if key not in dense_cos_lookup:
+                    dense_cos_lookup[key] = cos_val
+                else:
+                    # cosine is symmetric: both directions should be equal
+                    np.testing.assert_allclose(
+                        dense_cos_lookup[key],
+                        cos_val,
+                        atol=1e-14,
+                        err_msg=f"Dense cosine not symmetric for ({i},{na},{nb})",
+                    )
+
+    # Now compare each valid graph angle
+    for p in range(am.shape[0]):
+        if not am[p]:
+            continue
+        ea = int(ai[0, p])
+        eb = int(ai[1, p])
+        center = int(ei[1, ea])
+        assert center == int(ei[1, eb]), "Angle edges don't share center"
+        na = int(ei[0, ea])
+        nb = int(ei[0, eb])
+        key = (center, frozenset([na, nb]))
+        assert key in dense_cos_lookup, (
+            f"Graph angle (center={center}, na={na}, nb={nb}) not in dense"
+        )
+        cos_g = float(cos_graph[p])
+        cos_d = dense_cos_lookup[key]
+        np.testing.assert_allclose(
+            cos_g,
+            cos_d,
+            rtol=1e-12,
+            atol=1e-12,
+            err_msg=f"cos mismatch at (center={center}, na={na}, nb={nb})",
+        )
+
+
+# ---------------------------------------------------------------------------
+# se_t dot-product cross-check
+# ---------------------------------------------------------------------------
+
+
+def test_matches_se_t_dot_form():
+    """Cos * |va| * |vb| ≈ va·vb (the se_t unnormalized inner product).
+
+    se_t.py:428-437 uses `rr_i·rr_j` (unnormalized 3D vectors = edge_vec
+    without smoothing weights).  graph_angle_cos(eps=1e-6) satisfies:
+
+        cos ≈ (va/‖va‖) · (vb/‖vb‖) * (1-eps)
+
+    So:
+        va·vb ≈ cos * (‖va‖ + eps) * (‖vb‖ + eps) / (1 - eps)
+
+    We verify this identity up to the eps rounding (atol~1e-6 * ‖va‖*‖vb‖).
+    """
+    coord = np.array([[[0.0, 0, 0], [1.0, 0.5, 0.3], [0.2, 1.0, 0.8]]])
+    atype = np.array([[0, 0, 0]])
+    ng = attach_angles(build_neighbor_graph(coord, atype, None, 5.0), a_rcut=5.0)
+    ai = np.asarray(ng.angle_index)
+    am = np.asarray(ng.angle_mask)
+    ev = np.asarray(ng.edge_vec)
+    cos = np.asarray(graph_angle_cos(ng.angle_index, ng.edge_vec))
+
+    eps = 1e-6
+    for p in range(am.shape[0]):
+        if not am[p]:
+            continue
+        ea = int(ai[0, p])
+        eb = int(ai[1, p])
+        va = ev[ea]
+        vb = ev[eb]
+        # unnormalized dot product (se_t convention)
+        dot_se_t = float(np.dot(va, vb))
+        # reconstruct from graph cos
+        na_norm = np.linalg.norm(va)
+        nb_norm = np.linalg.norm(vb)
+        dot_reconstructed = (
+            float(cos[p]) * (na_norm + eps) * (nb_norm + eps) / (1 - eps)
+        )
+        # tolerance: eps * ‖va‖*‖vb‖ dominates
+        tol = eps * max(na_norm * nb_norm, 1e-12)
+        np.testing.assert_allclose(
+            dot_reconstructed,
+            dot_se_t,
+            atol=tol,
+            err_msg=f"se_t dot mismatch for angle {p}: va={va}, vb={vb}",
+        )
+
+
+# ---------------------------------------------------------------------------
+# torch namespace smoke test (TID253)
+# ---------------------------------------------------------------------------
+
+
+def test_graph_angle_cos_torch_matches_numpy():
+    """graph_angle_cos on torch tensors matches numpy output (array-API compat)."""
+    import torch
+
+    coord = np.array([[[0.0, 0, 0], [1.0, 0, 0], [0.0, 1.0, 0], [1.0, 1.0, 0]]])
+    atype = np.array([[0, 0, 0, 0]])
+    ng = attach_angles(build_neighbor_graph(coord, atype, None, 3.0), a_rcut=3.0)
+    angle_index_np = np.asarray(ng.angle_index)
+    edge_vec_np = np.asarray(ng.edge_vec)
+
+    cos_np = np.asarray(graph_angle_cos(angle_index_np, edge_vec_np))
+
+    angle_index_t = torch.from_numpy(angle_index_np)
+    edge_vec_t = torch.from_numpy(edge_vec_np)
+    cos_t = graph_angle_cos(angle_index_t, edge_vec_t)
+    cos_t_np = cos_t.numpy()
+
+    np.testing.assert_allclose(cos_t_np, cos_np, rtol=1e-14, atol=1e-14)

From df28cdc217ebf1ab21a613b18c1f88c015682721 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 12:59:04 +0800
Subject: [PATCH 08/12] test(dpmodel): real se_t-convention oracle for
 graph_angle_cos; fix antiparallel test name

---
 .../dpmodel/test_graph_angle_cos_parity.py    | 123 ++++++++++++------
 1 file changed, 80 insertions(+), 43 deletions(-)

diff --git a/source/tests/common/dpmodel/test_graph_angle_cos_parity.py b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
index 06ad8b1c80..09d3852db1 100644
--- a/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
+++ b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
@@ -36,14 +36,15 @@ def test_graph_angle_cos_matches_normalized_dot():
     np.testing.assert_allclose(vals, [0.0], atol=1e-6)
 
 
-def test_graph_angle_cos_parallel():
-    """Two neighbors in the SAME direction as seen from center: cos ~(1-eps).
-
-    Use rcut=1.5 so only center→neighbor1 (dist=1) and center→neighbor2 (dist=2)
-    but atom1→atom2 (dist=1) is NOT seen from atom2 as a second neighbor.
-    Actually place center at 1.0, so neighbors at 0.0 and 2.0 are collinear.
-    rcut=1.1 => center (at 1.0) sees atoms at 0.0 (dist=1.0) and 2.0 (dist=1.0),
-    but those two atoms don't see each other (dist=2.0 > 1.1).
+def test_graph_angle_cos_antiparallel():
+    """Two neighbors in ANTIPARALLEL directions as seen from the center: cos ≈ -1.
+
+    Center is placed at (1,0,0) with one neighbor at (0,0,0) and another at
+    (2,0,0).  The edge vectors from center are (-1,0,0) and (+1,0,0), which
+    point in opposite directions => cos ≈ -(1-eps).
+
+    rcut=1.1: center sees both neighbors (dist=1.0 each); the two neighbors
+    do not see each other (dist=2.0 > 1.1), so exactly ONE angle is formed.
     """
     coord = np.array([[[1.0, 0, 0], [0.0, 0, 0], [2.0, 0, 0]]])
     atype = np.array([[0, 0, 0]])
@@ -52,12 +53,8 @@ def test_graph_angle_cos_parallel():
     real = np.asarray(ng.angle_mask)
     vals = np.asarray(cos)[real]
     assert int(real.sum()) == 1
-    # edge_vec = neighbor - center; both neighbors are in opposite x-directions
-    # from the center: edge to atom1=(0,0,0) is (-1,0,0); edge to atom2=(2,0,0) is (+1,0,0)
-    # For unit-norm vectors (norm=1.0):
-    #   na = va / (1 + eps),  nb = vb / (1 + eps)
-    #   dot(na, nb) = -1 / (1+eps)^2
-    #   cos = -1 / (1+eps)^2 * (1 - eps)
+    # edge_vec = neighbor - center; edge to (0,0,0) is (-1,0,0); edge to (2,0,0) is (+1,0,0).
+    # Antiparallel unit vectors: dot(na, nb) = -1 / (1+eps)^2, scaled by (1-eps).
     eps = 1e-6
     expected = -1.0 / (1 + eps) ** 2 * (1 - eps)
     np.testing.assert_allclose(vals, [expected], rtol=1e-6)
@@ -243,49 +240,89 @@ def test_graph_angle_cos_parity_vs_dpa3_dense():
 
 
 def test_matches_se_t_dot_form():
-    """Cos * |va| * |vb| ≈ va·vb (the se_t unnormalized inner product).
+    """Cross-check graph_angle_cos against an independent coordinate-based oracle.
+
+    se_t.py:428-437 computes ``env_ij = sum(rr_i * rr_j, -1)`` where
+    ``rr_i = sw * diff / r^2`` (the 3-D columns of the env-mat).  The raw
+    unnormalized dot product ``va · vb`` (with ``va = r_a - r_center``) is the
+    numerator that graph_angle_cos normalizes:
 
-    se_t.py:428-437 uses `rr_i·rr_j` (unnormalized 3D vectors = edge_vec
-    without smoothing weights).  graph_angle_cos(eps=1e-6) satisfies:
+        graph_angle_cos = (1 - eps) * (va · vb) / ((|va| + eps) * (|vb| + eps))
 
-        cos ≈ (va/‖va‖) · (vb/‖vb‖) * (1-eps)
+    Inverting:
 
-    So:
-        va·vb ≈ cos * (‖va‖ + eps) * (‖vb‖ + eps) / (1 - eps)
+        graph_angle_cos * (|va| + eps) * (|vb| + eps) / (1 - eps) = va · vb
 
-    We verify this identity up to the eps rounding (atol~1e-6 * ‖va‖*‖vb‖).
+    **Why sw is factored out**: sw scales each env-mat vector by a scalar.
+    When all neighbor distances are *below* ``rcut_smth``, the smooth switch
+    function equals 1 exactly (``sw == 1``), so the sw factor contributes
+    nothing and ``env_ij`` reduces to the plain geometry.
+
+    **Why this test is not tautological**: the reference ``va``, ``vb``, and
+    ``env_ij = va · vb`` are computed DIRECTLY FROM COORDINATES in plain numpy,
+    independent of the graph code path.  The |va| and |vb| norms used to unwind
+    ``cos`` are also recomputed from coordinates, NOT read from ``edge_vec``.
+    This verifies that the graph stores ``edge_vec = neighbor - center``
+    correctly and that ``graph_angle_cos`` faithfully encodes the geometry.
+    With distances well below ``rcut_smth`` the identity holds to ``rtol=1e-12``
+    because it is exact algebra over fp64; eps-induced rounding is negligible
+    compared to fp64 relative precision.
     """
-    coord = np.array([[[0.0, 0, 0], [1.0, 0.5, 0.3], [0.2, 1.0, 0.8]]])
+    # All atoms within distance 0.5 of center; rcut_smth = 1.0 so sw == 1 for all.
+    rng = np.random.default_rng(7)
+    center = np.array([0.0, 0.0, 0.0])
+    r_a = rng.uniform(0.1, 0.4, 3)  # distance from center < rcut_smth=1.0
+    r_b = rng.uniform(0.1, 0.4, 3)
+    coord = np.array([[[*center], [*r_a], [*r_b]]])  # (1, 3, 3), single frame
     atype = np.array([[0, 0, 0]])
-    ng = attach_angles(build_neighbor_graph(coord, atype, None, 5.0), a_rcut=5.0)
+
+    rcut = 2.0
+    a_rcut = 2.0
+    ng = attach_angles(build_neighbor_graph(coord, atype, None, rcut), a_rcut=a_rcut)
     ai = np.asarray(ng.angle_index)
     am = np.asarray(ng.angle_mask)
-    ev = np.asarray(ng.edge_vec)
+    ei = np.asarray(ng.edge_index)
     cos = np.asarray(graph_angle_cos(ng.angle_index, ng.edge_vec))
 
     eps = 1e-6
-    for p in range(am.shape[0]):
-        if not am[p]:
-            continue
+
+    # At least one valid angle must exist (atom 0 is the only center with ≥2 nei)
+    valid_angles = [p for p in range(am.shape[0]) if am[p]]
+    assert len(valid_angles) >= 1, "No valid angles found — geometry problem"
+
+    for p in valid_angles:
         ea = int(ai[0, p])
         eb = int(ai[1, p])
-        va = ev[ea]
-        vb = ev[eb]
-        # unnormalized dot product (se_t convention)
-        dot_se_t = float(np.dot(va, vb))
-        # reconstruct from graph cos
-        na_norm = np.linalg.norm(va)
-        nb_norm = np.linalg.norm(vb)
-        dot_reconstructed = (
-            float(cos[p]) * (na_norm + eps) * (nb_norm + eps) / (1 - eps)
-        )
-        # tolerance: eps * ‖va‖*‖vb‖ dominates
-        tol = eps * max(na_norm * nb_norm, 1e-12)
+        center_node = int(ei[1, ea])
+        na_node = int(ei[0, ea])
+        nb_node = int(ei[0, eb])
+
+        # Reference: compute difference vectors FROM COORDINATES (independent of graph)
+        r_center = coord[0, center_node]
+        r_na = coord[0, na_node]
+        r_nb = coord[0, nb_node]
+        va_ref = r_na - r_center  # (3,)
+        vb_ref = r_nb - r_center  # (3,)
+
+        # Reference: unnormalized dot product from coordinates (se_t convention)
+        env_ij_ref = float(np.dot(va_ref, vb_ref))
+
+        # Reference norms — from coordinates, NOT from edge_vec
+        na_norm = float(np.linalg.norm(va_ref))
+        nb_norm = float(np.linalg.norm(vb_ref))
+
+        # Graph: unwind graph_angle_cos back to the unnormalized dot product
+        env_ij_graph = float(cos[p]) * (na_norm + eps) * (nb_norm + eps) / (1 - eps)
+
         np.testing.assert_allclose(
-            dot_reconstructed,
-            dot_se_t,
-            atol=tol,
-            err_msg=f"se_t dot mismatch for angle {p}: va={va}, vb={vb}",
+            env_ij_graph,
+            env_ij_ref,
+            rtol=1e-12,
+            err_msg=(
+                f"se_t dot mismatch at angle {p}: "
+                f"center={center_node}, na={na_node}, nb={nb_node}, "
+                f"va={va_ref}, vb={vb_ref}"
+            ),
         )
 
 

From d7628833a79bbcc5ae8cf31f1f99ff83f9d6fb57 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 13:14:21 +0800
Subject: [PATCH 09/12] feat(dpmodel): angle-force invariance test +
 padding-fraction report

---
 .../dpmodel/utils/neighbor_graph/__init__.py  |  2 +
 deepmd/dpmodel/utils/neighbor_graph/angles.py | 26 +++++++
 .../common/dpmodel/test_angle_builder.py      | 74 +++++++++++++++++++
 3 files changed, 102 insertions(+)

diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index 3f9f6def24..ca524991d2 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
@@ -10,6 +10,7 @@
 """
 
 from .angles import (
+    angle_padding_fraction,
     angle_to_edge_sum,
     angle_to_node_sum,
     attach_angles,
@@ -53,6 +54,7 @@
 __all__ = [
     "GraphLayout",
     "NeighborGraph",
+    "angle_padding_fraction",
     "angle_to_edge_sum",
     "angle_to_node_sum",
     "attach_angles",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/angles.py b/deepmd/dpmodel/utils/neighbor_graph/angles.py
index ebef42d72b..e684221945 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/angles.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/angles.py
@@ -214,3 +214,29 @@ def angle_to_node_sum(
     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
diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index 7efae7e5ce..1e817838bd 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -3,11 +3,13 @@
 
 from deepmd.dpmodel.utils.neighbor_graph import (
     GraphLayout,
+    angle_padding_fraction,
     angle_to_edge_sum,
     angle_to_node_sum,
     attach_angles,
     build_angle_index,
     build_neighbor_graph,
+    edge_force_virial,
     pad_and_guard_angles,
 )
 
@@ -358,3 +360,75 @@ def test_angle_aggregation_torch_namespace():
     # compare
     np.testing.assert_allclose(np.asarray(e_t), e_np)
     np.testing.assert_allclose(np.asarray(n_t), n_np)
+
+
+# ---------------------------------------------------------------------------
+# Task 6: angle-force invariance + angle_padding_fraction
+# ---------------------------------------------------------------------------
+
+
+def _small_graph_with_angles(a_rcut: float, layout: GraphLayout | None = None):
+    """Return (graph_no_angles, graph_with_angles) for a 3-atom, 1-frame system."""
+    # 3 atoms: 0,1,2 in a line along x; shape (nf=1, nloc=3, 3) / (nf=1, nloc=3)
+    coord = np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]])
+    atype = np.array([[0, 1, 0]], dtype=np.int64)
+    box = np.array([[[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]])
+    g = build_neighbor_graph(coord, atype, box, rcut=3.0)
+    g_with = attach_angles(g, a_rcut, layout=layout)
+    return g, g_with
+
+
+def test_edge_force_virial_ignores_angles():
+    """edge_force_virial output must be bit-identical with or without angles.
+
+    Angles add topology (angle_index/angle_mask) to the NeighborGraph but do
+    NOT change edge_vec, edge_index, or edge_mask — the only inputs to
+    edge_force_virial.  This test proves that the angle fields are truly
+    transparent to the force/virial assembly.
+    """
+    g_bare, g_with_angles = _small_graph_with_angles(a_rcut=1.5)
+
+    # Manufacture a fake per-edge gradient (same shape as edge_vec)
+    rng = np.random.default_rng(42)
+    n_edges = int(g_bare.edge_index.shape[1])
+    g_e = rng.standard_normal((n_edges, 3))
+
+    def run(graph):
+        return edge_force_virial(
+            g_e,
+            graph.edge_vec,
+            graph.edge_index,
+            graph.edge_mask,
+            graph.n_node,
+        )
+
+    force_bare, av_bare, vir_bare = run(g_bare)
+    force_with, av_with, vir_with = run(g_with_angles)
+
+    # Exact equality: same inputs → same computation → identical bits
+    np.testing.assert_array_equal(force_bare, force_with)
+    np.testing.assert_array_equal(av_bare, av_with)
+    np.testing.assert_array_equal(vir_bare, vir_with)
+
+
+def test_angle_padding_fraction():
+    """angle_padding_fraction returns 1 - A_real/A_max for a static layout.
+
+    We build with a fixed angle_capacity=A_max so the fraction is deterministic
+    (not influenced by the dynamic min_angles guard of pad_and_guard_angles).
+    """
+    A_max = 20  # static capacity, larger than any real angle count
+    layout = GraphLayout(angle_capacity=A_max)
+    g_bare, g_with = _small_graph_with_angles(a_rcut=1.5, layout=layout)
+
+    # Confirm angles are present
+    assert g_with.angle_mask is not None
+    A_real = int(np.sum(g_with.angle_mask))
+    assert 0 < A_real <= A_max, f"Expected 0 < A_real <= {A_max}, got {A_real}"
+
+    expected = 1.0 - A_real / A_max
+    got = angle_padding_fraction(g_with)
+    assert got == pytest.approx(expected), f"got {got}, expected {expected}"
+
+    # No angles → fraction is 0.0
+    assert angle_padding_fraction(g_bare) == 0.0

From 082798c307b1d0b4e32d3230b37a5c6044dfca12 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jul 2026 13:18:29 +0800
Subject: [PATCH 10/12] test(dpmodel): angle_padding_fraction total==0 branch

---
 .../common/dpmodel/test_angle_builder.py      | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index 1e817838bd..d6975d087e 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -432,3 +432,24 @@ def test_angle_padding_fraction():
 
     # No angles → fraction is 0.0
     assert angle_padding_fraction(g_bare) == 0.0
+
+
+def test_angle_padding_fraction_total_zero():
+    """angle_padding_fraction returns 0.0 when angle_mask.shape[0] == 0.
+
+    This exercises the `if total == 0: return 0.0` branch in angle_padding_fraction.
+    A graph with angle_capacity=0 and a_rcut too small for any angles produces
+    angle_mask with shape (0,).
+    """
+    # Create a layout with angle_capacity=0 and use a_rcut=0.01 (too small
+    # for any edge in the fixture to pass the distance gate)
+    layout = GraphLayout(angle_capacity=0)
+    g_bare, g_with = _small_graph_with_angles(a_rcut=0.01, layout=layout)
+
+    # Verify angle_mask exists and has shape (0,)
+    assert g_with.angle_mask is not None
+    assert g_with.angle_mask.shape[0] == 0
+
+    # angle_padding_fraction must return 0.0 for empty mask
+    result = angle_padding_fraction(g_with)
+    assert result == 0.0

From dd58c021eb30b087676dd230353e43cfe2bc3975 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 3 Jul 2026 05:52:30 +0000
Subject: [PATCH 11/12] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/dpmodel/utils/neighbor_graph/angles.py | 22 ++++++++++++++-----
 .../common/dpmodel/test_angle_builder.py      |  1 +
 .../dpmodel/test_graph_angle_cos_parity.py    |  5 +++--
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/deepmd/dpmodel/utils/neighbor_graph/angles.py b/deepmd/dpmodel/utils/neighbor_graph/angles.py
index e684221945..be208986ed 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/angles.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/angles.py
@@ -6,9 +6,13 @@
 center_edge_pairs; a_rcut filters the participating edges.
 """
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
-from typing import TYPE_CHECKING
+from typing import (
+    TYPE_CHECKING,
+)
 
 import array_api_compat
 
@@ -17,9 +21,17 @@
 
 import dataclasses
 
-from .graph import GraphLayout, NeighborGraph, pad_and_guard_angles
-from .pairs import center_edge_pairs
-from .segment import segment_sum
+from .graph import (
+    GraphLayout,
+    NeighborGraph,
+    pad_and_guard_angles,
+)
+from .pairs import (
+    center_edge_pairs,
+)
+from .segment import (
+    segment_sum,
+)
 
 
 def build_angle_index(
diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index d6975d087e..17a78862d9 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -1,3 +1,4 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
 import numpy as np
 import pytest
 
diff --git a/source/tests/common/dpmodel/test_graph_angle_cos_parity.py b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
index 09d3852db1..8040141919 100644
--- a/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
+++ b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Parity tests for graph_angle_cos vs dpa3 dense cosine_ij and se_t dot form."""
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
 import numpy as np
 
@@ -11,7 +13,6 @@
     graph_angle_cos,
 )
 
-
 # ---------------------------------------------------------------------------
 # Step 1/4: explicit geometry sanity
 # ---------------------------------------------------------------------------

From 6c2c9223f26b178851c564a9b3098bbd8a7f7383 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Tue, 7 Jul 2026 16:19:46 +0800
Subject: [PATCH 12/12] test(dpmodel): address review comments on angle tests

- Remove unused locals flagged by CodeQL (ai_def/ai_full in the
  ordered-superset test, ev in the angle-cos parity test); strengthen the
  ordered+include_self test to assert the default pairs are an actual
  subset of the full pair set (using the previously-dead ai_def/ai_full).
- Add masked-padding aggregation tests (CodeRabbit nitpick): pin the
  precondition that angle_to_edge_sum/angle_to_node_sum ignore guard
  angles only once data is zeroed via angle_mask, in numpy and torch.
---
 .../common/dpmodel/test_angle_builder.py      | 63 +++++++++++++++++++
 .../dpmodel/test_graph_angle_cos_parity.py    |  1 -
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/source/tests/common/dpmodel/test_angle_builder.py b/source/tests/common/dpmodel/test_angle_builder.py
index 17a78862d9..d1e460ad46 100644
--- a/source/tests/common/dpmodel/test_angle_builder.py
+++ b/source/tests/common/dpmodel/test_angle_builder.py
@@ -292,6 +292,18 @@ def test_attach_angles_ordered_include_self():
     n_full = int(am_full.sum())
     # ordered+include_self must produce at least as many angles as default
     assert n_full >= n_default
+    # ... and the full set must contain every default (edge_a, edge_b) pair
+    default_pairs = {
+        (int(ai_def[0, p]), int(ai_def[1, p]))
+        for p in range(ai_def.shape[1])
+        if am_def[p]
+    }
+    full_pairs = {
+        (int(ai_full[0, p]), int(ai_full[1, p]))
+        for p in range(ai_full.shape[1])
+        if am_full[p]
+    }
+    assert default_pairs <= full_pairs
 
 
 def test_attach_angles_with_layout_node_capacity():
@@ -363,6 +375,57 @@ def test_angle_aggregation_torch_namespace():
     np.testing.assert_allclose(np.asarray(n_t), n_np)
 
 
+def test_angle_aggregation_ignores_padding_when_masked():
+    """Padded angle slots must not pollute the sums once ``data`` is masked.
+
+    ``angle_to_edge_sum`` / ``angle_to_node_sum`` do not apply ``angle_mask``
+    internally: guard angles point at edge ``pad_value`` (an in-range real
+    edge), so their ``data`` would otherwise land on that edge/node. The
+    precondition is that callers zero ``data`` on padded slots (``data *
+    angle_mask``). This test pins that: with non-zero padding data, the masked
+    aggregation reproduces the real-only result, while the unmasked one does
+    not.
+    """
+    edge_index = np.array([[5, 5], [0, 0]], dtype=np.int64)
+    real_ai = np.array([[0, 0, 1], [0, 1, 0]], dtype=np.int64)  # 3 real angles
+    # pad to capacity 5 => 2 guard angles pointing at edge pad_value=0
+    ai, mask = pad_and_guard_angles(real_ai, angle_capacity=5, pad_value=0)
+    np.testing.assert_array_equal(mask, [True, True, True, False, False])
+    # non-zero data in the padded slots (would corrupt edge 0 / node 0 if used)
+    data = np.array([1.0, 2.0, 4.0, 99.0, 88.0])
+    masked = data * mask.astype(data.dtype)
+
+    # masked aggregation == real-only reference [3, 4] / [7]
+    e = angle_to_edge_sum(masked, ai, 2)
+    n = angle_to_node_sum(masked, ai, edge_index, 1)
+    np.testing.assert_allclose(e, [3.0, 4.0])
+    np.testing.assert_allclose(n, [7.0])
+
+    # sanity: without masking the padding DOES leak into edge/node 0
+    e_bad = angle_to_edge_sum(data, ai, 2)
+    assert not np.allclose(e_bad, [3.0, 4.0])
+
+
+def test_angle_aggregation_ignores_padding_torch_namespace():
+    """torch-namespace smoke test for the masked-padding aggregation."""
+    import torch
+
+    edge_index = np.array([[5, 5], [0, 0]], dtype=np.int64)
+    real_ai = np.array([[0, 0, 1], [0, 1, 0]], dtype=np.int64)
+    ai, mask = pad_and_guard_angles(real_ai, angle_capacity=5, pad_value=0)
+    data = np.array([1.0, 2.0, 4.0, 99.0, 88.0])
+    masked = data * mask.astype(data.dtype)
+
+    t_edge_index = torch.from_numpy(edge_index)
+    t_ai = torch.from_numpy(ai)
+    t_masked = torch.from_numpy(masked)
+
+    e_t = angle_to_edge_sum(t_masked, t_ai, 2)
+    n_t = angle_to_node_sum(t_masked, t_ai, t_edge_index, 1)
+    np.testing.assert_allclose(np.asarray(e_t), [3.0, 4.0])
+    np.testing.assert_allclose(np.asarray(n_t), [7.0])
+
+
 # ---------------------------------------------------------------------------
 # Task 6: angle-force invariance + angle_padding_fraction
 # ---------------------------------------------------------------------------
diff --git a/source/tests/common/dpmodel/test_graph_angle_cos_parity.py b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
index 8040141919..40e58ba82f 100644
--- a/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
+++ b/source/tests/common/dpmodel/test_graph_angle_cos_parity.py
@@ -147,7 +147,6 @@ def test_graph_angle_cos_parity_vs_dpa3_dense():
     am = np.asarray(ng.angle_mask)
     ai = np.asarray(ng.angle_index)
     ei = np.asarray(ng.edge_index)  # (2, E): [src, dst]
-    ev = np.asarray(ng.edge_vec)
 
     # No self-angles
     for p in range(am.shape[0]):