CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators

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 $\mathcal{O}(10^3)$. 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.

arXiv comment12 pages, 5 figures. Accepted for publication in The Open Journal of Astrophysics. CosmoPower-JAX is available at https://github.com/dpiras/cosmopower-jax

References in corpus (51)

• Variational Inference: A Review for Statisticians
• Cosmological parameters from CMB and other data: a Monte-Carlo approach
• MCMC using Hamiltonian dynamics
• MultiNest: an efficient and robust Bayesian inference tool for cosmology and particle physics
• The Cosmic Linear Anisotropy Solving System (CLASS) II: Approximation schemes
• Multimodal nested sampling: an efficient and robust alternative to MCMC methods for astronomical data analysis
• The Simons Observatory: Science goals and forecasts
• Dark Energy Survey Year 3 Results: Cosmological Constraints from Galaxy Clustering and Weak Lensing
• Wide-Field InfrarRed Survey Telescope-Astrophysics Focused Telescope Assets WFIRST-AFTA 2015 Report
• Importance Nested Sampling and the MultiNest Algorithm
• Conservative Constraints on Early Cosmology: an illustration of the Monte Python cosmological parameter inference code
• Cobaya: Code for Bayesian Analysis of hierarchical physical models
• Intrinsic alignment-lensing interference as a contaminant of cosmic shear
• Efficient sampling of fast and slow cosmological parameters
• PolyChord: nested sampling for cosmology
• PolyChord: next-generation nested sampling
• Extended Limber Approximation
• An accurate halo model for fitting non-linear cosmological power spectra and baryonic feedback models
• Systematic Errors in Future Weak Lensing Surveys: Requirements and Prospects for Self-Calibration
• Bayesian physical reconstruction of initial conditions from large scale structure surveys
• Core Cosmology Library: Precision Cosmological Predictions for LSST
• Accurate halo-model matter power spectra with dark energy, massive neutrinos and modified gravitational forces
• Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro
• Simultaneous measurement of cosmology and intrinsic alignments using joint cosmic shear and galaxy number density correlations
• COSMOPOWER: emulating cosmological power spectra for accelerated Bayesian inference from next-generation surveys
• Weak lensing shear calibration with simulations of the HSC survey
• Gravitational Lensing of Distant Field Galaxies by Rich Clusters: I. -- Faint Galaxy Redshift Distributions
• Bayesian non-linear large scale structure inference of the Sloan Digital Sky Survey data release 7
• Unmasking the Masked Universe: the 2M++ catalogue through Bayesian eyes
• Fast optimal CMB power spectrum estimation with Hamiltonian sampling
• Accelerating astronomical and cosmological inference with Preconditioned Monte Carlo
• Efficient Cosmological Parameter Estimation with Hamiltonian Monte Carlo
• JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library
• Euclid: The importance of galaxy clustering and weak lensing cross-correlations within the photometric Euclid survey
• Cosmology with the Wide-Field Infrared Survey Telescope -- Synergies with the Rubin Observatory Legacy Survey of Space and Time
• The Cosmic Graph: Optimal Information Extraction from Large-Scale Structure using Catalogues
• Accelerating Large-Scale-Structure data analyses by emulating Boltzmann solvers and Lagrangian Perturbation Theory
• CosmicNet II: Emulating extended cosmologies with efficient and accurate neural networks
• COMET: Clustering Observables Modelled by Emulated perturbation Theory
• Parameter Inference for Weak Lensing using Gaussian Processes and MOPED
• KiDS-1000 Cosmology: Constraints from density split statistics
• Hamiltonian Monte Carlo reconstruction from peculiar velocities
• KiDS-1000 Cosmology: machine learning -- accelerated constraints on Interacting Dark Energy with COSMOPOWER
• Almanac: Weak Lensing power spectra and map inference on the masked sphere
• On Scale-Dependent Cosmic Shear Systematic Effects
• Field-level inference of cosmic shear with intrinsic alignments and baryons
• Microcanonical Hamiltonian Monte Carlo
• Fast emulation of two-point angular statistics for photometric galaxy surveys
• Propagating spatially-varying multiplicative shear bias to cosmological parameter estimation for stage-IV weak-lensing surveys
• High-accuracy emulators for observables in $\Lambda$CDM, $N_{\mathrm{eff}}$, $\Sigma m_{\nu }$, and $w$ cosmologies
• Sampling from high-dimensional, multimodal distributions using automatically tuned, tempered Hamiltonian Monte Carlo

Cited by in corpus (17)

• The CosmoVerse White Paper: Addressing observational tensions in cosmology with systematics and fundamental physics
• SPT-3G D1: CMB temperature and polarization power spectra and cosmology from 2019 and 2020 observations of the SPT-3G Main field
• candl: Cosmic Microwave Background Analysis with a Differentiable Likelihood
• DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology
• LINX: A Fast, Differentiable, and Extensible Big Bang Nucleosynthesis Package
• A representation learning approach to probe for dynamical dark energy in matter power spectra
• On the degeneracies between baryons, massive neutrinos and f(R) gravity in Stage IV cosmic shear analyses
• Cosmic Cartography: Bayesian reconstruction of the galaxy density informed by large-scale structure
• Reducing nuisance prior sensitivity via non-linear reparameterization, with application to EFT analyses of large-scale structure
• Cosmic Cartography II: completing galaxy catalogs for gravitational-wave cosmology
• ${\text{SwiftC}}_{\ell }$: fast differentiable angular power spectra beyond Limber
• jaxspec : a fast and robust Python library for X-ray spectral fitting
• Differentiable Modeling of Planet and Substellar Atmosphere: High-Resolution Emission, Transmission, and Reflection Spectroscopy with ExoJAX2
• Constraining the primordial power spectrum using a differentiable likelihood
• Anchors no more: Using peculiar velocities to constrain $H_0$ and the primordial Universe without calibrators
• ABCMB: A Python+JAX Package for the Cosmic Microwave Background Power Spectrum
• Conditional variational autoencoders for cosmological model discrimination and anomaly detection in cosmic microwave background power spectra