xnn.cnn.featurizers.voxel.VoxelGrid#

class xnn.cnn.featurizers.voxel.VoxelGrid(species, cutoff=4.0, grid_size=17, sigma=None, cutoff_fn='cosine', include_center=True)[source]#

Bases: Featurizer

Per-atom voxel grids of the species-resolved neighbor density.

For atom i the grid covers the cube [-cutoff, cutoff]^3 around its position with grid_size voxels per axis (voxel width 2 * cutoff / grid_size, voxel centers at (k - (grid_size - 1) / 2) widths). Channel c holds the density of the neighbors of species species[c],

rho_c(x) = sum_{j in N(i), Z_j = species[c]} w(|r_ij|) exp(-|x - r_ij|^2 / (2 sigma^2)),

with r_ij the position of j relative to i and w the cutoff envelope. The atom itself is deposited at the origin of its own species channel when include_center is set, so the grid also encodes the species of the central atom.

Parameters:

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

  • cutoff (float, optional) – Half the side of the cube (the neighbor-list radius), by default 4.0.

  • grid_size (int, optional) – Voxels per axis s, by default 17 (0.47 Angstrom voxels at 4.0).

  • sigma (float or None, optional) – Standard deviation of the Gaussian density of one atom. None (default) uses the voxel width, which keeps the density sampled without aliasing so a rotation of the structure rotates the grids to the interpolation error; the paper’s CATH data set uses half the voxel width (sigma = cutoff / grid_size), a sharper but aliased choice.

  • cutoff_fn (str or None, optional) – "cosine" (default) multiplies each neighbor’s density by the CosineCutoff envelope so the fields are continuous at the cutoff; None deposits every neighbor with unit weight.

  • include_center (bool, optional) – Deposit the central atom at the origin of its species channel, by default True.

Variables:

  • species (list of int) – The channel species.

  • axis (Tensor) – The grid_size voxel-center coordinates along one axis.

property output_dim: int#

Number of channels (one per species).

Type:

int

property n_channels: int#

Number of channels (one per species).

Type:

int

channel_index(atomic_numbers)[source]#

Channel of every atom, raising for species without a channel.

Parameters:

atomic_numbers (Tensor)

Return type:

Tensor

voxelize(atomic_numbers, edge_index, edge_vec)[source]#

Build the grids from tensors (the core reused by forward()).

Parameters:

  • atomic_numbers (Tensor) – Per-atom atomic numbers, shape (N,).

  • edge_index (Tensor) – Edge index (2, E); row 0 the neighbor, row 1 the central atom.

  • edge_vec (Tensor) – pos[dst] - pos[src] per edge (with the periodic image shift), shape (E, 3); the neighbor sits at -edge_vec relative to the center.

Returns:

The grids, shape (N, n_channels, s, s, s) in the dtype of edge_vec.

Return type:

Tensor

forward(data)[source]#

Voxelize the environment of every atom of a batch.

Parameters:

data (AtomicGraph) – Batched atomic graph (atomic numbers, edge index, geometry).

Returns:

Density grids of shape (N, n_channels, s, s, s).

Return type:

Tensor

extra_repr()[source]#

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str