core.models.equiformer_v2.input_block#
Classes#
Module Contents#
- class core.models.equiformer_v2.input_block.EdgeDegreeEmbedding(sphere_channels: int, lmax_list: list[int], mmax_list: list[int], SO3_rotation, mappingReduced, max_num_elements: int, edge_channels_list, use_atom_edge_embedding: bool, 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
- sphere_channels#
- lmax_list#
- mmax_list#
- num_resolutions#
- SO3_rotation#
- mappingReduced#
- m_0_num_coefficients: int#
- m_all_num_coefficents: int#
- max_num_elements#
- edge_channels_list#
- use_atom_edge_embedding#
- rad_func#
- rescale_factor#
- forward(atomic_numbers, edge_distance, edge_index, num_nodes, node_offset=0)#