s2x: Accelerated and Differentiable Spherical Harmonic Transforms

Many fields of science and engineering encounter data on spherical manifolds, such as the celestial sphere in cosmology or the planet earth in climate forecasting. Modelling and analysis of these problems requires the spherical nature of the data to be considered and incorporated in any potential solutions. While codes which can be applied to spherical problems exist, they are severely limited in their usefulness due to two common factors. First, the codes are typically confined to the CPU and therefore lacking the ability to take advantage of powerful modern hardware accelerators, such as GPUs. Second, they do not support automatic differentiation which is critical for physical modelling and AI training. To remedy this, we develop the S2X code suite. Consisting of four software packages implemented in Python using the JAX framework with PyTorch support. Therefore supporting automatic differentiation and enabling scalable GPU deployment on HPC systems through parallelized algorithm design. S2FFT provides spherical harmonic transforms and forms the basis for S2WAV and S2SCAT. These contain spherical wavelet transforms and spherical scattering transforms respectively. S2AI provides scalable and equivariant spherical convolution neural networks.