deepmd.dpmodel.descriptor.dpa4_nn.indexing

deepmd.dpmodel.descriptor.dpa4_nn.indexing#

SO(3) packed-index and projection helpers for DPA4/SeZM.

This module defines the packed (l, m) indexing helpers and the projection utilities used by the DPA4 equivariant operators. It is the dpmodel port of deepmd.pt.model.descriptor.sezm_nn.indexing.

The index-table builders run at module-init time on static index data and are implemented in plain numpy by design (not array-API); they return np.int64 arrays. The torch-specific device/dtype keyword parameters of the pt versions are dropped for those builders. Only project_D_to_m and project_Dt_from_m operate on runtime tensors and are array-API compatible.

Functions#

`get_so3_dim_of_lmax`(→ int)	Return SO(3) representation dimension for given lmax.
`map_degree_idx`(→ numpy.ndarray)	Build degree (l) index for each position in the packed (l, m) layout.
`build_gie_zonal_index`(→ tuple[numpy.ndarray, ...)	Build node-level packed indices for GIE zonal coupling.
`project_D_to_m`(→ Any)	Row-project block-diagonal Wigner-D to the m-major truncated layout.
`project_Dt_from_m`(→ Any)	Column-project block-diagonal Wigner-D^T for inverse rotation.
`so3_packed_index`(→ int)	Compute packed (l, m) index for real spherical harmonics layout.
`build_l_major_index`(→ numpy.ndarray)	Build coefficient indices for l-major layout truncated by mmax.
`build_m_major_index`(→ numpy.ndarray)	Build coefficient indices for m-major layout truncated by mmax.
`build_m_major_l_index`(→ numpy.ndarray)	Build degree (l) index aligned with build_m_major_index.
`build_rotate_inv_rescale`(→ numpy.ndarray)	Build reduced-layout inverse-rotation rescale factors.

Module Contents#

deepmd.dpmodel.descriptor.dpa4_nn.indexing.get_so3_dim_of_lmax(lmax: int) → int[source]#

Return SO(3) representation dimension for given lmax.

The dimension equals:

sum_{l<=lmax} (2l+1) = (lmax+1)^2

which is the number of spherical harmonics basis functions.

Parameters:

lmax: Maximum spherical harmonic degree.

Returns:

int: The SO(3) dimension D = (lmax+1)^2.

deepmd.dpmodel.descriptor.dpa4_nn.indexing.map_degree_idx(lmax: int) → numpy.ndarray[source]#

Build degree (l) index for each position in the packed (l, m) layout.

For each spherical harmonic coefficient position in the packed tensor, returns the corresponding angular momentum quantum number l.

The torch version’s device parameter is dropped: the output is a static numpy table.

Parameters:

lmax: Maximum angular momentum degree.

Returns:

np.ndarray: np.int64 array with shape (D,), where D=(lmax+1)^2. Each element is the l value for that position.

Examples

For lmax=2, the packed layout has D=9 positions: - Position 0: l=0, m=0 - Positions 1-3: l=1, m=-1,0,+1 - Positions 4-8: l=2, m=-2,-1,0,+1,+2

Returns: [0, 1,1,1, 2,2,2,2,2]

deepmd.dpmodel.descriptor.dpa4_nn.indexing.build_gie_zonal_index(lmax: int) → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray][source]#

Build node-level packed indices for GIE zonal coupling.

The returned arrays are aligned row-wise for every non-scalar packed coefficient in the node representation. They select the local m=0 column of the matching degree from Dt_full or an equivalent zonal coupling table.

The torch version’s device parameter is dropped: the output is a static numpy table.

Parameters:

lmax: Maximum node degree used by the geometric initial embedding.

Returns:

tuple[np.ndarray, np.ndarray, np.ndarray]: (node_row_index, node_zonal_m0_col_index, node_radial_l_index). The first two index packed SO(3) rows/columns; the last one indexes radial features with degree slots l=1..lmax stored as 0..lmax-1. All are np.int64 arrays.

deepmd.dpmodel.descriptor.dpa4_nn.indexing.project_D_to_m(D_full: Any, coeff_index_m: Any, ebed_dim_full: int, cache: dict[str, Any] | None, key_lmax: int, key_mmax: int) → Any[source]#

Row-project block-diagonal Wigner-D to the m-major truncated layout.

This function operates on runtime tensors and is array-API compatible.

Parameters:

D_full: Block-diagonal Wigner-D with shape (E, D, D).
coeff_index_m: Indices for m-major reduced layout with shape (D_m_trunc,).
ebed_dim_full: Full SO(3) dimension D_full = (lmax+1)^2 to slice the block.
cache: Optional cache mapping (lmax, mmax) -> projected matrix.
key_lmax: lmax used to build coeff_index_m (cache key).
key_mmax: mmax used to build coeff_index_m (cache key).

Returns:

Array: Projected rotation matrix with shape (E, D_m_trunc, D).

Examples

For lmax=2, mmax=1 (D=9, D_m_trunc=7), coeff_index_m selects [0,2,6,1,5,3,7] in packed (l,m) order. The returned array keeps only those rows of D_full while retaining all columns, so that rotating and truncating is done in a single matmul: x_local = D_to_m @ x_global.

deepmd.dpmodel.descriptor.dpa4_nn.indexing.project_Dt_from_m(Dt_full: Any, coeff_index_m: Any, ebed_dim_full: int, cache: dict[str, Any] | None, key_lmax: int, key_mmax: int) → Any[source]#

Column-project block-diagonal Wigner-D^T for inverse rotation.

This function operates on runtime tensors and is array-API compatible.

Parameters:

Dt_full: Block-diagonal Wigner-D^T with shape (E, D, D).
coeff_index_m: Indices for m-major reduced layout with shape (D_m_trunc,).
ebed_dim_full: Full SO(3) dimension D_full = (lmax+1)^2 to slice the block.
cache: Optional cache mapping (lmax, mmax) -> projected matrix.
key_lmax: lmax used to build coeff_index_m (cache key).
key_mmax: mmax used to build coeff_index_m (cache key).

Returns:

Array: Projected inverse rotation matrix with shape (E, D, D_m_trunc).

Examples

Continuing lmax=2, mmax=1, the projection selects the same column subset [0,2,6,1,5,3,7] from Dt_full. This enables inverse rotation with missing coefficients implicitly zeroed: x_global = Dt_from_m @ x_local.

deepmd.dpmodel.descriptor.dpa4_nn.indexing.so3_packed_index(degree: int, m: int) → int[source]#

Compute packed (l, m) index for real spherical harmonics layout.

The packed layout is l-primary with m ordered as -l..+l inside each l-block. The index formula is:

idx(l, m) = l^2 + l + m

Parameters:

degree: Degree l.
m: Order m, must satisfy -l <= m <= l.

Returns:

int: Packed index.

deepmd.dpmodel.descriptor.dpa4_nn.indexing.build_l_major_index(lmax: int, mmax: int) → numpy.ndarray[source]#

Build coefficient indices for l-major layout truncated by mmax.

The returned indices select coefficients with |m| <= min(mmax, l) in the standard packed (l, m) layout. The order is l-major:

l = 0..lmax
within each l, m = -min(mmax, l) .. +min(mmax, l)

The torch version’s device parameter is dropped: the output is a static numpy table.

Parameters:

lmax: Maximum degree.
mmax: Maximum order (|m|). Must satisfy 0 <= mmax <= lmax.

Returns:

np.ndarray: np.int64 array of indices with shape (D_m_trunc,), selecting coefficients from the full packed layout with D=(lmax+1)^2, where D_m_trunc is the number of coefficients kept under |m| <= min(mmax, l).

Examples

For lmax=2, mmax=1: - Full packed layout: l=0(0), l=1(1-3), l=2(4-8) - Truncated by mmax=1: skip (l=2, m=±2) at indices 4,8 - Returns: [0, 1, 2, 3, 5, 6, 7]

deepmd.dpmodel.descriptor.dpa4_nn.indexing.build_m_major_index(lmax: int, mmax: int) → numpy.ndarray[source]#

Build coefficient indices for m-major layout truncated by mmax.

This layout minimizes rotation cost and avoids gather-heavy indexing:

m = 0: l = 0..lmax (single coefficient per l)
for each m = 1..mmax:
- negative part: l = m..lmax, coefficient (l, -m)
- positive part: l = m..lmax, coefficient (l, +m)

The torch version’s device parameter is dropped: the output is a static numpy table.

Parameters:

lmax: Maximum degree.
mmax: Maximum order (|m|). Must satisfy 0 <= mmax <= lmax.

Returns:

np.ndarray: np.int64 array of indices with shape (D_m_trunc,), selecting coefficients from the full packed layout with D=(lmax+1)^2, where D_m_trunc is the number of coefficients kept under |m| <= min(mmax, l).

Examples

For lmax=2, mmax=1: - m=0 group: (l=0,m=0)→0, (l=1,m=0)→2, (l=2,m=0)→6 - m=1 neg group: (l=1,m=-1)→1, (l=2,m=-1)→5 - m=1 pos group: (l=1,m=+1)→3, (l=2,m=+1)→7 - Returns: [0, 2, 6, 1, 5, 3, 7]

deepmd.dpmodel.descriptor.dpa4_nn.indexing.build_m_major_l_index(lmax: int, mmax: int) → numpy.ndarray[source]#

Build degree (l) index aligned with build_m_major_index.

The torch version’s device parameter is dropped: the output is a static numpy table.

Parameters:

lmax: Maximum degree.
mmax: Maximum order (|m|). Must satisfy 0 <= mmax <= lmax.

Returns:

np.ndarray: np.int64 array of degrees with shape (D_m_trunc,). Entry i is the degree l for the i-th coefficient in the m-major layout.

Examples

For lmax=2, mmax=1: - m=0 group: l=0,1,2 - m=1 neg group: l=1,2 - m=1 pos group: l=1,2 - Returns: [0, 1, 2, 1, 2, 1, 2]

deepmd.dpmodel.descriptor.dpa4_nn.indexing.build_rotate_inv_rescale(lmax: int, mmax: int, degree_index: numpy.ndarray, *, dtype: Any = np.float64) → numpy.ndarray[source]#

Build reduced-layout inverse-rotation rescale factors.

When mmax < lmax, the reduced local layout keeps only 2*mmax+1 orders for each degree l > mmax. The inverse rotation rescales those truncated degrees by sqrt((2*l+1)/(2*mmax+1)) so the reduced representation matches the amplitude expected by the full SO(3) basis.

The torch version’s device parameter is dropped: the output is a static numpy table. dtype is kept (as a numpy dtype) since the floating-point precision of the rescale vector is meaningful.

Parameters:

lmax: Maximum degree.
mmax: Maximum order (|m|). Must satisfy 0 <= mmax <= lmax.
degree_index: Degree index aligned with the reduced coefficient layout, typically returned by build_m_major_l_index.
dtype: Floating-point numpy dtype for the returned array.

Returns:

np.ndarray: Rescale vector with shape (D_m_trunc,), aligned with the reduced coefficient layout.