MolecularDiffusion.modules.layers.equiformer_v2.so3

Copyright (c) Meta Platforms, Inc. and 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 l <--> m and to get coefficients of specific

SO3_Embedding

Helper functions for performing operations on irreps embedding.

SO3_Grid

Helper functions for grid representation of the irreps.

SO3_Linear

SO3_LinearV2

SO3_Rotation

Helper functions for Wigner-D rotations.

Module Contents

class MolecularDiffusion.modules.layers.equiformer_v2.so3.CoefficientMappingModule(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

coefficient_idx(lmax, mmax)
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.

get_rotate_inv_rescale(lmax, mmax)
device = 'cpu'
lmax_list
mask_indices_cache = None
mmax_list
num_resolutions
rotate_inv_rescale_cache = None
class MolecularDiffusion.modules.layers.equiformer_v2.so3.SO3_Embedding(length, lmax_list, num_channels, device, 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

clone()
expand_edge(edge_index)
set_embedding(embedding)
set_lmax_mmax(lmax_list, mmax_list)
to_grid(SO3_grid, lmax=-1)
device
dtype
num_channels
num_coefficients = 0
num_resolutions
class MolecularDiffusion.modules.layers.equiformer_v2.so3.SO3_Grid(lmax, mmax, normalization='integral', resolution=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

from_grid(grid, lmax, mmax)
get_from_grid_mat(device)
get_to_grid_mat(device)
to_grid(embedding, lmax, mmax)
lat_resolution
lmax
mapping
mmax
class MolecularDiffusion.modules.layers.equiformer_v2.so3.SO3_Linear(in_features, out_features, lmax, bias=True)

Bases: torch.nn.Module

forward(input_embedding, output_scale=None)
in_features
linear_list
lmax
out_features
class MolecularDiffusion.modules.layers.equiformer_v2.so3.SO3_LinearV2(in_features, out_features, lmax, bias=True)

Bases: torch.nn.Module

  1. Use torch.einsum to prevent slicing and concatenation

  2. Need to specify some behaviors in no_weight_decay and weight initialization.

forward(input_embedding)
bias
in_features
lmax
out_features
weight
class MolecularDiffusion.modules.layers.equiformer_v2.so3.SO3_Rotation(lmax)

Bases: torch.nn.Module

Helper functions for Wigner-D rotations.

Parameters:

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

RotationToWignerDMatrix(edge_rot_mat, start_lmax, end_lmax)
rotate(embedding, out_lmax, out_mmax)
rotate_inv(embedding, in_lmax, in_mmax)
set_wigner(rot_mat3x3)
lmax
mapping