xnn.cnn.models.base.VoxelPotential#

class xnn.cnn.models.base.VoxelPotential(species=(1, 6, 8), cutoff=4.0, grid_size=17, n_features=32, sigma=None, cutoff_fn='cosine', include_center=True, readout_activation='ssp', energy_shift=0.0, energy_scale=1.0, atomic_energies=None, featurizer=None)[source]#

Bases: InteratomicPotential

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

A subclass implements trunk(), mapping the (N, C, ...) grids of the featurizer (voxel density grids (N, C, s, s, s) by default) to (N, n_features) invariant per-atom features; this class featurizes the environments, reads the features out into per-atom energies (atom-wise n_features -> n_features/2 -> 1 network with a zero-initialized last layer, the DTNN standardization and the per-element reference energies, as in SchNet) and sum-pools them per structure.

Parameters:

  • species (sequence of int) – Atomic numbers with a density channel.

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

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

  • n_features (int, optional) – Width of the pooled features (input of the readout), by default 32.

  • sigma (float or None, optional) – Width of the atomic Gaussians (None: one voxel).

  • cutoff_fn (str or None, optional) – Envelope of the neighbor densities ("cosine" or None).

  • include_center (bool, optional) – Deposit the central atom in its own channel, by default True.

  • readout_activation (str, optional) – Activation of the readout network, by default "ssp" (the exact shifted softplus, smooth).

  • energy_shift (float, optional) – Per-atom energy standardization (buffers), by default 0 and 1.

  • energy_scale (float, optional) – Per-atom energy standardization (buffers), by default 0 and 1.

  • atomic_energies (array-like or None, optional) – Per-species reference energies loaded into atom_ref.

  • featurizer (Featurizer or None, optional) – The grid featurizer to use instead of the VoxelGrid built from the voxel options (it must expose n_channels and grid_size).

Variables:
property grid_size: int#

Samples per axis of the environment grids.

Type:

int

property n_channels: int#

Input channels of the trunk (one per species).

Type:

int

set_energy_scale_shift(scale, shift)[source]#

Set the per-atom energy standardization from training statistics.

Parameters:

  • scale (float) – Standard deviation of the training-set energy per atom.

  • shift (float) – Mean training-set energy per atom.

Return type:

None

set_atomic_energies(values)[source]#

Initialize the per-element reference energies atom_ref.

Parameters:

values (array-like) – One reference energy per entry of species, in order.

Raises:

ValueError – If the number of values does not match the number of species.

Return type:

None

trunk(grid)[source]#

Map density grids to invariant per-atom features.

Parameters:

grid (Tensor) – Shape (N, n_channels, s, s, s).

Returns:

Shape (N, n_features).

Return type:

Tensor

node_features_energy(grid, atomic_numbers)[source]#

Features and per-atom energies of already voxelized environments.

Parameters:
Return type:

tuple[Tensor, 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]