Grid Signal on the Sphere

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

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

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

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

e3nn.o3.rfft(x, l)[source]

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)[source]

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='component', dtype=None, device=None)[source]

Bases: torch.nn.modules.module.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[source]

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

Type

torch.Tensor

Methods:

forward(x)

Evaluate

forward(x)[source]

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='component', lmax_in=None, dtype=None, device=None)[source]

Bases: torch.nn.modules.module.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[source]

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

Type

torch.Tensor

Methods:

forward(x)

Evaluate

forward(x)[source]

Evaluate

Parameters

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

Returns

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

Return type

torch.Tensor