xnn.cnn.models.spherical.SphericalCNN#

class xnn.cnn.models.spherical.SphericalCNN(species=(1, 6, 7, 8, 16), cutoff=10.0, radius=0.48, bandwidth=10, exponent=1.0, features=None, bandwidths=None, n_features=160, n_blocks=5, n_alpha=None, n_beta=2, n_gamma=2, max_beta=0.39269908169872414, normalization='batch', activation='relu', set_readout=None, cutoff_fn=None, readout_activation='ssp', energy_shift=0.0, energy_scale=1.0, atomic_energies=None)[source]#

Bases: VoxelPotential

Spherical CNN interatomic potential (Cohen et al., ICLR 2018, Sec. 5.4).

Every atom’s environment is the stack of spherical signals of SphericalGrid (the potential of the neighbors of each species on a sphere around the atom, sampled on the Driscoll-Healy grid of bandwidth bandwidth). len(features) SphericalResBlock blocks map it to SO(3) signals of decreasing bandwidth, so3_integrate() pools them into invariant per-atom features, and the shared readout gives the per-atom energies. The defaults are Table 3 of the paper: five blocks with 20, 40, 60, 80, 160 channels at bandwidths 10, 8, 6, 4, 2, batch normalization and ReLU. With set_readout the paper’s DeepSet readout is used instead: the per-atom features pass through an MLP, are summed over the structure and mapped to the energy by a second MLP (the per-atom energy is then the structure’s energy divided by its atom count, plus the atom’s reference energy).

Parameters:

  • species (sequence of int, optional) – Atomic numbers with a potential channel, by default (1, 6, 7, 8, 16) (the elements of QM7).

  • cutoff (float, optional) – Neighbor-list radius of the potential sums (Angstrom), by default 10.0.

  • radius (float, optional) – Radius of the sphere around each atom (Angstrom), by default 0.48 (half the shortest interatomic distance of QM7, so no atom lies on a sphere).

  • bandwidth (int, optional) – Bandwidth of the input signals, by default 10.

  • exponent (float, optional) – Power of the inverse distance of the potential, by default 1 (the paper’s 1/|x - p|; the reference data script uses 2).

  • features (sequence of int or None, optional) – Channels and output bandwidth of every block; None derives them from n_features, n_blocks and bandwidth (default_schedule()).

  • bandwidths (sequence of int or None, optional) – Channels and output bandwidth of every block; None derives them from n_features, n_blocks and bandwidth (default_schedule()).

  • n_features (int, optional) – Channels of the last block when features is None, by default 160.

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

  • n_alpha (Optional[int]) – Filter support of the correlations (see SphericalResBlock).

  • n_beta (int) – Filter support of the correlations (see SphericalResBlock).

  • n_gamma (int) – Filter support of the correlations (see SphericalResBlock).

  • max_beta (float) – Filter support of the correlations (see SphericalResBlock).

  • normalization (str or None, optional) – "batch" (default, the paper) or None.

  • activation (str, optional) – Activation inside the blocks, by default "relu" (the paper).

  • set_readout (sequence of int or None, optional) – (phi_hidden, latent, psi_hidden) of the DeepSet readout, e.g. the paper’s (150, 100, 50); None (default) reads every atom out separately with the shared atom-wise readout.

  • cutoff_fn (str or None, optional) – "cosine" multiplies each neighbor’s potential by the cosine envelope at cutoff (continuous energies for dynamics); None (default) is the paper’s untruncated potential.

  • readout_activation (str) – See VoxelPotential.

  • energy_shift (float) – See VoxelPotential.

  • energy_scale (float) – See VoxelPotential.

  • atomic_energies – See VoxelPotential.

Variables:

  • blocks (torch.nn.Sequential) – The ResNet blocks.

  • bandwidths (features,) – Channels and bandwidths of every block.

trunk(grid)[source]#

ResNet blocks followed by the integral over SO(3).

Parameters:

grid (Tensor)

Return type:

Tensor

forward(data)[source]#

Compute per-atom and total energies for a batch of structures.

Parameters:

data (AtomicGraph) – Batched atomic graph.

Returns:

"node_energy" (N,), "energy" (B,) and the pooled invariant "node_features" (N, n_features).

Return type:

dict[str, Tensor]

classmethod from_config(cfg)[source]#

Build a SphericalCNN from a configuration object.

cfg.cutoff, cfg.n_features (channels of the last block) and cfg.n_interactions (number of blocks) are the core fields; cfg.extra may set species, radius, bandwidth, exponent, features, bandwidths, n_alpha, n_beta, n_gamma, max_beta, normalization, activation, set_readout, cutoff_fn, readout_activation, energy_shift, energy_scale and atomic_energies.

Return type:

SphericalCNN