REVIEW 9 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 9 Pith papers
-
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.
-
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.