MolecularDiffusion.modules.layers.esen.nn.embedding¶

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

Classes¶

`ChgSpinDatasetEmbedding`
`EdgeDegreeEmbedding`

Module Contents¶

class MolecularDiffusion.modules.layers.esen.nn.embedding.ChgSpinDatasetEmbedding(embedding_type, embedding_target, embedding_size, grad, scale=1.0)¶

Bases: torch.nn.Module

forward(x)¶

embedding_target¶

embedding_type¶

class MolecularDiffusion.modules.layers.esen.nn.embedding.EdgeDegreeEmbedding(sphere_channels: int, lmax: int, mmax: int, max_num_elements: int, edge_channels_list, rescale_factor, cutoff, mappingReduced, out_mask, use_envelope)¶

Bases: torch.nn.Module

Parameters:

sphere_channels (int) – Number of spherical channels
lmax (int) – degrees (l)
mmax (int) – orders (m)
max_num_elements (int) – Maximum number of atomic numbers
(list (edge_channels_list) – int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels]. The last one will be used as hidden size when use_atom_edge_embedding is True.
use_atom_edge_embedding (bool) – Whether to use atomic embedding along with relative distance for edge scalar features
rescale_factor (float) – Rescale the sum aggregation
cutoff (float) – Cutoff distance for the radial function
mappingReduced (CoefficientMapping) – Class to convert l and m indices once node embedding is rotated
out_mask (torch.Tensor) – Mask to select the output irreps
use_envelope (bool) – Whether to use envelope function

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

edge_channels_list¶

lmax¶

m_0_num_coefficients: int¶

m_all_num_coefficents: int¶

mappingReduced¶

max_num_elements¶

mmax¶

out_mask¶

rad_func¶

rescale_factor¶

sphere_channels¶

use_envelope¶