Grid Signal on the Sphere
- e3nn.o3.s2_grid(res_beta, res_alpha, dtype=None, device=None)[source]
grid on the sphere
- Parameters
- 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)[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
- 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)[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
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
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
Examples
>>> m = ToS2Grid(6, (100, 101)) >>> x = torch.randn(3, 49) >>> m(x).shape torch.Size([3, 100, 101])
ToS2Grid
andFromS2Grid
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
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
- 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
Examples
>>> m = FromS2Grid((100, 101), 6) >>> x = torch.randn(3, 100, 101) >>> m(x).shape torch.Size([3, 49])
ToS2Grid
andFromS2Grid
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
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