schnetpack.nn

Basic layers

Dense

Fully connected linear layer with activation function.

Equivariant layers

Cartesian:

GatedEquivariantBlock

Gated equivariant block as used for the prediction of tensorial properties by PaiNN.

Irreps:

RealSphericalHarmonics	Generates the real spherical harmonics for a batch of vectors.
SO3TensorProduct	SO3-equivariant Clebsch-Gordon tensor product.
SO3Convolution	SO3-equivariant convolution using Clebsch-Gordon tensor product.
SO3GatedNonlinearity	SO3-equivariant gated nonlinearity.
SO3ParametricGatedNonlinearity	SO3-equivariant parametric gated nonlinearity.

Radial basis

GaussianRBF	Gaussian radial basis functions.
GaussianRBFCentered	Gaussian radial basis functions centered at the origin.
BesselRBF	Sine for radial basis functions with coulomb decay (0th order bessel).

Cutoff

CosineCutoff	Behler-style cosine cutoff module.
MollifierCutoff	Mollifier cutoff module scaled to have a value of 1 at \(r=0\).

Activations

shifted_softplus

Compute shifted soft-plus activation function.

Ops

scatter_add

Sum over values with the same indices.

Factory functions

build_mlp	Build multiple layer fully connected perceptron neural network.
build_gated_equivariant_mlp	Build neural network analog to MLP with `GatedEquivariantBlock`s instead of dense layers.
replicate_module