MolecularDiffusion.modules.layers.equiformer_v2_s.input_block¶
Classes¶
Module Contents¶
- class MolecularDiffusion.modules.layers.equiformer_v2_s.input_block.EdgeDegreeEmbedding(sphere_channels, lmax_list, mmax_list, SO3_rotation, mappingReduced, max_num_elements, edge_channels_list, use_atom_edge_embedding, rescale_factor)¶
Bases:
torch.nn.Module- Parameters:
sphere_channels (int) – Number of spherical channels
(list (edge_channels_list) – int): List of degrees (l) for each resolution
(list – int): List of orders (m) for each resolution
(list – SO3_Rotation): Class to calculate Wigner-D matrices and rotate embeddings
mappingReduced (CoefficientMappingModule) – Class to convert l and m indices once node embedding is rotated
max_num_elements (int) – Maximum number of atomic numbers
(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
- forward(atomic_numbers, edge_distance, edge_index)¶
- SO3_rotation¶
- edge_channels_list¶
- lmax_list¶
- m_0_num_coefficients¶
- m_all_num_coefficents¶
- mappingReduced¶
- max_num_elements¶
- mmax_list¶
- num_resolutions¶
- rad_func¶
- rescale_factor¶
- sphere_channels¶
- use_atom_edge_embedding¶