S2BALL is a JAX package which provides support for automatically
differentiable and greatly accelerated generalized Fourier transforms on the sphere,
rotation group, and ball. Furthermore, we support extremely efficient and differentiable
Wigner-Laguerre wavelets which can be up to 22,000 times faster
than their C counterparts, discussed in related works by Leistedt & McEwen 2012 and Price & McEwen 2021.
The transforms S2BALL provides are optimally fast but come with a substantial memory overhead and cannot be used above a harmonic bandlimit of L ~ 256, at least with current GPU memory limitations. That being said, many applications are more than comfortable at these resolutions, for which these JAX transforms are ideally suited, e.g. geophysical modelling, diffusion tensor imaging, multiscale molecular modelling. For those with machine learning in mind, it should be explicitly noted that these transforms are indeed equivariant to their respective groups.
Read the full documentation here.
Should this code be used in any way, we kindly request that the following article is referenced. A BibTeX entry for this reference may look like:
@article{price:s2ball,
author = "Matthew A. Price and Alicja Polanska and Jessica Whitney and Jason D. McEwen",
title = "Differentiable and accelerated directional wavelet transform on the sphere and ball",
year = "2024",
eprint = "arXiv:2402.01282"
}
S2BALL is released under the MIT license (see LICENSE.txt).
We provide this code under an MIT open-source licence with the hope that it will be of use to a wider community. Copyright 2024 Matthew Price, Jason McEwen and contributors. `S2BALL` is free software made available under the MIT License. For details see the LICENSE file.