torch_harmonics.distributed.DistributedRealSHT#

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

Bases: Module

Distributed version of the forward (real-valued) SHT. Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes. The SHT is applied to the last two dimensions of the input.

Distribution scheme. The input tensor has shape (B, C, nlat_local, nlon_local) where latitudes and longitudes are split across the polar and azimuth process groups respectively. All leading dimensions are flattened into a single axis N = B * C which is used as the redistribution currency during the all-to-all transposes. The forward pass proceeds as follows:

Azimuth transpose (nlon ↔ N) — each rank trades its local longitude chunk for a slice of the channel axis, making nlon fully local so the real FFT can be applied.
Real FFT along the (now local) longitude dimension.
Azimuth transpose (N ↔ mmax) — redistribute so that spectral orders m are split across azimuth ranks and channels are local again.
Polar transpose (N ↔ nlat) — trade channel slices for the full latitude axis, making nlat local for the Legendre contraction.
Legendre contraction — local matrix multiply with the quadrature weights, producing spectral degrees l.
Polar transpose (l ↔ N) — redistribute so that degrees l are split across polar ranks.

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

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.