test(dpmodel): build_angle_index multi-center, static-layout, ordered-no-self coverage

Author: Han Wang

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