JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library - AstroParticule et Cosmologie Accéder directement au contenu
Article Dans Une Revue Open J.Astrophys. Année : 2023

JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library

Joe Zuntz
  • Fonction : Auteur
Santiago Casas
  • Fonction : Auteur
Minas Karamanis
  • Fonction : Auteur
David Kirkby
  • Fonction : Auteur
Yin Li
  • Fonction : Auteur
Austin Peel
  • Fonction : Auteur

Résumé

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.

Dates et versions

hal-04007879 , version 1 (28-02-2023)

Identifiants

Citer

Jean-Eric Campagne, François Lanusse, Joe Zuntz, Alexandre Boucaud, Santiago Casas, et al.. JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library. Open J.Astrophys., 2023, 6, pp.1-15. ⟨10.21105/astro.2302.05163⟩. ⟨hal-04007879⟩
49 Consultations
0 Téléchargements

Altmetric

Partager

Gmail Facebook X LinkedIn More