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

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

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