S2WAV: Differentiable and accelerated spherical wavelets

Web Github

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.

code

Publications

Differentiable and accelerated wavelet transforms on the sphere and ball

M. A. Price, A. Polanska, J. Whitney, J. D. McEwen

Preprint PDF Code Code

S2WAV: Differentiable and accelerated spherical wavelets

Publications

Talks