xnn.cnn.featurizers.spherical.SphericalGrid#

class xnn.cnn.featurizers.spherical.SphericalGrid(species, cutoff=10.0, radius=0.48, bandwidth=10, exponent=1.0, cutoff_fn=None)[source]#

Bases: Featurizer

Per-atom spherical signals of the species-resolved neighbor potential.

Parameters:

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

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

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

  • bandwidth (int, optional) – Bandwidth b of the grid (2b samples per axis), by default 10.

  • exponent (float, optional) – Power of the inverse distance, by default 1.0.

  • cutoff_fn (str or None, optional) – "cosine" multiplies each neighbor’s potential by the CosineCutoff envelope at cutoff; None (default) keeps the untruncated potential of the paper.

Variables:

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

  • directions (Tensor) – Unit vectors of the grid points, (2b, 2b, 3).

property output_dim: int#

Number of channels (one per species).

Type:

int

property n_channels: int#

Number of channels (one per species).

Type:

int

property grid_size: int#

Samples per axis of the grid, 2b.

Type:

int

channel_index(atomic_numbers)[source]#

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

Parameters:

atomic_numbers (Tensor)

Return type:

Tensor

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

Build the signals 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, shape (E, 3); the neighbor sits at -edge_vec relative to the center.

Returns:

Shape (N, n_channels, 2b, 2b) in the dtype of edge_vec.

Return type:

Tensor

forward(data)[source]#

Signals of every atom of a batch, (N, n_channels, 2b, 2b).

Parameters:

data (AtomicGraph)

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