xnn.gnn.models.painn#

PaiNN (Schuett, Unke, Gastegger, ICML 2021) – equivariant message passing.

A faithful implementation of the polarizable atom interaction neural network of

  • K. T. Schuett, O. T. Unke, M. Gastegger, “Equivariant message passing for the prediction of tensorial properties and molecular spectra”, ICML 2021 (arXiv:2102.03150),

built from the paper’s equations on the xnn abstractions (BesselRBF, CosineCutoff, scatter_sum(), …) and consistent with the conventions of the authors’ reference code (the schnetpack package: feature layout, initialization, numerical stabilization), without copying it.

Representation (paper section IV.A). Every atom carries F rotationally invariant scalar features s_i and F equivariant vector features v_i in R^3. The scalars start as a learned embedding of the nuclear charge, s_i^0 = a_{Z_i}, the vectors as zero. T pairs of a message and an update block (Fig. 2) then refine both, each adding a residual.

Message block (eq 7 and 8, Fig. 2b):

ds_i = sum_j phi_s(s_j) o W_s(r_ij)
dv_i = sum_j v_j o phi_vv(s_j) o W_vv(r_ij)
     + sum_j phi_vs(s_j) o W_vs(r_ij) r_ij / |r_ij|

where r_ij = r_j - r_i, the shared network phi = W silu(W s + b) + b maps F to 3F features that are split three ways, and the filters W(r) = W rbf(r) f_cut(r) are linear combinations of the radial basis sin(n pi r / r_cut) / r, 1 <= n <= 20, times the cosine cutoff.

Update block (eq 9 and 10, Fig. 2c):

ds_i = a_ss(s_i, |V v_i|) + a_sv(s_i, |V v_i|) <U v_i, V v_i>
dv_i = a_vv(s_i, |V v_i|) U v_i

with the linear (bias-free) maps U and V of the vector features and the shared network a = W silu(W [s_i, |V v_i|] + b) + b from 2F to 3F features, split into a_vv, a_sv, a_ss.

Readout (Fig. 2a): E = sum_i W silu(W s_i + b) + b with widths F -> F/2 -> 1, plus the per-atom standardization and per-element reference of every xnn model (energy_scale, energy_shift, atom_ref).

Tensorial properties (section IV.B). With dipole=True the dipole moment is built from latent charges and atomic dipoles (eq 13):

mu = sum_i mu_i(v_i) + q_i(s_i) r_i

and with polarizability=True the polarizability tensor from the rank-1 decomposition of eq 14:

alpha = sum_i alpha_i(s_i) I + nu_i(v_i) (x) r_i + r_i (x) nu_i(v_i),

where the scalars and vectors come from a stack of two gated equivariant blocks (Fig. 3) and r_i are the positions relative to the center of mass of the structure (the paper assumes the center of mass at the origin). The latent charges are shifted so that they sum to the structure’s net charge (total_charge, zero by default), as in the reference code.

Conventions worth knowing:

  • xnn edge vectors are pos[dst] - pos[src] = r_i - r_j; the paper’s direction r_ij / |r_ij| with r_ij = r_j - r_i is their negative, so the model flips the sign. The sign is observable only through the vector features (the dipoles and polarizabilities), not the energy.

  • The norm in the update block is the stabilized sqrt(|V v|^2 + epsilon) of the reference code (epsilon = 1e-8), so the gradient of an atom without neighbors (v = 0) is finite.

  • The radial basis is the paper’s Bessel form by default; the schnetpack default, Gaussians every r_cut / (n_rbf - 1) with the spacing as width, is available as radial_basis="gaussian".

  • The ablations of the paper’s Table IV are options: scalar_product (the <U v, V v> term of eq 9) and vector_propagation (the v_j o phi_vv term of eq 8).

Works for molecules and periodic systems through data.edge_vectors() (the energy, forces and stress); the dipole and polarizability heads use absolute positions and are meant for molecules.

Classes

GatedEquivariantBlock(n_in_s, n_in_v, ...[, ...])

Gated equivariant block (paper Fig.

PaiNN([n_features, n_interactions, n_rbf, ...])

PaiNN interatomic potential with optional dipole and polarizability heads.
Show Source