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:

Legendre synthesis — a distributed matrix multiply: degrees stay split across the polar group and each rank synthesizes from the degrees it owns, producing a partial sum over all latitudes.
Reduce-scatter over the polar group along nlat — completes the synthesis sum and leaves the latitudes partitioned in one collective.
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.

Keeping l distributed is what makes the precomputed Legendre polynomials scale: they are partitioned as (mmax_local, nlat, lmax_local), so the tensor is split across the full process grid rather than replicated over the polar group. The cost is that the synthesis sum is now accumulated across ranks, so results are not bitwise identical to the serial transform.

If N < 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.