feat(dpmodel): pad_and_guard_angles angle-axis padder

Author: Han Wang

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

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)``.
 

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)