torch_harmonics.SpectralConvS2#

class torch_harmonics.SpectralConvS2(
in_shape,
out_shape,
in_channels,
out_channels,
num_groups=1,
grid_in='equiangular',
grid_out='equiangular',
bias=False,
)[source]#

Bases: Module

Spectral convolution layer on \(S^2\) implemented via real SHT (Driscoll–Healy formulation, see https://api.semanticscholar.org/CorpusID:122817218).

Given a multi-channel input signal \(u^{c_i}(\theta, \lambda)\) on the sphere, the layer computes the output channels \(v^{c_o}(\theta, \lambda)\) in three steps:

  1. Forward SHT (cf. RealSHT) – transform each input channel to spectral space:

\[\hat{u}_l^{m,\,c_i} = \text{SHT}\!\left[\, u^{c_i}(\theta, \lambda) \,\right]\]

  1. Spectral contraction – mix channels with learnable weights \(K_l^{c_o,\,c_i}\) that are diagonal in \((l, m)\) (i.e.the same weight is applied to every order \(m\) at a given degree \(l\)):

\[\hat{v}_l^{m,\,c_o} = \sum_{c_i} K_l^{c_o,\,c_i}\; \hat{u}_l^{m,\,c_i}\]

  1. Inverse SHT (cf. InverseRealSHT) – transform back to the spatial domain:

\[v^{c_o}(\theta, \lambda) = \text{ISHT}\!\left[\, \hat{v}_l^{m,\,c_o} \,\right]\]

Because the spectral weights depend only on degree \(l\) and not on order \(m\), this corresponds to an isotropic (azimuthally symmetric) convolution kernel on the sphere. When num_groups > 1, the channel contraction is performed independently within each group (grouped convolution).

Spectral bias. When bias=True, a learnable spectral bias \(b_l^{m,\,c_i}\) is added to the SHT coefficients before the channel contraction. The bias is modulated by the spatial integral (zeroth moment) of each input channel:

\[I^{c_i} = \int_0^{2\pi}\!\int_0^{\pi} u^{c_i}(\theta,\lambda)\,\sin\theta\;d\theta\;d\lambda\]
\[\hat{u}_l^{m,\,c_i} \;\leftarrow\; \hat{u}_l^{m,\,c_i} + I^{c_i}\, b_l^{m,\,c_i}\]

This allows the layer to learn a spectral response that depends on the global mean of each input channel, effectively coupling the zero-frequency content into all spectral modes.

Parameters:

  • in_shape (Tuple[int]) – Spatial input grid shape (nlat, nlon).

  • out_shape (Tuple[int]) – Spatial output grid shape (nlat, nlon).

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • num_groups (int, optional) – Number of channel groups for grouped spectral weights, by default 1.

  • grid_in (str, optional) – Grid used for the forward SHT ("equiangular", "legendre-gauss", "lobatto", "equiangular-trapezoidal"), by default "equiangular".

  • grid_out (str, optional) – Grid used for the inverse SHT, same options as grid_in.

  • bias (bool, optional) – If True, adds a learnable spectral bias computed from the spatial integral, by default False.

Examples

>>> import torch
>>> import torch_harmonics as th
>>> conv = th.SpectralConvS2(
...     in_shape=(128, 256), out_shape=(128, 256),
...     in_channels=16, out_channels=32,
... ).cuda()
>>> x = torch.randn(4, 16, 128, 256, device="cuda")
>>> y = conv(x)
>>> y.shape
torch.Size([4, 32, 128, 256])
Raises:

AssertionError – If in_channels or out_channels is not divisible by num_groups.

Parameters:

Notes

The SHT truncation lmax/mmax is the minimum of the input and output truncations.

forward(x)[source]#

Apply the spectral convolution.

Parameters:

x (torch.Tensor) – Input signal of shape (batch, in_channels, nlat_in, nlon_in).

Returns:

Convolved signal of shape (batch, out_channels, nlat_out, nlon_out).

Return type:

torch.Tensor

On this page