PaiNN¶

class graph_pes.models.PaiNN(channels=32, radial_features=20, layers=3, cutoff=5.0)[source]¶

Bases: GraphPESModel

The Polarizable Atom Interaction Neural Network (PaiNN) model.

Alternating Interaction and Update blocks are used to residually update both vector and scalar per-atom embeddings.

Citation:

@misc{Schutt-21-06,
    title = {
        Equivariant Message Passing for the Prediction
        of Tensorial Properties and Molecular Spectra
    },
    author = {
        Sch{\"u}tt, Kristof T. and Unke, Oliver T.
        and Gastegger, Michael
    },
    year = {2021},
    doi = {10.48550/arXiv.2102.03150},
}

Parameters:

channels (int) – The number of channels of the internal representations.
radial_features (int) – The number of radial features to expand bond distances into.
layers (int) – The number of (interaction + update) layers to use.
cutoff (float) – The cutoff distance for the radial features.

Examples

Configure a PaiNN model for use with graph-pes-train:

model:
  +PaiNN:
    channels: 32
    layers: 3
    cutoff: 5.0

class graph_pes.models.painn.Interaction(radial_features, channels, cutoff)[source]¶

The interaction block of the PaiNN model.

Continuous filters generated from neighbour distances are convolved with existing scalar embeddings to create messages \(x_{j \rightarrow i}\) for each neighbour \(j\) of atom \(i\).

Scalar total messages, \(\Delta s_i\), are created by summing over neighbours, while vector total messages, \(\Delta v_i\), incorporate directional information from neighbour unit vectors and existing vector embeddings.

The code aims to follow Figure 2b of the PaiNN paper as closely as possible.

Parameters:

radial_features (int) – The number of radial features to expand bond distances into.
channels (int) – The number of channels of the internal representations.
cutoff (float) – The cutoff distance for the radial features.

class graph_pes.models.painn.Update(channels)[source]¶

The update block of the PaiNN model.

Projections of vector embeddings are used to update the scalar embeddings, and vice versa. The code aims to follow Figure 2c of the PaiNN paper as closely as possible.

Parameters:: channels (int) – The number of channels of the internal representations.