deepmd.dpmodel.descriptor.dpa4_nn.projection

deepmd.dpmodel.descriptor.dpa4_nn.projection#

S2 grid projection helpers for DPA4/SeZM function-space nonlinearities.

This module is the dpmodel port of deepmd.pt.model.descriptor.sezm_nn.projection, restricted to the Lebedev S2 quadrature path used by the core DPA4 configuration (lebedev_quadrature=True). The projectors only handle basis transforms: a projector maps coefficient tensors to a fixed quadrature grid, and maps grid fields back to coefficients with the matching quadrature rule.

Ported names: BaseGridProjector, S2GridProjector (Lebedev branch), resolve_s2_grid_resolution (as-is, both methods — pure arithmetic), and _normalize_s2_grid_resolution.

SO(3) Wigner-D grid machinery (consumed by SO3GridNet in grid_net.py, which backs the node_wise_so3, message_node_so3, and ffn_so3_grid paths): resolve_so3_grid, _build_so3_frame_set, and SO3GridProjector are ported here. The SO(3) projection matrices are assembled at init time with pure numpy via the Wigner-D quadrature (WignerDCalculator) over a Lebedev x gamma rotation grid, matching the pt float64 buffers to machine precision.

Not ported (guarded): the e3nn product-grid branch of S2GridProjector (grid_method="e3nn", i.e. lebedev_quadrature=False) raises NotImplementedError at construction. Only the Lebedev path reproduces to-grid/from-grid roundtrip identities at machine precision.

The Lebedev projection matrices are assembled at init time with pure numpy: load_lebedev_rule replaces the pt Lebedev loader (same packaged data) and real_spherical_harmonics exactly replaces the e3nn call spherical_harmonics(list(range(lmax+1)), points, normalize=True, normalization="norm"), so the buffers match the pt float64 buffers to machine precision.

Classes#

`BaseGridProjector`	Base class for fixed coefficient-to-grid projection matrices.
`S2GridProjector`	Project SO(3) coefficients to/from a flattened S2 grid (Lebedev only).
`SO3GridProjector`	Project SO(3) coefficients to/from a Wigner-D grid with frame indices.

Functions#

`resolve_s2_grid_resolution`(→ list[int])	Resolve the default S2 grid resolution.
`_normalize_s2_grid_resolution`(→ list[int])	Resolve default grids or validate already-resolved low-level grids.
`resolve_so3_grid`(→ tuple[int, int, int])	Resolve the default SO(3) quadrature as Lebedev sphere times gamma samples.
`_build_so3_frame_set`(→ list[int])	Build the symmetric frame-index set with zero first.

Module Contents#

class deepmd.dpmodel.descriptor.dpa4_nn.projection.BaseGridProjector(*, lmax: int, mmax: int | None, precision: str = DEFAULT_PRECISION, n_frames: int, coefficient_layout: str)[source]#

Bases: deepmd.dpmodel.NativeOP

Base class for fixed coefficient-to-grid projection matrices.

Subclasses build to_grid_mat with shape (G, J) and from_grid_mat with shape (J, G), where G is the number of grid samples and J is the flattened coefficient axis consumed by the grid net. For ordinary S2 projections, J is the SO(3) feature coefficient axis: D = (lmax + 1)^2 in packed layout, or the retained D_m axis in m-major layout.

lmax[source]#

mmax = 0[source]#

coefficient_layout = ''[source]#

precision = 'float64'[source]#

n_frames[source]#

packed_dim[source]#

coeff_dim[source]#

grid_size[source]#

to_grid_mat[source]#

from_grid_mat[source]#

abstractmethod call(*args: Any, **kwargs: Any) → Any[source]#: Projectors expose to_grid/from_grid; there is no forward.

to_grid(embedding: Any) → Any[source]#: Project flattened coefficients (N, J, C) to grid fields (N, G, C).

from_grid(grid: Any) → Any[source]#: Project grid fields (N, G, C) back to flattened coefficients (N, J, C).

_build_coefficient_index() → numpy.ndarray[source]#: Build the coefficient subset consumed by the projector matrices.

abstractmethod _build_projection_mats(coeff_index: numpy.ndarray) → tuple[numpy.ndarray, numpy.ndarray][source]#: Build to_grid_mat (G, J) and from_grid_mat (J, G).

class deepmd.dpmodel.descriptor.dpa4_nn.projection.S2GridProjector(*, lmax: int, mmax: int | None = None, precision: str = DEFAULT_PRECISION, grid_resolution_list: list[int] | None = None, coefficient_layout: str = 'packed', grid_method: str = 'e3nn')[source]#

Bases: BaseGridProjector

Project SO(3) coefficients to/from a flattened S2 grid (Lebedev only).

Parameters:

lmax: Maximum spherical harmonic degree.
mmax: Maximum order kept in the coefficient layout. If None, use lmax.
precision: Buffer precision used by the projection matrices.
grid_resolution_list: Two-element resolution list [precision, n_points] for grid_method='lebedev'. If None, resolved automatically.
coefficient_layout: Coefficient ordering expected by the caller: - "packed": packed (l, m) order, optionally truncated by mmax. - "m_major": reduced m-major order used inside SO2Convolution.
grid_method: S2 quadrature backend. Must be "e3nn" or "lebedev"; only "lebedev" (lebedev_quadrature=True) is ported to dpmodel.

grid_method = ''[source]#

grid_resolution_list[source]#

phi_resolution = 0[source]#

theta_resolution = 0[source]#

_rescale_truncated_matrix(mat: numpy.ndarray) → None[source]#

_build_projection_mats(coeff_index: numpy.ndarray) → tuple[numpy.ndarray, numpy.ndarray][source]#: Build to_grid_mat (G, J) and from_grid_mat (J, G).

serialize() → dict[str, Any][source]#: Serialize the S2GridProjector to a dict (pt-compatible format).

classmethod deserialize(data: dict[str, Any]) → S2GridProjector[source]#: Deserialize an S2GridProjector from a dict.

class deepmd.dpmodel.descriptor.dpa4_nn.projection.SO3GridProjector(*, lmax: int, mmax: int | None = None, kmax: int = 1, precision: str = DEFAULT_PRECISION, lebedev_precision: int | None = None, coefficient_layout: str = 'packed')[source]#

Bases: BaseGridProjector

Project SO(3) coefficients to/from a Wigner-D grid with frame indices.

The coefficient axis is packed as (l, m, k) with ordinary SeZM (l, m) order outside and the configured frame set inside each row. A frame index outside [-l, l] is kept as a zero column/row. This keeps the tensor layout regular while preserving the exact per-degree frame support.

Parameters:

lmax: Maximum spherical harmonic degree.
mmax: Maximum order kept in the coefficient layout. If None, use lmax.
kmax: Frame-index half-width; the frame set is {0, -1, 1, ..., -kmax, kmax}.
precision: Buffer precision used by the projection matrices.
lebedev_precision: Explicit Lebedev rule precision. If None, resolved automatically.
coefficient_layout: Coefficient ordering expected by the caller: - "packed": packed (l, m) order, optionally truncated by mmax. - "m_major": reduced m-major order used inside SO2Convolution.

kmax = 1[source]#

frame_set[source]#

frame_zero_index[source]#

frame_values[source]#

_build_projection_mats(coeff_index: numpy.ndarray) → tuple[numpy.ndarray, numpy.ndarray][source]#: Build to_grid_mat (G, J) and from_grid_mat (J, G).

serialize() → dict[str, Any][source]#: Serialize the SO3GridProjector to a dict (pt-compatible format).

classmethod deserialize(data: dict[str, Any]) → SO3GridProjector[source]#: Deserialize an SO3GridProjector from a dict.

deepmd.dpmodel.descriptor.dpa4_nn.projection.resolve_s2_grid_resolution(lmax: int, mmax: int, *, method: str = 'e3nn') → list[int][source]#

Resolve the default S2 grid resolution.

For method='e3nn', the automatic default uses even azimuthal sampling R_phi = 2 * mmax + 4 and even polar sampling R_theta = ceil_even(3 * lmax + 2).

For method='lebedev', the automatic default picks the smallest packaged Lebedev rule whose algebraic precision is at least 3 * lmax and returns [precision, n_points].

deepmd.dpmodel.descriptor.dpa4_nn.projection._normalize_s2_grid_resolution(lmax: int, mmax: int, grid_resolution_list: list[int] | None, *, method: str) → list[int][source]#: Resolve default grids or validate already-resolved low-level grids.

deepmd.dpmodel.descriptor.dpa4_nn.projection.resolve_so3_grid(lmax: int, *, kmax: int = 1, lebedev_precision: int | None = None) → tuple[int, int, int][source]#

Resolve the default SO(3) quadrature as Lebedev sphere times gamma samples.

The Lebedev precision follows the same conservative 3*lmax rule used by the S2 grid path. The gamma grid is chosen for the quadratic grid products used by the SO(3) grid nets, whose third-angle frequency can reach k1 + k2 - kout.

deepmd.dpmodel.descriptor.dpa4_nn.projection._build_so3_frame_set(kmax: int) → list[int][source]#: Build the symmetric frame-index set with zero first.