Differentiable and accelerated wavelet transforms on the sphere and ball

TL;DR

This work designs new highly distributable and automatically differentiable directional wavelet transforms on the sphere and ball that dramatically accelerate existing spherical wavelet transforms, and unlocks many data-driven analysis techniques previously not possible for these spaces.

Abstract

Directional wavelet dictionaries are hierarchical representations which efficiently capture and segment information across scale, location and orientation. Such representations demonstrate a particular affinity to physical signals, which often exhibit highly anisotropic, localised multiscale structure. Many physically important signals are observed over spherical domains, such as the celestial sky in cosmology. Leveraging recent advances in computational harmonic analysis, we design new highly distributable and automatically differentiable directional wavelet transforms on the $2$-dimensional sphere $\mathbb{S}^2$ and $3$-dimensional ball $\mathbb{B}^3 = \mathbb{R}^+ \times \mathbb{S}^2$ (the space formed by augmenting the sphere with the radial half-line). We observe up to a $300$-fold and $21800$-fold acceleration for signals on the sphere and ball, respectively, compared to existing software, whilst maintaining 64-bit machine precision. Not only do these algorithms dramatically accelerate existing spherical wavelet transforms, the gradient information afforded by automatic differentiation unlocks many data-driven analysis techniques previously not possible for these spaces. We publicly release both S2WAV and S2BALL, open-sourced JAX libraries for our transforms that are automatically differentiable and readily deployable both on and over clusters of hardware accelerators (e.g. GPUs & TPUs).

Figures (5)

• Figure 1: Directional ($L=256, N=5$) scale-discretised wavelet filters on the sphere, for wavelet scales $j\in\lbrace1,\dots,5\rbrace$ from left to right. As the wavelet scale increases the filter becomes more highly localised, both in pixel and frequency space mcewen:s2let_localisation. Notice the clearly elongated structure which gives rise to the desired directional sensitivity mcewen:s2let_spin. With an azimuthal bandlimit of $N=5$ this filter bank is augmented by $9$ rotated filters, which are not shown here.
• Figure 2: Directional ($N=5$) scale-discretised wavelet filters on the ball ($P=L=256$), for angular wavelet scales $j\in[3,7]$ from left to right, and radial wavelet scale $j^\prime = 3$ for radial nodes $r\in[1,4]$. As the angular wavelet scale increases the filter becomes more highly localised, both in pixel and frequency space mcewen:s2let_localisation. Equally, as the radial scale increases the localisation along the radial half-line increases. For a given radial scale (as shown here) the energy of a given filter $\Psi(r)$ decays exponentially with $r$ as expected. Notice the clearly elongated structure which gives rise to the desired directional sensitivity.
• Figure 3: Surface visualisation of a subset of our tiling of Fourier-Laguerre space with infinitely differentiable Cauchy-Schwartz functions (see Appendix \ref{['sec:tiling_ball']}). Here we present the axisymmetric $(m=0)$ components of our wavelet filters for $j=j^\prime$. Notice that if one were to consider all $j<j^\prime$ and vice versa this filtering scheme would span the entire domain.
• Figure : Directional wavelet transform on $\mathbb{S}^2$
• Figure : Directional wavelet transform on $\mathbb{B}^3$