MolecularDiffusion.modules.layers.equiformer_v2.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.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.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.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¶