REVIEW 8 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
JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library
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 8 Pith papers
-
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.
-
Alleviating prior dependencies for DESI DR1 clustering fits through reparameterization
A hybrid Jeffreys+baseline prior removes multi-σ prior-volume projection in DESI DR1 full-shape fits of H0, w0, and wa, yielding late-time expansion constraints consistent with HOD-informed Bayesian and frequentist analyses.
-
HIcosmo: a differentiable JAX-based framework for cosmology inference
HIcosmo is a new JAX-based differentiable framework for background cosmology inference that matches Cobaya results while delivering 8.7x CPU and up to 20x GPU speedups.
-
Systematic LLM Translation of Legacy Scientific Code to Differentiable Frameworks: Application to a Land Surface Model
An LLM-based pipeline systematically converts the CLM-ml-v2 Fortran land surface model to JAX, delivering full differentiability, 8x fewer optimization steps, and 24x wall-clock speedup at N=2048 while releasing the c...
-
A unified harmonic framework for dark siren cosmology
The GW-galaxy cross-correlation method, unified with spectral sirens in a harmonic framework, can measure H0 to 1% and Omega_m to 5% precision with 2 years of data from next-generation detectors like Einstein Telescop...
-
SPT-3G D1: CMB temperature and polarization power spectra and cosmology from 2019 and 2020 observations of the SPT-3G Main field
SPT-3G delivers the most precise CMB EE and TE spectra at high multipoles to date, giving LCDM parameters with H0 = 66.66 ± 0.60 km/s/Mpc from ground-based data alone and reaching Planck-level constraints when combine...
-
Symbolic Emulators for Cosmology: Accelerating Cosmological Analyses Without Sacrificing Precision
Symbolic emulators approximate key Lambda CDM functions to 0.001-0.05% accuracy across relevant redshifts and Omega_m values, enabling faster 3x2pt inference with consistent results.
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.