Differentiable and accelerated wavelet transform on the sphere

S2WAV is a python package for computing wavelet transforms on the sphere and rotation group, both in JAX and PyTorch. 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. S2WAV is a sister package of S2FFT, both of which are part of the SAX project, which aims to provide comprehensive support for differentiable transforms on the sphere and rotation group.

As of version 1.0.0 S2WAV also provides partial frontend support for PyTorch. In future this will be expanded to full support. Also note that this release also provides JAX support for existing C spherical harmonic libraries, specifically SSHT. This works be wrapping python bindings with custom JAX frontends. Note that currently this C to JAX interoperability is limited to CPU.

Documentation

Read the full documentation here.

Attribution

A BibTeX entry for s2wav is:

@article{price:s2wav,
    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"
}

we also request that you cite the following paper

@article{price:s2fft,
    author      = "Matthew A. Price and Jason D. McEwen",
    title       = "Differentiable and accelerated spherical harmonic and Wigner transforms",
    journal     = "Journal of Computational Physics, submitted",
    year        = "2023",
    eprint      = "arXiv:2311.14670"
}

in which the core underlying algorithms for the spherical harmonic and Wigner transforms are developed.

License

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, Jessica Whtiney, Alicja Polanska, Jason McEwen and contributors.

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