xnn.gnn.featurizers.cartesian.CartesianAngularBasis#

class xnn.gnn.featurizers.cartesian.CartesianAngularBasis(l_max)[source]#

Bases: Module

Cartesian monomial angular basis x^lx y^ly z^lz (CACE, paper eq. 2).

A drop-in Cartesian alternative to spherical harmonics for encoding edge directions: evaluates every monomial with lx + ly + lz <= l_max on the (normalized) edge vectors. Entries are grouped contiguously by total l in ascending order, matching lxlylz_list().

Parameters:

l_max (int) – Maximum total angular momentum of the basis.

Variables:

  • l_max (int) – The maximum total angular momentum.

  • lxlylz (list of tuple of int) – The (lx, ly, lz) triple of every output entry, in order.

  • l_of_entry (Tensor) – Long buffer of shape (n_angular,) giving the total l of each entry (used by the CACE blocks that share weights per l group).

  • n_angular (int) – Output width, (l_max+1)(l_max+2)(l_max+3)/6.

forward(unit_vec)[source]#

Evaluate the monomials on (unit) edge vectors.

Parameters:

unit_vec (Tensor) – Normalized edge vectors of shape (E, 3).

Returns:

The angular basis of shape (E, n_angular); column i is the monomial of self.lxlylz[i].

Return type:

Tensor