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.

Numerical Algorithms for the Computation of Generalized Prolate Spheroidal Functions

Generalized prolate spheroidal functions: algorithms and analysis

Fast Fourier Transforms for Spherical Gauss-Laguerre Basis Functions

Spherical Slepian functions and the polar gap in geodesy

Numerical reconstruction from the Fourier transform on the ball using prolate spheroidal wave functions

Generalized Fourier-Bessel operator and almost-periodic interpolation and approximation

Strictly and non-strictly positive definite functions on spheres

Spherical higher order Fourier analysis over finite fields III: a spherical Gowers inverse theorem

Improved Algorithm for Efficient Computation of Slepian Functions Over Arbitrary Regions on the Sphere

Isotropic Positive Definite Functions on Spheres Generated from Those in Euclidean Spaces

Efficient Methods for the Chebyshev-Type Prolate Spheroidal Wave Functions and Corresponding Eigenvalues

Grid method for computation of generalized spheroidal wave functions based on discrete variable representation.

N ov 2 01 4 Generalized Poincaŕe Sphere

Approximation in Hankel Sobolev Space by Circular Prolate Spheroidal Series

A package on orthogonal polynomials and special functions

Discrete Hankel Prolate Spheroidal Wave Functions: Spectral Analysis and Application

Generalized Poincaré Sphere.

Differentiable and accelerated spherical harmonic and Wigner transforms

Generalized Prolate Spheroidal Wave Functions Associated with Linear Canonical Transform

Numerical optimisation of Dirac eigenvalues

Technologies for supporting high-order geodesic mesh frameworks for computational astrophysics and space sciences