MolecularDiffusion.modules.layers.esen.esen_block¶

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes¶

`Edgewise`
`GridAtomwise`
`SpectralAtomwise`
`eSEN_Block`

Module Contents¶

class MolecularDiffusion.modules.layers.esen.esen_block.Edgewise(sphere_channels: int, hidden_channels: int, lmax: int, mmax: int, edge_channels_list, mappingReduced, SO3_grid, cutoff, act_type='gate', use_envelope: bool = True)¶

Bases: torch.nn.Module

forward(x, x_edge, edge_distance, edge_index, wigner, wigner_inv, node_offset: int = 0)¶

SO3_grid¶

act_type = 'gate'¶

edge_channels_list¶

hidden_channels¶

lmax¶

mappingReduced¶

mmax¶

out_mask¶

so2_conv_1¶

so2_conv_2¶

sphere_channels¶

use_envelope = True¶

class MolecularDiffusion.modules.layers.esen.esen_block.GridAtomwise(sphere_channels: int, hidden_channels: int, lmax: int, mmax: int, SO3_grid)¶

Bases: torch.nn.Module

forward(x)¶

SO3_grid¶

grid_mlp¶

hidden_channels¶

lmax¶

mmax¶

sphere_channels¶

class MolecularDiffusion.modules.layers.esen.esen_block.SpectralAtomwise(sphere_channels: int, hidden_channels: int, lmax: int, mmax: int, SO3_grid)¶

Bases: torch.nn.Module

forward(x)¶

SO3_grid¶

act¶

hidden_channels¶

lmax¶

mmax¶

scalar_mlp¶

so3_linear_1¶

so3_linear_2¶

sphere_channels¶

class MolecularDiffusion.modules.layers.esen.esen_block.eSEN_Block(sphere_channels: int, hidden_channels: int, lmax: int, mmax: int, mappingReduced, SO3_grid, edge_channels_list: list[int], cutoff: float, norm_type: str, act_type: str, mlp_type: str, use_envelope: bool)¶

Bases: torch.nn.Module

forward(x, x_edge, edge_distance, edge_index, wigner, wigner_inv, node_offset: int = 0)¶

edge_wise¶

hidden_channels¶

lmax¶

mmax¶

norm_1¶

norm_2¶

sphere_channels¶

use_envelope¶