xnn.gnn.models.aimnet2.ShellConvolution#

class xnn.gnn.models.aimnet2.ShellConvolution(n_channels, n_shells, n_combinations)[source]#

Bases: Module

The AIMNet2 convolution of per-atom features over the radial shells (paper eqs 4-5).

For center i with neighbors j carrying features a_jcs (one value per channel c and radial shell s) and the radial basis g_ijs of the pair, the scalar features are the shell-wise sums

v_ics = sum_j g_ijs a_jcs

and the vector features the squared norms of n_combinations learned combinations of the shells along the unit bond vectors u_ij,

v_ich = | sum_s w_csh sum_j g_ijs u_ij a_jcs |^2.

Both are flattened and concatenated (channel-major) into a vector of n_channels * (n_shells + n_combinations) entries per atom. The atomic embeddings (one value per shell) and the partial charges (the same value in every shell) go through the same block.

Parameters:

  • n_channels (int) – Feature channels c of the convolved quantity.

  • n_shells (int) – Radial shells s (the number of Gaussian basis functions).

  • n_combinations (int) – Learned shell combinations h of the vector features.

Variables:

  • weight (torch.nn.Parameter) – The combination weights w_csh, shape (n_channels, n_shells, n_combinations). Initialized to standardized random vectors over the shells (the reference code picks a maximally spread set of sinusoids; any spread set serves, and trained weights transplant).

  • output_dim (int) – n_channels * (n_shells + n_combinations).

forward(feats, g, gu, dst, num_nodes)[source]#

Convolve neighbor features into per-center environment features.

Parameters:

  • feats (Tensor) – Features of the neighbor (source) atom of every edge, shape (E, n_channels, n_shells).

  • g (Tensor) – Radial basis of every edge, shape (E, n_shells).

  • gu (Tensor) – Radial basis times the unit bond vector, shape (E, n_shells, 3).

  • dst (Tensor) – Center atom of every edge, shape (E,).

  • num_nodes (int) – Number of atoms N.

Returns:

Environment features, shape (N, output_dim).

Return type:

Tensor