S2FFT: Differentiable and accelerated spherical transforms

Web Github

S2FFT is a JAX package for computing Fourier transforms on the sphere and rotation group. It leverages autodiff to provide differentiable transforms, which are also deployable on modern hardware accelerators (e.g. GPUs and TPUs).

More specifically, S2FFT provides support for spin spherical harmonic and Wigner transforms (for both real and complex signals), with support for adjoint transformations where needed, and comes with different optimisations (precompute or not) that one may select depending on available resources and desired angular resolution.

code

Publications

Differentiable and accelerated spherical harmonic and Wigner transforms

M. A. Price, J. D. McEwen

Preprint PDF Code DOI

Physical machine learning for astrophysics: differentiable spherical harmonics; harmonic Bayesian evidence; spherical scattering networks

Nov 2023 Paris

PDF Code

Geometric deep learning on the sphere for the physical sciences

Jul 2023 London

PDF Code Code

S2FFT: Differentiable and accelerated spherical transforms

Publications

Talks