deepmd.pt.model.descriptor.sezm_nn.indexing

deepmd.pt.model.descriptor.sezm_nn.indexing#

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

This module defines the packed (l, m) indexing helpers and the projection utilities used by the SeZM equivariant operators.

Functions#

`get_so3_dim_of_lmax`(→ int)	Return SO(3) representation dimension for given lmax.
`map_degree_idx`(→ torch.Tensor)	Build degree (l) index for each position in the packed (l, m) layout.
`build_gie_zonal_index`(→ tuple[torch.Tensor, ...)	Build node-level packed indices for GIE zonal coupling.
`project_D_to_m`(→ torch.Tensor)	Row-project block-diagonal Wigner-D to the m-major truncated layout.
`project_Dt_from_m`(→ torch.Tensor)	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`(→ torch.Tensor)	Build coefficient indices for l-major layout truncated by mmax.
`build_m_major_index`(→ torch.Tensor)	Build coefficient indices for m-major layout truncated by mmax.
`build_m_major_l_index`(→ torch.Tensor)	Build degree (l) index aligned with build_m_major_index.
`build_rotate_inv_rescale`(→ torch.Tensor)	Build reduced-layout inverse-rotation rescale factors.

Module Contents#

deepmd.pt.model.descriptor.sezm_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.pt.model.descriptor.sezm_nn.indexing.map_degree_idx(lmax: int, *, device: torch.device) → torch.Tensor[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.

Parameters:

lmax: Maximum angular momentum degree.
device: Device for the returned tensor.

Returns:

torch.Tensor: Integer tensor 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.pt.model.descriptor.sezm_nn.indexing.build_gie_zonal_index(lmax: int, *, device: torch.device) → tuple[torch.Tensor, torch.Tensor, torch.Tensor][source]#

Build node-level packed indices for GIE zonal coupling.

The returned tensors 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.

Parameters:

lmax: Maximum node degree used by the geometric initial embedding.
device: Device for the returned tensors.

Returns:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]: (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.

deepmd.pt.model.descriptor.sezm_nn.indexing.project_D_to_m(D_full: torch.Tensor, coeff_index_m: torch.Tensor, ebed_dim_full: int, cache: dict[str, torch.Tensor] | None, key_lmax: int, key_mmax: int) → torch.Tensor[source]#

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

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:

torch.Tensor: 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 tensor keeps only those rows of D_full while retaining all columns, so that rotating and truncating is done in a single bmm: x_local = D_to_m @ x_global.

deepmd.pt.model.descriptor.sezm_nn.indexing.project_Dt_from_m(Dt_full: torch.Tensor, coeff_index_m: torch.Tensor, ebed_dim_full: int, cache: dict[str, torch.Tensor] | None, key_lmax: int, key_mmax: int) → torch.Tensor[source]#

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

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:

torch.Tensor: 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.pt.model.descriptor.sezm_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.pt.model.descriptor.sezm_nn.indexing.build_l_major_index(lmax: int, mmax: int, *, device: torch.device) → torch.Tensor[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)

Parameters:

lmax: Maximum degree.
mmax: Maximum order (|m|). Must satisfy 0 <= mmax <= lmax.
device: Device for the returned tensor.

Returns:

torch.Tensor: Long tensor 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.pt.model.descriptor.sezm_nn.indexing.build_m_major_index(lmax: int, mmax: int, *, device: torch.device) → torch.Tensor[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)

Parameters:

lmax: Maximum degree.
mmax: Maximum order (|m|). Must satisfy 0 <= mmax <= lmax.
device: Device for the returned tensor.

Returns:

torch.Tensor: Long tensor 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.pt.model.descriptor.sezm_nn.indexing.build_m_major_l_index(lmax: int, mmax: int, *, device: torch.device) → torch.Tensor[source]#

Build degree (l) index aligned with build_m_major_index.

Parameters:

lmax: Maximum degree.
mmax: Maximum order (|m|). Must satisfy 0 <= mmax <= lmax.
device: Device for the returned tensor.

Returns:

torch.Tensor: Long tensor 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.pt.model.descriptor.sezm_nn.indexing.build_rotate_inv_rescale(lmax: int, mmax: int, degree_index: torch.Tensor, *, device: torch.device, dtype: torch.dtype) → torch.Tensor[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.

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.
device: Device for the returned tensor.
dtype: Floating-point dtype for the returned tensor.

Returns:

torch.Tensor: Rescale vector with shape (D_m_trunc,), aligned with the reduced coefficient layout.