torch_harmonics.distributed.DistributedInverseRealSHT#

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

Bases: Module

Distributed version of the inverse (real-valued) SHT. Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.

Distribution scheme. The input tensor has shape (B, C, lmax_local, mmax_local) where spectral degrees and orders are split across the polar and azimuth process groups. All leading dimensions are flattened into N = B * C for redistribution. The forward pass proceeds as follows:

Polar transpose (N ↔ lmax) — trade channel slices for the full degree axis, making l local for the Legendre synthesis.
Legendre synthesis — local matrix multiply with the associated Legendre polynomials, producing latitude points.
Polar transpose (nlat ↔ N) — redistribute so that latitudes are split across polar ranks and channels are local.
Azimuth transpose (N ↔ mmax) — make spectral orders m fully local for the inverse FFT.
Inverse real FFT along the (now local) m / longitude dimension.
Azimuth transpose (nlon ↔ N) — redistribute so that longitudes are split across azimuth ranks.

The output has shape (B, C, nlat_local, nlon_local) with the spatial grid partitioned in the same way as the input spectral modes.

If N < max(polar_group_size, azimuth_group_size), the leading axis is zero-padded before the transposes and the padding is removed afterwards; since the transform is linear this is exact.