MolecularDiffusion.modules.layers.esen.nn.so2_layers

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

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_Conv

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

Module Contents

class MolecularDiffusion.modules.layers.esen.nn.so2_layers.SO2_Convolution(sphere_channels: int, m_output_channels: int, lmax: int, mmax: 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

  • lmax (int) – degrees (l)

  • mmax (int) – orders (m)

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

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

  • (list (edge_channels_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 and extra_m0_features (Tensor).

forward(x: torch.Tensor, x_edge: torch.Tensor)
edge_channels_list
extra_m0_output_channels = None
fc_m0
internal_weights = True
lmax
m_output_channels
mappingReduced
mmax
rad_func = None
so2_m_conv
sphere_channels
class MolecularDiffusion.modules.layers.esen.nn.so2_layers.SO2_Linear(sphere_channels: int, m_output_channels: int, lmax: int, mmax: 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

  • lmax (int) – degrees (l)

  • mmax (int) – orders (m)

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

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

  • (list (edge_channels_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
m_output_channels
mappingReduced
mmax
rad_func = None
so2_m_fc
sphere_channels
class MolecularDiffusion.modules.layers.esen.nn.so2_layers.SO2_m_Conv(m: int, sphere_channels: int, m_output_channels: int, lmax: int, mmax: 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

  • lmax (int) – degrees (l)

  • mmax (int) – orders (m)

forward(x_m)
fc
lmax
m
m_output_channels
mmax
out_channels_half
sphere_channels