core.models.equiformer_v2.so3

core.models.equiformer_v2.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.

Classes#

`CoefficientMappingModule`	Helper module for coefficients used to reshape lval <--> m and to get coefficients of specific degree or order
`SO3_Embedding`	Helper functions for performing operations on irreps embedding
`SO3_Rotation`	Helper functions for Wigner-D rotations
`SO3_Grid`	Helper functions for grid representation of the irreps
`SO3_Linear`	Base class for all neural network modules.
`SO3_LinearV2`	Base class for all neural network modules.

Module Contents#

class core.models.equiformer_v2.so3.CoefficientMappingModule(lmax_list: list[int], mmax_list: list[int])#

Bases: torch.nn.Module

Helper module for coefficients used to reshape lval <–> 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

lmax_list#

mmax_list#

num_resolutions#

device = 'cpu'#

l_harmonic#

m_harmonic#

m_complex#

res_size#

offset = 0#

num_coefficients#

to_m#

m_size#

mask_indices_cache = None#

rotate_inv_rescale_cache = None#

complex_idx(m: int, lmax: int, m_complex, l_harmonic)#: Add m_complex and l_harmonic to the input arguments since we cannot use self.m_complex.

coefficient_idx(lmax: int, mmax: int)#

get_rotate_inv_rescale(lmax: int, mmax: int)#

__repr__() → str#

class core.models.equiformer_v2.so3.SO3_Embedding(length: int, lmax_list: list[int], num_channels: int, device: torch.device, dtype: torch.dtype)#

Helper functions for performing operations on 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: list[int], mmax_list: list[int]) → None#

_expand_edge(edge_index: torch.Tensor) → None#

expand_edge(edge_index: torch.Tensor)#

_reduce_edge(edge_index: torch.Tensor, num_nodes: int)#

_m_primary(mapping)#

_l_primary(mapping)#

_rotate(SO3_rotation, lmax_list: list[int], mmax_list: list[int])#

_rotate_inv(SO3_rotation, mappingReduced)#

_grid_act(SO3_grid, act, mappingReduced)#

to_grid(SO3_grid, lmax=-1)#

_from_grid(x_grid, SO3_grid, lmax: int = -1)#

class core.models.equiformer_v2.so3.SO3_Rotation(lmax: int)#

Bases: torch.nn.Module

Helper functions for Wigner-D rotations

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

lmax#

mapping#

set_wigner(rot_mat3x3)#

rotate(embedding, out_lmax: int, out_mmax: int)#

rotate_inv(embedding, in_lmax: int, in_mmax: int)#

RotationToWignerDMatrix(edge_rot_mat, start_lmax: int, end_lmax: int) → torch.Tensor#

class core.models.equiformer_v2.so3.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)#

get_from_grid_mat(device)#

class core.models.equiformer_v2.so3.SO3_Linear(in_features: int, out_features: int, lmax: int, bias: bool = True)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

in_features#

out_features#

lmax#

linear_list#

forward(input_embedding, output_scale=None)#

__repr__() → str#

class core.models.equiformer_v2.so3.SO3_LinearV2(in_features: int, out_features: int, lmax: int, bias: bool = True)#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

in_features#

out_features#

lmax#

weight#

bound#

bias#

expand_index#

forward(input_embedding)#

__repr__() → str#