-
Differentiable and accelerated wavelet transforms on the sphere and ball
Authors:
Matthew A. Price,
Alicja Polanska,
Jessica Whitney,
Jason D. McEwen
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…
▽ More
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).
△ Less
Submitted 14 March, 2024; v1 submitted 2 February, 2024;
originally announced February 2024.
-
Differentiable and accelerated spherical harmonic and Wigner transforms
Authors:
Matthew A. Price,
Jason D. McEwen
Abstract:
Many areas of science and engineering encounter data defined on spherical manifolds. Modelling and analysis of spherical data often necessitates spherical harmonic transforms, at high degrees, and increasingly requires efficient computation of gradients for machine learning or other differentiable programming tasks. We develop novel algorithmic structures for accelerated and differentiable computa…
▽ More
Many areas of science and engineering encounter data defined on spherical manifolds. Modelling and analysis of spherical data often necessitates spherical harmonic transforms, at high degrees, and increasingly requires efficient computation of gradients for machine learning or other differentiable programming tasks. We develop novel algorithmic structures for accelerated and differentiable computation of generalised Fourier transforms on the sphere $\mathbb{S}^2$ and rotation group $\text{SO}(3)$, i.e. spherical harmonic and Wigner transforms, respectively. We present a recursive algorithm for the calculation of Wigner $d$-functions that is both stable to high harmonic degrees and extremely parallelisable. By tightly coupling this with separable spherical transforms, we obtain algorithms that exhibit an extremely parallelisable structure that is well-suited for the high throughput computing of modern hardware accelerators (e.g. GPUs). We also develop a hybrid automatic and manual differentiation approach so that gradients can be computed efficiently. Our algorithms are implemented within the JAX differentiable programming framework in the S2FFT software code. Numerous samplings of the sphere are supported, including equiangular and HEALPix sampling. Computational errors are at the order of machine precision for spherical samplings that admit a sampling theorem. When benchmarked against alternative C codes we observe up to a 400-fold acceleration. Furthermore, when distributing over multiple GPUs we achieve very close to optimal linear scaling with increasing number of GPUs due to the highly parallelised and balanced nature of our algorithms. Provided access to sufficiently many GPUs our transforms thus exhibit an unprecedented effective linear time complexity.
△ Less
Submitted 20 May, 2024; v1 submitted 24 November, 2023;
originally announced November 2023.
-
Bayesian model comparison for simulation-based inference
Authors:
A. Spurio Mancini,
M. M. Docherty,
M. A. Price,
J. D. McEwen
Abstract:
Comparison of appropriate models to describe observational data is a fundamental task of science. The Bayesian model evidence, or marginal likelihood, is a computationally challenging, yet crucial, quantity to estimate to perform Bayesian model comparison. We introduce a methodology to compute the Bayesian model evidence in simulation-based inference (SBI) scenarios (also often called likelihood-f…
▽ More
Comparison of appropriate models to describe observational data is a fundamental task of science. The Bayesian model evidence, or marginal likelihood, is a computationally challenging, yet crucial, quantity to estimate to perform Bayesian model comparison. We introduce a methodology to compute the Bayesian model evidence in simulation-based inference (SBI) scenarios (also often called likelihood-free inference). In particular, we leverage the recently proposed learnt harmonic mean estimator and exploit the fact that it is decoupled from the method used to generate posterior samples, i.e. it requires posterior samples only, which may be generated by any approach. This flexibility, which is lacking in many alternative methods for computing the model evidence, allows us to develop SBI model comparison techniques for the three main neural density estimation approaches, including neural posterior estimation (NPE), neural likelihood estimation (NLE), and neural ratio estimation (NRE). We demonstrate and validate our SBI evidence calculation techniques on a range of inference problems, including a gravitational wave example. Moreover, we further validate the accuracy of the learnt harmonic mean estimator, implemented in the HARMONIC software, in likelihood-based settings. These results highlight the potential of HARMONIC as a sampler-agnostic method to estimate the model evidence in both likelihood-based and simulation-based scenarios.
△ Less
Submitted 8 November, 2023; v1 submitted 8 July, 2022;
originally announced July 2022.
