MolecularDiffusion.modules.layers.equiformer_v2_s.so2_ops

Classes

SO2_Convolution

SO(2) Block: Perform SO(2) convolutions for all m (orders)

SO2_Linear

SO(2) Linear: Perform SO(2) linear for all m (orders).

SO2_m_Convolution

SO(2) Conv: Perform an SO(2) convolution on features corresponding to +- m

Module Contents

class MolecularDiffusion.modules.layers.equiformer_v2_s.so2_ops.SO2_Convolution(sphere_channels, m_output_channels, lmax_list, mmax_list, mappingReduced, internal_weights=True, edge_channels_list=None, extra_m0_output_channels=None)

Bases: torch.nn.Module

SO(2) Block: Perform SO(2) convolutions for all m (orders)

Parameters:

  • sphere_channels (int) – Number of spherical channels

  • m_output_channels (int) – Number of output channels used during the SO(2) conv

  • (list (edge_channels_list) – int): List of degrees (l) for each resolution

  • (list – int): List of orders (m) for each resolution

  • mappingReduced (CoefficientMappingModule) – Used to extract a subset of m components

  • internal_weights (bool) – If True, not using radial function to multiply inputs features

  • (list – int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels].

  • extra_m0_output_channels (int) – If not None, return out_embedding (SO3_Embedding) and extra_m0_features (Tensor).

forward(x, x_edge)
edge_channels_list
extra_m0_output_channels = None
fc_m0
internal_weights = True
lmax_list
m_output_channels
mappingReduced
mmax_list
num_resolutions
rad_func = None
so2_m_conv
sphere_channels
class MolecularDiffusion.modules.layers.equiformer_v2_s.so2_ops.SO2_Linear(sphere_channels, m_output_channels, lmax_list, mmax_list, mappingReduced, internal_weights=False, edge_channels_list=None)

Bases: torch.nn.Module

SO(2) Linear: Perform SO(2) linear for all m (orders).

Parameters:

  • sphere_channels (int) – Number of spherical channels

  • m_output_channels (int) – Number of output channels used during the SO(2) conv

  • (list (edge_channels_list) – int): List of degrees (l) for each resolution

  • (list – int): List of orders (m) for each resolution

  • mappingReduced (CoefficientMappingModule) – Used to extract a subset of m components

  • internal_weights (bool) – If True, not using radial function to multiply inputs features

  • (list – int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels].

forward(x, x_edge)
edge_channels_list
fc_m0
internal_weights = False
lmax_list
m_output_channels
mappingReduced
mmax_list
num_resolutions
rad_func = None
so2_m_fc
sphere_channels
class MolecularDiffusion.modules.layers.equiformer_v2_s.so2_ops.SO2_m_Convolution(m, sphere_channels, m_output_channels, lmax_list, mmax_list)

Bases: torch.nn.Module

SO(2) Conv: Perform an SO(2) convolution on features corresponding to +- m

Parameters:

  • m (int) – Order of the spherical harmonic coefficients

  • sphere_channels (int) – Number of spherical channels

  • m_output_channels (int) – Number of output channels used during the SO(2) conv

  • (list (mmax_list) – int): List of degrees (l) for each resolution

  • (list – int): List of orders (m) for each resolution

forward(x_m)
fc
lmax_list
m
m_output_channels
mmax_list
num_resolutions
sphere_channels