REVIEW 13 cited by
JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library
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
read the original abstract
We present jax-cosmo, a library for automatically differentiable cosmological theory calculations. It uses the JAX library, which has created a new coding ecosystem, especially in probabilistic programming. As well as batch acceleration, just-in-time compilation, and automatic optimization of code for different hardware modalities (CPU, GPU, TPU), JAX exposes an automatic differentiation (autodiff) mechanism. Thanks to autodiff, jax-cosmo gives access to the derivatives of cosmological likelihoods with respect to any of their parameters, and thus enables a range of powerful Bayesian inference algorithms, otherwise impractical in cosmology, such as Hamiltonian Monte Carlo and Variational Inference. In its initial release, jax-cosmo implements background evolution, linear and non-linear power spectra (using halofit or the Eisenstein and Hu transfer function), as well as angular power spectra with the Limber approximation for galaxy and weak lensing probes, all differentiable with respect to the cosmological parameters and their other inputs. We illustrate how autodiff can be a game-changer for common tasks involving Fisher matrix computations, or full posterior inference with gradient-based techniques. In particular, we show how Fisher matrices are now fast, exact, no longer require any fine tuning, and are themselves differentiable. Finally, using a Dark Energy Survey Year 1 3x2pt analysis as a benchmark, we demonstrate how jax-cosmo can be combined with Probabilistic Programming Languages to perform posterior inference with state-of-the-art algorithms including a No U-Turn Sampler, Automatic Differentiation Variational Inference,and Neural Transport HMC. We further demonstrate that Normalizing Flows using Neural Transport are a promising methodology for model validation in the early stages of analysis.
Forward citations
Cited by 13 Pith papers
-
Practical and Scalable Hamiltonian Monte Carlo Without the Metropolis Test
For Gaussian targets the per-dimension variance of the integrator energy error provides a sharp bound on the asymptotic bias, and the resulting automatic step-size selection makes unadjusted HMC and Langevin samplers ...
-
End-to-end differentiable retrieval of molecular spectra using hydrodynamics, chemistry, and radiative transfer
An end-to-end differentiable JAX pipeline couples 1D hydrodynamics, time-dependent chemistry, and radiative transfer, and recovers shock and rate parameters from synthetic HCO+ spectra.
-
Per Astronomix ad Astra: High-Order Differentiable (Magneto)hydrodynamics with Energy-Conserving Self-Gravity
A new differentiable MHD+self-gravity simulator combines a fifth-order finite-difference WENO scheme with a fourth-order semi-discretely energy-conserving self-gravity module.
-
Diffhalos: A Generative Model of Cosmological Lightcones of Dark Matter Halos
Diffhalos generates statistically accurate Monte-Carlo and quasi-Monte-Carlo lightcones of halos, subhalos and Diffmah mass-assembly histories, enabling autodiff gradients of the mass functions.
-
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.
-
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.
-
Diagnosing Systematic Effects Using the Inferred Initial Power Spectrum
The inferred initial matter power spectrum from the SELFI algorithm reveals misspecified galaxy bias, selection, mask, redshift, and gravity models, exposing a >2σ cosmological bias before parameter 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.
-
Savage-Dickey density ratio estimation with normalizing flows for Bayesian model comparison
A normalizing flow estimates the normalized marginal posterior in the Savage-Dickey density ratio, enabling Bayes factors for nested models with many extra parameters.
-
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.
-
Fisher Score Matching for Simulation-Based Forecasting and Inference
A neural network trained to predict latent parameter scores can approximate the Fisher score and support Fisher forecasting and gradient-based Bayesian inference from simulations.
-
Unexplored Opportunities for Automatic Differentiation in Astrophysics
A position paper proposing automatic differentiation for nine astrophysical domains under the GRASP framework, with no working code or benchmark results.
Discussion (0). Continue with ORCID to comment.