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.

Skipped names (SO(3) Wigner-D grid machinery; consumed only by SO3GridNet in pt grid_net.py, which backs the node_wise_so3, message_node_so3, and ffn_so3_grid paths — all disabled in the core DPA4 config): SO3GridProjector, resolve_so3_grid, _build_so3_frame_set.

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

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.

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.

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.