core.models.escn.so3_exportable

core.models.escn.so3_exportable#

Classes#

CoefficientMapping

Helper module for coefficients used to reshape l <--> m and to get coefficients of specific degree or order

SO3_Grid

Helper functions for grid representation of the irreps

Functions#

wigner_D(→ torch.Tensor)

_z_rot_mat(→ torch.Tensor)

rotation_to_wigner(→ torch.Tensor)

Module Contents#

core.models.escn.so3_exportable.wigner_D(lv: int, alpha: torch.Tensor, beta: torch.Tensor, gamma: torch.Tensor, _Jd: list[torch.Tensor]) → torch.Tensor#
core.models.escn.so3_exportable._z_rot_mat(angle: torch.Tensor, lv: int) → torch.Tensor#
core.models.escn.so3_exportable.rotation_to_wigner(edge_rot_mat: torch.Tensor, start_lmax: int, end_lmax: int, Jd: list[torch.Tensor]) → torch.Tensor#
class core.models.escn.so3_exportable.CoefficientMapping(lmax_list, mmax_list)#

Bases: torch.nn.Module

Helper module for coefficients used to reshape l <–> m and to get coefficients of specific degree or order

Parameters:

  • (list (mmax_list) – int): List of maximum degree of the spherical harmonics

  • (list – int): List of maximum order of the spherical harmonics

  • use_rotate_inv_rescale (bool) – Whether to pre-compute inverse rotation rescale matrices

lmax_list#
mmax_list#
num_resolutions#
l_harmonic#
m_harmonic#
m_complex#
res_size#
offset = 0#
num_coefficients#
to_m#
m_size#
complex_idx(m, lmax, m_complex, l_harmonic)#

Add m_complex and l_harmonic to the input arguments since we cannot use self.m_complex.

pre_compute_coefficient_idx()#

Pre-compute the results of coefficient_idx() and access them with prepare_coefficient_idx()

prepare_coefficient_idx()#

Construct a list of buffers

coefficient_idx(lmax: int, mmax: int)#
pre_compute_rotate_inv_rescale()#
__repr__()#
class core.models.escn.so3_exportable.SO3_Grid(lmax: int, mmax: int, normalization: str = 'integral', resolution: int | None = None)#

Bases: torch.nn.Module

Helper functions for grid representation of the irreps

Parameters:

  • lmax (int) – Maximum degree of the spherical harmonics

  • mmax (int) – Maximum order of the spherical harmonics

lmax#
mmax#
lat_resolution#
mapping#
device = 'cpu'#
to_grid#
to_grid_mat#
from_grid#
from_grid_mat#
get_to_grid_mat(device=None)#
get_from_grid_mat(device=None)#
Contents