Differentiable and accelerated spherical harmonic and Wigner transforms

Differentiable and accelerated spherical harmonic and Wigner transforms Matthew A. Price, Jason D. McEwen *Journal of Computational Physics (2024)* * This work introduces novel algorithmic structures for the **accelerated and differentiable computation** of generalized Fourier transforms on the sphere ($S^2$) and the rotation group ($SO(3)$), specifically spherical harmonic and Wigner transforms. * A key component is a **recursive algorithm for Wigner d-functions** designed to be stable to high harmonic degrees and extremely parallelizable, making the algorithms well-suited for high throughput computing on modern hardware accelerators such as GPUs. * The transforms support efficient computation of gradients, which is critical for machine learning and other differentiable programming tasks, achieved through a **hybrid automatic and manual differentiation approach** to avoid the memory overhead associated with full automatic differentiation. * Implemented in the open-source **S2FFT** software code (within the JAX differentiable programming framework), the algorithms support various sampling schemes, including equiangular samplings that admit exact spherical harmonic transforms. * Benchmarking results demonstrate **up to a 400-fold acceleration** compared to alternative C codes, and the transforms exhibit **very close to optimal linear scaling** when distributed over multiple GPUs, yielding an unprecedented effective linear time complexity (O(L)) given sufficient computational resources.