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#

res_size#

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, rescale: bool = False)#

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#

get_to_grid_mat(device=None)#

get_from_grid_mat(device=None)#

to_grid(embedding, lmax: int, mmax: int)#

from_grid(grid, lmax: int, mmax: int)#