xnn.cnn.models.base

Shared body of the grid-input models of the cnn family.

VoxelPotential composes a grid featurizer of the atomic environments (VoxelGrid by default; the spherical signals of SphericalGrid for the spherical CNN) with a convolutional trunk that a subclass supplies (a conventional 3D CNN in cnn3d, the SE(3)-equivariant 3D steerable CNN of Weiler et al., NeurIPS 2018, in steerable, the spherical CNN of Cohen et al., ICLR 2018, in spherical) and the per-atom energy readout shared with SchNet: a two-layer atom-wise network on the pooled features, the per-atom energy standardization E_i = energy_scale * E^hat_i + energy_shift and the per-element reference energy atom_ref.

The module also holds the grid operations both trunks use: the Gaussian low-pass filter the paper applies before every strided downsampling (Sec. 4.4.2), the global average pooling of its architectures, the bookkeeping of field multiplicities, and the exact rotation of a voxel grid by a symmetry of the cube (what the equivariance tests and notebooks use).

Functions

default_fields(n_features, n_blocks[, l_max])	The default field multiplicities of the voxel potentials.
default_strides(n_blocks)	Stride 2 on every block but the first and the last (the paper's Tetris network downsamples between layers 1-2 and 2-3).
field_dim(fields)	Number of channels of a stack of fields.
global_average_pool(x)	Average the spatial axes of (B, C, X, Y, Z) fields into (B, C).
grid_coordinates(size[, dtype, device])	Integer-centered coordinates of a cubic grid.
low_pass_filter(x, scale[, stride])	Gaussian blur of the spatial axes, optionally followed by a stride.
rotate_voxels(x, R)	Rotate scalar voxel grids by a symmetry of the cube, exactly.

Classes

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

Shared body of the grid-input models: featurizer, trunk, readout.