SchNet

class graph_pes.models.SchNet(
cutoff=5.0,
channels=64,
expansion_features=50,
layers=3,
expansion=None,
)[source]

Bases: GraphPESModel

The SchNet model: a pairwise, scalar, message passing GNN.

A stack of SchNetInteraction blocks are used to update the node features of each atom, \(x_i\) by sequentially aggregating over neighbouring atom’s features, \(x_j\):, and distances, \(r_{ij}\):.

Citation:

@article{Schutt-18-03,
    title = {{{SchNet}} {\textendash} {{A}} Deep Learning Architecture for Molecules and Materials},
    author = {Sch{\"u}tt, K. T. and Sauceda, H. E. and Kindermans, P.-J. and Tkatchenko, A. and M{\"u}ller, K.-R.},
    year = {2018},
    journal = {The Journal of Chemical Physics},
    volume = {148},
    number = {24},
    pages = {241722},
    doi = {10.1063/1.5019779},
}
Parameters:

  • channels (int) – Number of features per node.

  • expansion_features (int) – Number of features used for the radial basis expansion.

  • cutoff (float) – Neighborhood cutoff radius.

  • layers (int) – Number of interaction blocks to apply.

  • expansion (type[DistanceExpansion] | None) – The type of radial basis expansion to use. Defaults to GaussianSmearing as in the original paper.

Components:

class graph_pes.models.schnet.SchNetInteraction(channels, expansion_features, cutoff, basis_type)[source]

SchNet interaction block.

Updates the embedding of each atom, \(x_i\) by sequentially applying the following:

  1. a linear transform of each node’s features \(x_i \leftarrow W x_i\)

  2. message creation \(m_{ij} = x_j \odot \mathbb{F}(r_{ij})\)

  3. message aggregation \(m_i = \sum_{j \in \mathcal{N}(i)} m_{ij}\)

  4. a multi-layer perceptron that further embeds these new node features \(x_i \leftarrow \mathrm{MLP}(h_i)\)

Parameters:

  • channels (int) – Number of channels of the internal representations.

  • expansion_features (int) – Number of features used for the radial basis expansion.

  • cutoff (float) – Neighborhood cutoff radius.

  • basis_type (type[DistanceExpansion]) – The type of radial basis expansion to use.

class graph_pes.models.schnet.CFConv(filter_generator)[source]

CFConv: The Continous Filter Convolution

Originally proposed as part of SchNet, this message passing layer creates messages of the form:

\[m_{ij} = x_j \odot \mathbb{F}(r_{ij})\]

where \(\mathbb{F}\) is some learnable filter function which expands the radial distance \(r_{ij}\) into a set of features that match the dimensionality of the node features \(x_j\).

Example

from graph_pes.models.components.distances import GaussianSmearing
from graph_pes.utils.nn import MLP
from torch import nn

filter_generator = nn.Sequential(
    GaussianSmearing(8, 5.0),
    MLP([8, 16, 16, 8]),
)
cfconv = CFConv(filter_generator)
Parameters:

filter_generator (torch.nn.Module) – The filter function \(\mathbb{F}\).