Piras, Davide;
Spurio Mancini, Alessio;
(2023)
CosmoPower-JAX: high-dimensional Bayesian inference
with differentiable cosmological emulators.
The Open Journal of Astrophysics
, 6
10.21105/astro.2305.06347.
Preview |
Text
2305.06347.pdf - Published Version Download (2MB) | Preview |
Abstract
We present CosmoPower-JAX, a JAX-based implementation of the CosmoPower framework, which accelerates cosmological inference by building neural emulators of cosmological power spectra. We show how, using the automatic differentiation, batch evaluation and just-in-time compilation features of JAX, and running the inference pipeline on graphics processing units (GPUs), parameter estimation can be accelerated by orders of magnitude with advanced gradient-based sampling techniques. These can be used to efficiently explore high-dimensional parameter spaces, such as those needed for the analysis of next-generation cosmological surveys. We showcase the accuracy and computational efficiency of CosmoPower-JAX on two simulated Stage IV configurations. We first consider a single survey performing a cosmic shear analysis totalling 37 model parameters. We validate the contours derived with CosmoPower-JAX and a Hamiltonian Monte Carlo sampler against those derived with a nested sampler and without emulators, obtaining a speed-up factor of O(103 ). We then consider a combination of three Stage IV surveys, each performing a joint cosmic shear and galaxy clustering (3x2pt) analysis, for a total of 157 model parameters. Even with such a high-dimensional parameter space, CosmoPower-JAX provides converged posterior contours in 3 days, as opposed to the estimated 6 years required by standard methods. CosmoPower-JAX is fully written in Python, and we make it publicly available to help the cosmological community meet the accuracy requirements set by next-generation surveys.
| Type: | Article |
|---|---|
| Title: | CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators |
| Open access status: | An open access version is available from UCL Discovery |
| DOI: | 10.21105/astro.2305.06347 |
| Publisher version: | https://doi.org/10.21105/astro.2305.06347 |
| Language: | English |
| Additional information: | This is an Open Access article published under a Creative Commons Attribution 4.0 International (CC BY 4.0) Licence (https://creativecommons.org/licenses/by/4.0/). |
| Keywords: | Bayesian inference, machine learning, cosmology, JAX, data analysis, power spectra |
| UCL classification: | UCL UCL > Provost and Vice Provost Offices > UCL BEAMS UCL > Provost and Vice Provost Offices > UCL BEAMS > Faculty of Maths and Physical Sciences UCL > Provost and Vice Provost Offices > UCL BEAMS > Faculty of Maths and Physical Sciences > Dept of Space and Climate Physics |
| URI: | https://discovery-pp.ucl.ac.uk/id/eprint/10173280 |
Archive Staff Only
 |
View Item |
