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.