xnn.cnn.models.steerable.SteerableCNN#

class xnn.cnn.models.steerable.SteerableCNN(species=(1, 6, 8), cutoff=4.0, grid_size=17, fields=None, n_features=32, n_blocks=3, l_max=2, kernel_size=5, strides=None, padding=None, activation='ssp', gate_activation='sigmoid', normalization=None, smooth_stride=True, bandlimit='compromise', shell_width=0.6, sigma=None, cutoff_fn='cosine', include_center=True, readout_activation='ssp', energy_shift=0.0, energy_scale=1.0, atomic_energies=None)[source]#

Bases: VoxelPotential

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

The species density grids of VoxelGrid (one scalar field per species) pass through len(fields) GatedBlock blocks; the last block outputs scalar fields only, which a global average pool turns into the invariant per-atom features of the shared readout. The defaults are the paper’s recipe scaled to environments: hidden blocks carry fields of order 0, 1 and 2 with multiplicities doubling block by block, kernels of side 5, a stride of 2 (low-pass filtered) in the inner blocks, gated nonlinearities, no batch normalization, the compromise bandlimits and shells of width 0.6 voxels.

Parameters:

  • species (sequence of int, optional) – Atomic numbers with a density channel, by default (1, 6, 8).

  • cutoff (float, optional) – Radius of the voxelized environment (Angstrom), by default 4.0.

  • grid_size (int, optional) – Voxels per axis, by default 17.

  • fields (sequence of sequence of int or None, optional) – Multiplicities (m_0, m_1, ...) of the output stack of each block; the last entry must be scalar only. None derives them from n_features, n_blocks and l_max (default_fields()).

  • n_features (int, optional) – Scalar fields of the last block when fields is None (and the readout width), by default 32.

  • n_blocks (int, optional) – Number of blocks when fields is None, by default 3.

  • l_max (int, optional) – Highest field order of the default hidden blocks, by default 2.

  • kernel_size (int, optional) – Side of the cubic kernels, by default 5.

  • strides (sequence of int or None, optional) – Downsampling per block; None downsamples by 2 in every block but the first and the last.

  • padding (int or None, optional) – Zero padding; None keeps the grid size.

  • activation (str, optional) – Activation of the scalar fields, by default "ssp" (smooth, for a smooth energy surface; the paper’s networks use "relu").

  • gate_activation (str, optional) – Activation of the gates, by default "sigmoid".

  • normalization (str or None, optional) – "batch" for SteerableBatchNorm in every block, by default None.

  • smooth_stride (bool, optional) – Low-pass filter before every downsampling, by default True.

  • bandlimit (str or sequence of int, optional) – Shell bandlimits of the kernels, by default "compromise".

  • shell_width (float, optional) – Width of the radial Gaussian shells in voxels, by default 0.6.

  • sigma (Optional[float]) – See VoxelPotential.

  • cutoff_fn (Optional[str]) – See VoxelPotential.

  • include_center (bool) – See VoxelPotential.

  • readout_activation (str) – See VoxelPotential.

  • energy_shift (float) – See VoxelPotential.

  • energy_scale (float) – See VoxelPotential.

  • atomic_energies – See VoxelPotential.

Variables:
trunk(grid)[source]#

Gated blocks followed by global average pooling of the scalar fields.

Parameters:

grid (Tensor)

Return type:

Tensor

classmethod from_config(cfg)[source]#

Build a SteerableCNN from a configuration object.

cfg.cutoff, cfg.n_features (scalar fields of the last block) and cfg.n_interactions (number of blocks) are the core fields; cfg.extra may set species, grid_size, sigma, cutoff_fn, include_center, fields, l_max, kernel_size, strides, padding, activation, gate_activation, normalization, smooth_stride, bandlimit, shell_width, readout_activation, energy_shift, energy_scale and atomic_energies.

Parameters:

cfg (xnn.common.config.schema.ModelConfig) – The model config.

Returns:

The model.

Return type:

SteerableCNN