Differentiable and Accelerated Spherical Harmonic Transforms

In JAX and PyTorch

Many areas of science and engineering encounter data defined on the sphere. Modelling and analysis of such data often requires the spherical counterpart to the Fourier transform — the spherical harmonic transform. We provide a brief overview of the spherical harmonic transform and present a new differentiable algorithm tailored towards acceleration on GPUs [1]. This algorithm is implemented in the recently released S2FFT python package, which supports both JAX and PyTorch.

Increasingly …

algorithm analysis and analysis backpropagation data differentiable engineering fourier geometric-deep-learning jax machine learning mathematics modelling overview physics science sphere