Convolution Layers
Graph Convolution Operations
E3-equivariant convolution block (ConvBlockE3) with tensor products and skips.
Wraps message packing, radial networks, and optional KAN / gate nonlinearities for graph convolution on atomic data.
- class hamgnn.nn.convolution.ConvBlockE3(irreps_in, irreps_out, irreps_node_attrs, irreps_edge_attrs, irreps_edge_embed, radial_MLP=None, use_skip_connections=True, use_kan=False, nonlinearity_type='gate', nonlinearity_scalars={'e': 'ssp', 'o': 'tanh'}, nonlinearity_gates={'e': 'ssp', 'o': 'abs'}, lite_mode=False)[source]
Bases:
ModuleAn equivariant convolutional block for processing node features using tensor products with optional skip connections.
Parameters: - irreps_in (o3.Irreps): Input irreducible representations. - irreps_out (o3.Irreps): Output irreducible representations. - irreps_edge_attrs (o3.Irreps): Edge attribute irreducible representations. - irreps_edge_embed (o3.Irreps): Edge embedding irreducible representations. - radial_MLP (Optional[List[int]]): MLP architecture for radial embeddings. Defaults to [64, 64, 64]. - use_skip_connections (bool): Whether to use skip connections. Defaults to True. - use_kan (bool): Whether to use the FastKAN module for weight generation. Defaults to False. - nonlinearity_type (str): Type of nonlinearity to use (“gate” or “norm”). Defaults to “gate”. - nonlinearity_scalars (Dict[int, Callable]): Nonlinearity for scalar channels. - nonlinearity_gates (Dict[int, Callable]): Nonlinearity for gate channels. - lite_mode (bool): The mode with the fewest model parameters and the fastest running speed.