Differentiable and accelerated spherical wavelets

S2WAV is a JAX package for computing wavelet 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), and can be mapped across multiple accelerators.

More specifically, S2WAV provides support for scale-discretised wavelet transforms on the sphere and rotation group (for both real and complex signals), with support for adjoints where needed, and comes with a variety of different optimisations (e.g. precompute or not, multi-resolution algorithms) that one may select depending on available resources and desired angular resolution \(L\).

Wavelet Transform ⚡

S2WAV is an updated implementation of the scale-discretised wavelet transform on the sphere, which builds upon the papers of Leistedt et al 2013 and McEwen et al 2017. This wavelet transform is designed to have excellent localisation and uncorrelation properties, and has been successfully adopted for various applications e.g. scattering transforms on the sphere McEwen et al 2022. The wavelet dictionary is constructed by tiling the harmonic line with infinitely differentiable Cauchy-Schwartz functions, which can straightforwardly be performed in an efficient multiresolution manner, as in the Euclidean case. For example the directional wavelet decomposition of a topographic map of the Earth can be seen below

Contributors ✨

We strongly encourage contributions from any interested developers; a simple example would be adding support for new wavelet filters e.g. spherical needlets Chan et al 2016 or spherical ridgelets McEwen & Price 2020! Thanks goes to these wonderful people (emojikey):

_{Matt Price}
💻 👀 📖 🎨

_{Jason McEwen}
👀 🎨

_{Alicja Polanska}
💻 👀

_{Jessica Whitney}
💻 👀

Attribution 📚

A BibTeX entry for S2WAV is:

@article{price:s2wav,
   AUTHOR = {Author names},
    TITLE = {"TBA"},
   EPRINT = {arXiv:0000.00000},
     YEAR = {2023}
}

License 📝

S2WAV is free software made available under the MIT License. For details see the LICENSE file.