xnn.cnn.models.steerable#

3D steerable CNN (Weiler, Geiger, Welling, Boomsma, Cohen, NeurIPS 2018).

An SE(3)-equivariant convolutional network over fields on a cubic grid, built from the manuscript (arXiv:1807.02547v2) and the conventions of its reference implementation, without using its code:

  • feature spaces are stacks of fields transforming under the irreducible representations D^l of SO(3) (Sec. 3.2); a stack is described by its multiplicities (m_0, m_1, m_2, ...) and laid out channel-wise as the m_l fields of each order in turn, each with its 2l + 1 components;

  • an equivariant linear map is a cross-correlation with a rotation-steerable kernel (Theorem 2), kappa(r x) = D^j(r) kappa(x) D^l(r)^{-1} for every block mapping order l to order j (eq 12);

  • the solutions are, per angular frequency J from |j - l| to j + l, the spherical harmonics Y^J of the direction, carried back from the coupled basis to the (2j + 1) x (2l + 1) block by the change of basis Q of eq 14, times a free radial profile (eqs 15-16). Q is made of Clebsch-Gordan coefficients (footnote 4; the real-basis wigner_3j of e3nn), so no numerical solve is needed; the tests solve eq 14 numerically as the paper does and compare;

  • on the discrete grid (Sec. 4.4.1) the radial profiles are Gaussian shells of width 0.6 voxels at the integer radii 0 .. s // 2, and every shell only carries the frequencies J <= J_max of a radius-dependent bandlimit that suppresses aliasing; every basis kernel is scaled to unit norm. A convolution learns one weight per basis kernel and per (output field, input field) pair and stacks the kernel it spans for a standard conv3d (Sec. 4.4.3);

  • nonlinearities are the gated kind (Sec. 4.3, supplement Fig. 5): scalar fields pass through an element-wise activation (ReLU in the paper), every non-scalar field is multiplied by the sigmoid of an extra scalar field the same convolution produces;

  • downsampling is preceded by a Gaussian low-pass filter (Sec. 4.4.2) and batch normalization, when used, normalizes non-scalar fields by their norm (supplement eq 17).

SteerableCNN is the interatomic potential built from these blocks on the voxelized environments of VoxelGrid: the energy is a pooled scalar field, so it is exactly invariant under the rotations of the grid onto itself, approximately (to the bandlimit) under every other rotation, and the forces co-rotate. The paper’s networks are SE(3)- but not O(3)-equivariant (a reflection is not a rotation), and so is this potential: mirror images are not constrained to the same energy.

Requires e3nn (pip install "xnn[gnn]") for the spherical harmonics and Clebsch-Gordan coefficients.

Module Attributes

BANDLIMITS

radius-dependent angular bandlimits J_max of the shells at radii 0, 1, 2, ... voxels (Sec.

Functions

angular_kernel_basis(l_in, l_out, J, points)

The angular part of the steerable kernel basis (eqs 13-16).

field_slices(fields)

(l, start, stop) channel ranges of the non-empty orders of a stack.

n_basis_kernels(l_in, l_out, size[, bandlimit])

Number of basis kernels B between two orders (Sec.

rotate_fields(x, R, fields)

Rotate a stack of fields by a symmetry of the cube: grid and fibers.

shell_bandlimits(size[, bandlimit])

Radii and angular bandlimits of the Gaussian shells of a kernel.

steerable_kernel_basis(l_in, l_out, size[, ...])

Sampled basis of the steerable kernels between two field orders (Sec.

Classes

GatedBlock(fields_in, fields_out[, ...])

Steerable convolution with the gated nonlinearity (Sec.

SteerableBatchNorm(fields[, eps, momentum, ...])

Equivariant batch normalization (supplement Sec.

SteerableCNN([species, cutoff, grid_size, ...])

3D steerable CNN interatomic potential (Weiler et al., NeurIPS 2018).

SteerableConv3d(fields_in, fields_out[, ...])

SE(3)-equivariant 3D convolution between stacks of fields.
Show Source