xnn.cnn.featurizers.voxel#
Voxelized atomic environments: the input fields of the 3D convolution models.
The volumetric models of the cnn family (CNN3D
and the 3D steerable SteerableCNN of Weiler
et al., NeurIPS 2018) read scalar fields sampled on a cubic grid. This
featurizer builds one such grid per atom: a cube of side 2 * cutoff
centered on the atom, with one channel per chemical species holding the
density that the neighbors of that species deposit on the voxels. The
density of a neighbor is a Gaussian of width sigma placed at its
position (the construction of the paper’s amino-acid and CATH experiments,
where “the values of the voxels were set to the densities arising from
placing a Gaussian at each atom position”), multiplied by a smooth envelope
that vanishes at the cutoff so the fields, and with them the energy, stay
continuous when a neighbor enters or leaves the cutoff sphere.
The grids are differentiable with respect to the positions, so the autograd
forces and stress of ForceStressOutput
work unchanged, and periodic images enter only through the edge vectors.
Module Attributes
voxels per block of edges splatted at once (bounds the |
Classes
|
Per-atom voxel grids of the species-resolved neighbor density. |