torch_harmonics.RealVectorSHT#

class torch_harmonics.RealVectorSHT(
nlat,
nlon,
lmax=None,
mmax=None,
grid='equiangular',
norm='ortho',
csphase=True,
)[source]#

Bases: Module

Defines a module for computing the forward (real) vector SHT. Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. The SHT is applied to the last three dimensions of the input.

Decomposes a tangential vector field \(\mathbf{v} = v_\theta\,\hat{e}_\theta + v_\lambda\,\hat{e}_\lambda\) into spheroidal and toroidal spectral coefficients \(\hat{s}_l^m\) and \(\hat{t}_l^m\) using the derivatives of the associated Legendre polynomials.

See also

Spherical harmonic transforms

User guide with the full mathematical derivation of the vector SHT formulas, normalization conventions, and worked examples.

Parameters:

  • nlat (int) – Number of latitude points

  • nlon (int) – Number of longitude points

  • lmax (int) – Maximum spherical harmonic degree

  • mmax (int) – Maximum spherical harmonic order

  • grid (str) – Grid type ("equiangular", "legendre-gauss", "lobatto", "equiangular-trapezoidal"), by default "equiangular"

  • norm (str) – Normalization convention ("ortho", "schmidt", "unnorm"), by default "ortho".

  • csphase (bool) – Whether to include the Condon–Shortley phase factor \((-1)^m\), by default True.

Examples

>>> import torch
>>> import torch_harmonics as th
>>> nlat, nlon = 128, 256
>>> vsht = th.RealVectorSHT(nlat, nlon).cuda()
>>> vector_field = torch.randn(1, 2, nlat, nlon, device="cuda")
>>> coeffs = vsht(vector_field)   # shape (1, 2, lmax, mmax), complex
>>> coeffs.shape
torch.Size([1, 2, 128, 129])

Note

This module uses cuFFT (via torch.fft.rfft()) to compute the longitudinal Fourier transform efficiently. When running in float16 or bfloat16 precision, cuFFT requires the transformed dimension (nlon) to be a power of two. If your grid does not satisfy this constraint and the module is called inside a torch.autocast context, guard it with torch.autocast(device_type="cuda", enabled=False):

with torch.autocast(device_type="cuda", dtype=torch.float16):
    # ... other half-precision work ...
    with torch.autocast(device_type="cuda", enabled=False):
        coeffs = vsht(vector_field.float())

References

[2], [3]

forward(x)[source]#

Compute the forward (real) vector spherical harmonic transform.

Parameters:

x (torch.Tensor) – Real-valued tangential vector field of shape (..., 2, nlat, nlon), where the size-2 dimension holds the two tangential (colatitude, longitude) components.

Returns:

Complex vector harmonic coefficients of shape (..., 2, lmax, mmax), where the size-2 dimension holds the spheroidal and toroidal components.

Return type:

torch.Tensor

On this page