xnn.cnn.models.cnn3d.CNN3D#

class xnn.cnn.models.cnn3d.CNN3D(species=(1, 6, 8), cutoff=4.0, grid_size=17, channels=None, n_features=32, n_blocks=3, kernel_size=5, strides=None, padding=None, activation='ssp', normalization=None, smooth_stride=True, 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

Conventional 3D convolutional potential over voxelized environments.

len(channels) Conv3dBlock blocks map the species density grids of VoxelGrid to channels[-1] feature maps, which a global average pool turns into the per-atom features of the shared readout. The defaults mirror the steerable model (default_fields()): the channel count of each block equals the number of components of the steerable fields of the same block, the paper’s recipe for its CNN baselines.

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.

  • channels (sequence of int or None, optional) – Output channels of each block; None derives them from n_features and n_blocks.

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

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

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

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

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

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

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

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

  • 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:

blocks (torch.nn.Sequential) – The convolution blocks.

trunk(grid)[source]#

Convolution blocks followed by global average pooling.

Parameters:

grid (Tensor)

Return type:

Tensor

classmethod from_config(cfg)[source]#

Build a CNN3D 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, grid_size, sigma, cutoff_fn, include_center, channels, kernel_size, strides, padding, activation, normalization, smooth_stride, readout_activation, energy_shift, energy_scale and atomic_energies.

Parameters:

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

Returns:

The model.

Return type:

CNN3D

On this page
Show Source