Differentiable and accelerated spherical harmonic and Wigner transforms (2311.14670v2)

Abstract: Many areas of science and engineering encounter data defined on spherical manifolds. Modelling and analysis of spherical data often necessitates spherical harmonic transforms, at high degrees, and increasingly requires efficient computation of gradients for machine learning or other differentiable programming tasks. We develop novel algorithmic structures for accelerated and differentiable computation of generalised Fourier transforms on the sphere $\mathbb{S}^2$ and rotation group $\text{SO}(3)$, i.e. spherical harmonic and Wigner transforms, respectively. We present a recursive algorithm for the calculation of Wigner $d$-functions that is both stable to high harmonic degrees and extremely parallelisable. By tightly coupling this with separable spherical transforms, we obtain algorithms that exhibit an extremely parallelisable structure that is well-suited for the high throughput computing of modern hardware accelerators (e.g. GPUs). We also develop a hybrid automatic and manual differentiation approach so that gradients can be computed efficiently. Our algorithms are implemented within the JAX differentiable programming framework in the S2FFT software code. Numerous samplings of the sphere are supported, including equiangular and HEALPix sampling. Computational errors are at the order of machine precision for spherical samplings that admit a sampling theorem. When benchmarked against alternative C codes we observe up to a 400-fold acceleration. Furthermore, when distributing over multiple GPUs we achieve very close to optimal linear scaling with increasing number of GPUs due to the highly parallelised and balanced nature of our algorithms. Provided access to sufficiently many GPUs our transforms thus exhibit an unprecedented effective linear time complexity.

• The paper develops recursive algorithms for spherical harmonic and Wigner transforms that enable stable, high-performance computations at very high harmonic degrees.
• It employs a hybrid differentiation approach combining manual and automatic methods to reduce memory overhead during gradient computations.
• The implementation in the S2FFT library demonstrates up to 400-fold acceleration on GPUs and TPUs, facilitating scalable spherical data analysis.

Overview of "Differentiable and Accelerated Spherical Harmonic and Wigner Transforms"

The research presented in "Differentiable and Accelerated Spherical Harmonic and Wigner Transforms" by Price and McEwen addresses the computational demands for efficient and differentiable spherical harmonic and Wigner transforms. These transforms are critical for data analysis on spherical manifolds, relevant in a wide array of scientific disciplines such as cosmology, quantum chemistry, and geophysics. This paper tackles the emerging need for transforms that can leverage modern hardware accelerators like GPUs and TPUs, and which can also support differentiable programming platforms that are increasingly critical in the field of machine learning.

Algorithmic Contributions

The primary contribution of the paper is the development of novel algorithms for the computation of Fourier transforms on the sphere ($\mathbb{S}^2$) and the rotation group (SO(3)). The authors introduce a recursive algorithm for Wigner $d$-functions computation, which is key for both spin spherical harmonic and Wigner transforms. Notably, the recursion they develop is stable to very high harmonic degrees and exhibits a highly parallelizable structure, allowing for significant computational acceleration on modern hardware accelerators.

The authors achieve this by constructing a recursion that iterates over the order $m$, independently from degree $\ell$, spin $n$, and angle $\beta$. This independence allows for extreme parallelization, a critical feature for deployment on hardware accelerators. They further mitigate the numerical overflow issues by implementing an on-the-fly normalisation scheme throughout the recursion.

Hybrid Differentiation Approach

Price and McEwen develop a hybrid approach to differentiation that combines both manual differentiation for core recursive components and automatic differentiation for the encompassing operations. This approach effectively mitigates the computational memory overhead typically associated with reverse-mode automatic differentiation while still reaping the software maintainability and complexity reduction benefits AD provides.

Practical Implementation: S2FFT

The research leads to the creation of S2FFT, a Python-based software library leveraging the JAX framework. S2FFT supports varying sampling schemes such as Driscoll & Healy (DH), McEwen & Wiaux (MW), and the HEALPix sampling scheme. S2FFT is designed to transparently deploy over multiple hardware accelerators, and for high-bandwidth settings, reveals close to linear scalability with increasing numbers of GPUs.

Numerical Evaluation and Performance

The performance of the proposed algorithms is rigorously evaluated against existing implementations. For bandlimits up to $L=8192$, the algorithms demonstrate computational errors within machine precision for sampling theorems, confirming the effectiveness of implementing exact quadrature approaches. Moreover, S2FFT achieves speed-ups up to 400-fold compared to alternative C implementations for spherical harmonic and Wigner transforms, showcasing its potential to significantly reduce computational time in extensive applications.

Implications and Future Directions

The results presented in this paper have significant implications for computational science and data analysis on spherical domains. By making these transforms both fast and differentiable, the paper opens up new possibilities for their integration with machine learning frameworks, particularly for gradient-based optimization processes. Potential applications extend to precise global climate modeling, cosmological data analysis, and various physics-based machine learning models.

Furthermore, this work could underpin advances in geometric deep learning, where leveraging spherical data with efficient harmonic analysis plays a crucial role. Future work might explore extending these algorithms for other spherical transformations or adapting them to emerging scientific computing environments and hardware platforms. The development of S2FFT as an open-source project invites further community engagement and extension, fostering further innovation and application in diverse scientific settings.