core.models.equiformer_v2.so2_ops

core.models.equiformer_v2.so2_ops#

Classes#

SO2_m_Convolution

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

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).

Module Contents#

class core.models.equiformer_v2.so2_ops.SO2_m_Convolution(m: int, sphere_channels: int, m_output_channels: int, lmax_list: list[int], mmax_list: list[int])#

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

m#
sphere_channels#
m_output_channels#
lmax_list#
mmax_list#
num_resolutions: int#
num_channels = 0#
fc#
forward(x_m)#
class core.models.equiformer_v2.so2_ops.SO2_Convolution(sphere_channels: int, m_output_channels: int, lmax_list: list[int], mmax_list: list[int], mappingReduced, internal_weights: bool = True, edge_channels_list: list[int] | None = None, extra_m0_output_channels: int | None = 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).

sphere_channels#
m_output_channels#
lmax_list#
mmax_list#
mappingReduced#
num_resolutions#
internal_weights#
edge_channels_list#
extra_m0_output_channels#
num_channels_rad = 0#
num_channels_m0 = 0#
m0_output_channels#
fc_m0#
so2_m_conv#
rad_func = None#
forward(x, x_edge)#
class core.models.equiformer_v2.so2_ops.SO2_Linear(sphere_channels: int, m_output_channels: int, lmax_list: list[int], mmax_list: list[int], mappingReduced, internal_weights: bool = False, edge_channels_list: list[int] | None = 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].

sphere_channels#
m_output_channels#
lmax_list#
mmax_list#
mappingReduced#
internal_weights#
edge_channels_list#
num_resolutions#
num_channels_rad = 0#
num_channels_m0 = 0#
fc_m0#
so2_m_fc#
rad_func = None#
forward(x, x_edge)#
Contents