core.models.escn

core.models.escn#

Submodules#

Classes#

eSCN

Equivariant Spherical Channel Network

Package Contents#

class core.models.escn.eSCN(use_pbc: bool = True, use_pbc_single: bool = False, regress_forces: bool = True, otf_graph: bool = False, max_neighbors: int = 40, cutoff: float = 8.0, max_num_elements: int = 90, num_layers: int = 8, lmax_list: list[int] | None = None, mmax_list: list[int] | None = None, sphere_channels: int = 128, hidden_channels: int = 256, edge_channels: int = 128, num_sphere_samples: int = 128, distance_function: str = 'gaussian', basis_width_scalar: float = 1.0, distance_resolution: float = 0.02, show_timing_info: bool = False, resolution: int | None = None, activation_checkpoint: bool | None = False)#

Bases: torch.nn.Module, fairchem.core.models.base.GraphModelMixin

Equivariant Spherical Channel Network Paper: Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs

Parameters:

  • use_pbc (bool) – Use periodic boundary conditions

  • use_pbc_single (bool) – Process batch PBC graphs one at a time

  • regress_forces (bool) – Compute forces

  • otf_graph (bool) – Compute graph On The Fly (OTF)

  • max_neighbors (int) – Maximum number of neighbors per atom

  • cutoff (float) – Maximum distance between nieghboring atoms in Angstroms

  • max_num_elements (int) – Maximum atomic number

  • num_layers (int) – Number of layers in the GNN

  • lmax_list (int) – List of maximum degree of the spherical harmonics (1 to 10)

  • mmax_list (int) – List of maximum order of the spherical harmonics (0 to lmax)

  • sphere_channels (int) – Number of spherical channels (one set per resolution)

  • hidden_channels (int) – Number of hidden units in message passing

  • num_sphere_samples (int) – Number of samples used to approximate the integration of the sphere in the output blocks

  • edge_channels (int) – Number of channels for the edge invariant features

  • distance_function ("gaussian", "sigmoid", "linearsigmoid", "silu") – Basis function used for distances

  • basis_width_scalar (float) – Width of distance basis function

  • distance_resolution (float) – Distance between distance basis functions in Angstroms

  • show_timing_info (bool) – Show timing and memory info

activation_checkpoint#
regress_forces#
use_pbc#
use_pbc_single#
cutoff#
otf_graph#
show_timing_info#
max_num_elements#
hidden_channels#
num_layers#
num_atoms = 0#
num_sphere_samples#
sphere_channels#
max_neighbors#
edge_channels#
distance_resolution#
grad_forces = False#
lmax_list#
mmax_list#
num_resolutions: int#
sphere_channels_all: int#
basis_width_scalar#
distance_function#
counter = 0#
act#
sphere_embedding#
num_gaussians#
SO3_grid#
layer_blocks#
energy_block#
sphere_points#
sphharm_weights#
forward(data)#
_init_edge_rot_mat(data, edge_index, edge_distance_vec)#
property num_params: int#