core.models.escn.escn

Contents

core.models.escn.escn#

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes#

eSCN

Equivariant Spherical Channel Network

eSCNBackbone

Equivariant Spherical Channel Network

eSCNEnergyHead

Base class for all neural network modules.

eSCNForceHead

Base class for all neural network modules.

LayerBlock

Layer block: Perform one layer (message passing and aggregation) of the GNN

MessageBlock

Message block: Perform message passing

SO2Block

SO(2) Block: Perform SO(2) convolutions for all m (orders)

SO2Conv

SO(2) Conv: Perform an SO(2) convolution

EdgeBlock

Edge Block: Compute invariant edge representation from edge diatances and atomic numbers

EnergyBlock

Energy Block: Output block computing the energy

ForceBlock

Force Block: Output block computing the per atom forces

Module Contents#

class core.models.escn.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#
class core.models.escn.escn.eSCNBackbone(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: eSCN, fairchem.core.models.base.BackboneInterface

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

forward(data: torch_geometric.data.batch.Batch) → dict[str, torch.Tensor]#

Backbone forward.

Parameters:

data (DataBatch) – Atomic systems as input

Returns:

embedding – Return backbone embeddings for the given input

Return type:

dict[str->torch.Tensor]

class core.models.escn.escn.eSCNEnergyHead(backbone, reduce='sum')#

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

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

reduce#
energy_block#
forward(data: torch_geometric.data.batch.Batch, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

  • data (DataBatch) – Atomic systems as input

  • emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

class core.models.escn.escn.eSCNForceHead(backbone)#

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

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

force_block#
forward(data: torch_geometric.data.batch.Batch, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]#

Head forward.

Parameters:

  • data (DataBatch) – Atomic systems as input

  • emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

class core.models.escn.escn.LayerBlock(layer_idx: int, sphere_channels: int, hidden_channels: int, edge_channels: int, lmax_list: list[int], mmax_list: list[int], distance_expansion, max_num_elements: int, SO3_grid: fairchem.core.models.escn.so3.SO3_Grid, act)#

Bases: torch.nn.Module

Layer block: Perform one layer (message passing and aggregation) of the GNN

Parameters:

  • layer_idx (int) – Layer number

  • sphere_channels (int) – Number of spherical channels

  • hidden_channels (int) – Number of hidden channels used during the SO(2) conv

  • edge_channels (int) – Size of invariant edge embedding

  • (list (mmax_list) – int): List of degrees (l) for each resolution

  • (list – int): List of orders (m) for each resolution

  • distance_expansion (func) – Function used to compute distance embedding

  • max_num_elements (int) – Maximum number of atomic numbers

  • SO3_grid (SO3_grid) – Class used to convert from grid the spherical harmonic representations

  • act (function) – Non-linear activation function

layer_idx#
act#
lmax_list#
mmax_list#
num_resolutions#
sphere_channels#
sphere_channels_all#
SO3_grid#
message_block#
fc1_sphere#
fc2_sphere#
fc3_sphere#
forward(x, atomic_numbers, edge_distance, edge_index, SO3_edge_rot, mappingReduced)#
class core.models.escn.escn.MessageBlock(layer_idx: int, sphere_channels: int, hidden_channels: int, edge_channels: int, lmax_list: list[int], mmax_list: list[int], distance_expansion, max_num_elements: int, SO3_grid: fairchem.core.models.escn.so3.SO3_Grid, act)#

Bases: torch.nn.Module

Message block: Perform message passing

Parameters:

  • layer_idx (int) – Layer number

  • sphere_channels (int) – Number of spherical channels

  • hidden_channels (int) – Number of hidden channels used during the SO(2) conv

  • edge_channels (int) – Size of invariant edge embedding

  • (list (mmax_list) – int): List of degrees (l) for each resolution

  • (list – int): List of orders (m) for each resolution

  • distance_expansion (func) – Function used to compute distance embedding

  • max_num_elements (int) – Maximum number of atomic numbers

  • SO3_grid (SO3_grid) – Class used to convert from grid the spherical harmonic representations

  • act (function) – Non-linear activation function

layer_idx#
act#
hidden_channels#
sphere_channels#
SO3_grid#
num_resolutions#
lmax_list#
mmax_list#
edge_channels#
edge_block#
so2_block_source#
so2_block_target#
forward(x, atomic_numbers, edge_distance, edge_index, SO3_edge_rot, mappingReduced)#
class core.models.escn.escn.SO2Block(sphere_channels: int, hidden_channels: int, edge_channels: int, lmax_list: list[int], mmax_list: list[int], act)#

Bases: torch.nn.Module

SO(2) Block: Perform SO(2) convolutions for all m (orders)

Parameters:

  • sphere_channels (int) – Number of spherical channels

  • hidden_channels (int) – Number of hidden channels used during the SO(2) conv

  • edge_channels (int) – Size of invariant edge embedding

  • (list (mmax_list) – int): List of degrees (l) for each resolution

  • (list – int): List of orders (m) for each resolution

  • act (function) – Non-linear activation function

sphere_channels#
hidden_channels#
lmax_list#
mmax_list#
num_resolutions: int#
act#
fc1_dist0#
fc1_m0#
fc2_m0#
so2_conv#
forward(x, x_edge, mappingReduced)#
class core.models.escn.escn.SO2Conv(m: int, sphere_channels: int, hidden_channels: int, edge_channels: int, lmax_list: list[int], mmax_list: list[int], act)#

Bases: torch.nn.Module

SO(2) Conv: Perform an SO(2) convolution

Parameters:

  • m (int) – Order of the spherical harmonic coefficients

  • sphere_channels (int) – Number of spherical channels

  • hidden_channels (int) – Number of hidden channels used during the SO(2) conv

  • edge_channels (int) – Size of invariant edge embedding

  • (list (mmax_list) – int): List of degrees (l) for each resolution

  • (list – int): List of orders (m) for each resolution

  • act (function) – Non-linear activation function

hidden_channels#
lmax_list#
mmax_list#
sphere_channels#
num_resolutions: int#
m#
act#
fc1_dist#
fc1_r#
fc2_r#
fc1_i#
fc2_i#
forward(x_m, x_edge) → torch.Tensor#
class core.models.escn.escn.EdgeBlock(edge_channels, distance_expansion, max_num_elements, act)#

Bases: torch.nn.Module

Edge Block: Compute invariant edge representation from edge diatances and atomic numbers

Parameters:

  • edge_channels (int) – Size of invariant edge embedding

  • distance_expansion (func) – Function used to compute distance embedding

  • max_num_elements (int) – Maximum number of atomic numbers

  • act (function) – Non-linear activation function

in_channels#
distance_expansion#
act#
edge_channels#
max_num_elements#
fc1_dist#
source_embedding#
target_embedding#
fc1_edge_attr#
forward(edge_distance, source_element, target_element)#
class core.models.escn.escn.EnergyBlock(num_channels: int, num_sphere_samples: int, act)#

Bases: torch.nn.Module

Energy Block: Output block computing the energy

Parameters:

  • num_channels (int) – Number of channels

  • num_sphere_samples (int) – Number of samples used to approximate the integral on the sphere

  • act (function) – Non-linear activation function

num_channels#
num_sphere_samples#
act#
fc1#
fc2#
fc3#
forward(x_pt) → torch.Tensor#
class core.models.escn.escn.ForceBlock(num_channels: int, num_sphere_samples: int, act)#

Bases: torch.nn.Module

Force Block: Output block computing the per atom forces

Parameters:

  • num_channels (int) – Number of channels

  • num_sphere_samples (int) – Number of samples used to approximate the integral on the sphere

  • act (function) – Non-linear activation function

num_channels#
num_sphere_samples#
act#
fc1#
fc2#
fc3#
forward(x_pt, sphere_points) → torch.Tensor#
Contents