Grid Signal on the Sphere

e3nn.o3.s2_grid(res_beta, res_alpha, dtype=None, device=None)

grid on the sphere

Parameters:

• res_beta (int) – \(N\)

• res_alpha (int) – \(M\)

• dtype (torch.dtype or None) – dtype of the returned tensors. If None then set to torch.get_default_dtype().

• device (torch.device or None) – device of the returned tensors. If None then set to the default device of the current context.

Returns:

• betas (torch.Tensor) – tensor of shape (res_beta)

• alphas (torch.Tensor) – tensor of shape (res_alpha)

e3nn.o3.spherical_harmonics_s2_grid(lmax, res_beta, res_alpha, dtype=None, device=None)

spherical harmonics evaluated on the grid on the sphere

\[ \begin{align}\begin{aligned}f(x) = \sum_{l=0}^{l_{\mathit{max}}} F^l \cdot Y^l(x)\\f(\beta, \alpha) = \sum_{l=0}^{l_{\mathit{max}}} F^l \cdot S^l(\alpha) P^l(\cos(\beta))\end{aligned}\end{align} \]

Parameters:

• lmax (int) – \(l_{\mathit{max}}\)

• res_beta (int) – \(N\)

• res_alpha (int) – \(M\)

Returns:

• betas (torch.Tensor) – tensor of shape (res_beta)

• alphas (torch.Tensor) – tensor of shape (res_alpha)

• shb (torch.Tensor) – tensor of shape (res_beta, (lmax + 1)**2)

• sha (torch.Tensor) – tensor of shape (res_alpha, 2 lmax + 1)

e3nn.o3.rfft(x, l) → Tensor

Real fourier transform

Parameters:

• x (torch.Tensor) – tensor of shape (..., 2 l + 1)

• res (int) – output resolution, has to be an odd number

Returns:

tensor of shape (..., res)

Return type:

torch.Tensor

Examples

>>> lmax = 8
>>> res = 101
>>> _betas, _alphas, _shb, sha = spherical_harmonics_s2_grid(lmax, res, res)
>>> x = torch.randn(res)
>>> (rfft(x, lmax) - x @ sha).abs().max().item() < 1e-4
True

e3nn.o3.irfft(x, res)

Inverse of the real fourier transform

Parameters:

• x (torch.Tensor) – tensor of shape (..., 2 l + 1)

• res (int) – output resolution, has to be an odd number

Returns:

positions on the sphere, tensor of shape (..., res, 3)

Return type:

torch.Tensor

Examples

>>> lmax = 8
>>> res = 101
>>> _betas, _alphas, _shb, sha = spherical_harmonics_s2_grid(lmax, res, res)
>>> x = torch.randn(2 * lmax + 1)
>>> (irfft(x, res) - sha @ x).abs().max().item() < 1e-4
True

class e3nn.o3.ToS2Grid(lmax=None, res=None, normalization: str = 'component', dtype=None, device=None)

Bases: Module

Transform spherical tensor into signal on the sphere

The inverse transformation of FromS2Grid

Parameters:

• lmax (int)

• res (int, tuple of int) – resolution in beta and in alpha

• normalization ({'norm', 'component', 'integral'})

• dtype (torch.dtype or None, optional)

• device (torch.device or None, optional)

Examples

>>> m = ToS2Grid(6, (100, 101))
>>> x = torch.randn(3, 49)
>>> m(x).shape
torch.Size([3, 100, 101])

ToS2Grid and FromS2Grid are inverse of each other

>>> m = ToS2Grid(6, (100, 101))
>>> k = FromS2Grid((100, 101), 6)
>>> x = torch.randn(3, 49)
>>> y = k(m(x))
>>> (x - y).abs().max().item() < 1e-4
True

grid

positions on the sphere, tensor of shape (res_beta, res_alpha, 3)

Type:

torch.Tensor

Methods:

forward(x)

Evaluate

forward(x)

Evaluate

Parameters:

x (torch.Tensor) – tensor of shape (..., (l+1)^2)

Returns:

tensor of shape [..., beta, alpha]

Return type:

torch.Tensor

class e3nn.o3.FromS2Grid(res=None, lmax=None, normalization: str = 'component', lmax_in=None, dtype=None, device=None)

Bases: Module

Transform signal on the sphere into spherical tensor

The inverse transformation of ToS2Grid

Parameters:

• res (int, tuple of int) – resolution in beta and in alpha

• lmax (int)

• normalization ({'norm', 'component', 'integral'})

• lmax_in (int, optional)

• dtype (torch.dtype or None, optional)

• device (torch.device or None, optional)

Examples

>>> m = FromS2Grid((100, 101), 6)
>>> x = torch.randn(3, 100, 101)
>>> m(x).shape
torch.Size([3, 49])

ToS2Grid and FromS2Grid are inverse of each other

>>> m = FromS2Grid((100, 101), 6)
>>> k = ToS2Grid(6, (100, 101))
>>> x = torch.randn(3, 100, 101)
>>> x = k(m(x))  # remove high frequencies
>>> y = k(m(x))
>>> (x - y).abs().max().item() < 1e-4
True

grid

positions on the sphere, tensor of shape (res_beta, res_alpha, 3)

Type:

torch.Tensor

Methods:

forward(x)

Evaluate

forward(x) → Tensor

Evaluate

Parameters:

x (torch.Tensor) – tensor of shape [..., beta, alpha]

Returns:

tensor of shape (..., (l+1)^2)

Return type:

torch.Tensor