core.models.escn.so3

core.models.escn.so3#

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Attributes#

_Jd

Classes#

CoefficientMapping

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

SO3_Embedding

Helper functions for irreps embedding

SO3_Rotation

Helper functions for Wigner-D rotations

SO3_Grid

Helper functions for grid representation of the irreps

Module Contents#

core.models.escn.so3._Jd#
class core.models.escn.so3.CoefficientMapping(lmax_list: list[int], mmax_list: list[int], device)#

Helper functions 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

  • device – Device of the output

lmax_list#
mmax_list#
num_resolutions#
device#
l_harmonic#
m_harmonic#
m_complex#
res_size#
offset = 0#
num_coefficients#
to_m#
m_size#
complex_idx(m, lmax: int = -1)#
coefficient_idx(lmax: int, mmax: int) → torch.Tensor#
class core.models.escn.so3.SO3_Embedding(length: int, lmax_list: list[int], num_channels: int, device: torch.device, dtype: torch.dtype)#

Bases: torch.nn.Module

Helper functions for irreps embedding

Parameters:

  • length (int) – Batch size

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

  • num_channels (int) – Number of channels

  • device – Device of the output

  • dtype – type of the output tensors

num_channels#
device#
dtype#
num_resolutions#
num_coefficients = 0#
embedding#
clone() → SO3_Embedding#
set_embedding(embedding) → None#
set_lmax_mmax(lmax_list, mmax_list) → None#
_expand_edge(edge_index) → None#
expand_edge(edge_index) → SO3_Embedding#
_reduce_edge(edge_index, num_nodes: int) → None#
_m_primary(mapping) → None#
_l_primary(mapping) → None#
_rotate(SO3_rotation, lmax_list, mmax_list) → None#
_rotate_inv(SO3_rotation, mappingReduced) → None#
_grid_act(SO3_grid, act, mappingReduced) → None#
to_grid(SO3_grid, lmax: int = -1) → torch.Tensor#
_from_grid(x_grid, SO3_grid, lmax: int = -1) → None#
class core.models.escn.so3.SO3_Rotation(rot_mat3x3: torch.Tensor, lmax: list[int])#

Bases: torch.nn.Module

Helper functions for Wigner-D rotations

Parameters:

  • rot_mat3x3 (tensor) – Rotation matrix

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

device#
dtype#
wigner#
wigner_inv#
set_lmax(lmax) → None#
rotate(embedding, out_lmax, out_mmax) → torch.Tensor#
rotate_inv(embedding, in_lmax, in_mmax) → torch.Tensor#
RotationToWignerDMatrix(edge_rot_mat: torch.Tensor, start_lmax: int, end_lmax: int) → torch.Tensor#
wigner_D(lval, alpha, beta, gamma)#
_z_rot_mat(angle: torch.Tensor, lv: int) → torch.Tensor#
class core.models.escn.so3.SO3_Grid(lmax: int, mmax: int, 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#
initialized = False#
_initialize(device: torch.device) → None#
get_to_grid_mat(device: torch.device)#
get_from_grid_mat(device: torch.device)#
to_grid(embedding: torch.Tensor, lmax: int, mmax: int) → torch.Tensor#
from_grid(grid: torch.Tensor, lmax: int, mmax: int) → torch.Tensor#
Contents