COSMOPOWER-JAX: HIGH-DIMENSIONAL BAYESIAN INFERENCE WITH DIFFERENTIABLE COSMOLOGICAL EMULATORS

5Citations
Citations of this article
7Readers
Mendeley users who have this article in their library.

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.

Cite

CITATION STYLE

APA

Piras, D., & Mancini, A. S. (2023). COSMOPOWER-JAX: HIGH-DIMENSIONAL BAYESIAN INFERENCE WITH DIFFERENTIABLE COSMOLOGICAL EMULATORS. Open Journal of Astrophysics, 6. https://doi.org/10.21105/astro.2305.06347

Register to see more suggestions

Mendeley helps you to discover research relevant for your work.

Already have an account?

Save time finding and organizing research with Mendeley

Sign up for free