REVIEW 9 cited by
CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators
Not yet reviewed by Pith; the record is open.
This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.
SPECIMEN: schema-true, not a live event
T0 review · schema-true
One-sentence machine reading of the paper's core claim.
pith:XXXXXXXX · record.json · timestamp
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.
Forward citations
Cited by 9 Pith papers
-
Alleviating prior dependencies for DESI DR1 clustering fits through reparameterization
Jeffreys prior over EFTofLSS coefficients mitigates projection effects in DESI DR1 power spectrum multipole fits, recentering posteriors for late-time expansion parameters.
-
Using Neural Emulators and Hamiltonian Monte Carlo to constrain the Epoch of Reionization's History with the Ly$\alpha$ Forest Power Spectrum
A JAX pipeline that emulates both the Lyα forest power spectrum and its model-dependent covariance with neural networks, then runs Hamiltonian Monte Carlo, recovers reionization-history parameters from mock observations.
-
Differentiable Halo Mass Prediction and the Cosmology-Dependence of Halo Mass Functions
A differentiable U-Net predicts halo mass functions and their cosmology derivatives from initial density fields, matching finite-difference gradients of simulations and emulators to within model scatter.
-
Tuning the cosmic instrument: robust cosmology through combined probes
Low-redshift large-scale structure data pull the dark-energy equation-of-state posterior toward a cosmological constant and yield S8=0.777±0.017.
-
Effort: a fast and differentiable emulator for the Effective Field Theory of the Large Scale Structure of the Universe
A fast, differentiable emulator for EFTofLSS galaxy power spectra, validated against pybird on simulations and BOSS data, enables gradient-based MCMC inference.
-
Reducing nuisance prior sensitivity via non-linear reparameterization, with application to EFT analyses of large-scale structure
A GAM-based reparameterization is shown to reduce the sensitivity of cosmological constraints from EFT-of-LSS analyses to the choice of nuisance parameter priors.
-
Attention-based Neural Network Emulators for Multi-Probe Data Vectors Part III: Modeling The Next Generation Surveys
A transformer-based emulator reproduces CAMB CMB TT, TE, and EE power spectra within cosmic variance errors across a wide Lambda-CDM parameter space, with outlier fractions below 10% for future survey configurations.
-
halox: Dark matter halo properties and large-scale structure calculations using JAX
A new JAX library provides differentiable, GPU-accelerated implementations of NFW halo profiles, the Tinker et al. halo mass function and halo bias, validated against colossus and Astropy.
-
Cosmological Parameter Estimation with Sequential Linear Simulation-based Inference
LSBI approximates simulated likelihoods as linear Gaussians and iteratively refines posteriors, reaching accuracy comparable to neural SBI on CMB power spectrum emulations.
Discussion (0). Continue with ORCID to comment.