feat(dpmodel): attach_angles post-hoc angle attachment

Author: Han Wang

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",

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)

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