SphericalHarmonics#

class cuequivariance_torch.SphericalHarmonics#

Compute the spherical harmonics of the input vectors as a torch module.

Forward Pass

forward(vectors)#
Parameters:

vectors (torch.Tensor) – Input vectors of shape (batch, 3).

Returns:

The spherical harmonics of the input vectors of shape (batch, dim), where dim is the sum of 2*l+1 for l in ls.

Return type:

torch.Tensor

__init__(
ls,
normalize=True,
device=None,
math_dtype=None,
use_fallback=None,
method=None,
)#
Parameters:

  • ls (list of int) – List of spherical harmonic degrees.

  • normalize (bool, optional) – Whether to normalize the input vectors. Defaults to True.

  • device (torch.device, optional) – The device to use for the operation.

  • math_dtype (torch.dtype, optional) – The dtype to use for the operation.

  • method (str, optional) – The method to use for the operation, by default “uniform_1d” (using a CUDA kernel).

  • use_fallback (bool, optional, deprecated) – Whether to use a “fallback” implementation, now maps to method: If True the “naive” method is used. If False or None (default) the “uniform_1d” method is used.

classmethod __new__(*args, **kwargs)#