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 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.
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. This is what the directional wavelet filters look like in pixel space.
The Python dependencies for the S2WAV
package are listed in the file
requirements/requirements-core.txt
and will be automatically installed
into the active python environment by pip when running
pip install .
from the root directory of the repository. Unit tests can then be executed to ensure the installation was successful by running
pytest tests/
In the near future one will be able to install S2WAV
directly from
PyPi by pip install s2wav
but this is not yet supported.
Note that to run JAX
on NVIDIA GPUs you will need to follow the
guide outlined by Google.
To import and use S2WAV
is as simple follows:
import s2wav
# Compute wavelet coefficients
f_wav, f_scal = s2wav.analysis(f, L, N)
# Map back to signal on the sphere
f = s2wav.synthesis(f_wav, f_scal, L, N)
however we strongly recommend that the multiresolution argument is set to true, as this will accelerate the transform by a factor of the total number of wavelet scales, which can be around an order of magnitude.
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 (emoji key):
Matt Price π» π π π¨ |
Jason McEwen π π¨ |
Alicja Polanska π» π |
Jessica Whitney π» π |
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"},
eprint = {arXiv:2402.01282},
year = {2024}
}
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.