Differentiable and accelerated spherical harmonic and Wigner transforms

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.

31 pages, 7 figures, accepted by Journal of Computational Physics, code available at https://github.com/astro-informatics/s2fft

References in corpus (26)

• HEALPix -- a Framework for High Resolution Discretization, and Fast Analysis of Data Distributed on the Sphere
• LSST: from Science Drivers to Reference Design and Anticipated Data Products
• Efficient Spherical Harmonic Transforms aimed at pseudo-spectral numerical simulations
• A novel sampling theorem on the sphere
• Transformations of asymptotic gravitational-wave data
• Dark Energy Survey Year 3 results: curved-sky weak lensing mass map reconstruction
• Libsharp - spherical harmonic transforms revisited
• Fast directional continuous spherical wavelet transform algorithms
• Weak lensing scattering transform: dark energy and neutrino mass sensitivity
• An Optimal-Dimensionality Sampling Scheme on the Sphere with Fast Spherical Harmonic Transforms
• S2LET: A code to perform fast wavelet analysis on the sphere
• Localisation of directional scale-discretised wavelets on the sphere
• Euclid: Reconstruction of Weak Lensing mass maps for non-Gaussianity studies
• Algorithm for the evaluation of reduced Wigner matrices
• Directional spin wavelets on the sphere
• Sparse image reconstruction on the sphere: implications of a new sampling theorem
• Spherical Bayesian mass-mapping with uncertainties: full sky observations on the celestial sphere
• A novel sampling theorem on the rotation group
• Numerical evolutions of fields on the 2-sphere using a spectral method based on spin-weighted spherical harmonics
• Differentiable Cosmological Simulation with Adjoint Method
• Sparse image reconstruction on the sphere: analysis and synthesis
• Differentiable and accelerated spherical harmonic and Wigner transforms
• Posterior sampling for inverse imaging problems on the sphere in seismology and cosmology
• Learned Interferometric Imaging for the SPIDER Instrument
• Gauss-Legendre Sampling on the Rotation Group
• Sparse image reconstruction on the sphere: a general approach with uncertainty quantification

Cited by in corpus (2)

• Generative models of astrophysical fields with scattering transforms on the sphere
• Differentiable and accelerated spherical harmonic and Wigner transforms