core.models.equiformer_v2.input_block

core.models.equiformer_v2.input_block#

Classes#

EdgeDegreeEmbedding

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)#
Contents