Differentiable and accelerated spherical harmonic and Wigner transforms

Differentiable and accelerated spherical harmonic and Wigner transformsMatthew 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.

There are no comments yet.

Please log in to write the first comment.