DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology

Abstract

We present the Einstein-Boltzmann module of the DISCO-DJ (DIfferentiable Simulations for COsmology - Done with JAX) software package. This module implements a fully differentiable solver for the linearised cosmological Einstein-Boltzmann equations in the JAX framework, and allows computing Jacobian matrices of all solver output with respect to all input parameters using automatic differentiation. This implies that along with the solution for a given set of parameters, the tangent hyperplane in parameter space is known as well, which is a key ingredient for cosmological inference and forecasting problems as well as for many other applications. We discuss our implementation and demonstrate that our solver agrees at the per-mille level with the existing non-differentiable solvers CAMB and CLASS, including massive neutrinos and a dark energy fluid with parameterised equation of state. We illustrate the dependence of various summary statistics in large-scale structure cosmology on model parameters using the differentiable solver, and finally demonstrate how it can be easily used for Fisher forecasting, with a forecast for Euclid as an example. Since the implementation is significantly shorter and more modular than existing solvers, we believe it will be more straightforward to extend our solver to include additional physics, such as additional dark energy and dark matter models, modified gravity, or other non-standard physics in the future.

32 pages, 8 figures, published in JCAP: improved Fisher forecast and minor fixes improving agreement with CLASS/CAMB; DISCO-EB code is available from https://github.com/ohahn/DISCO-EB

References in corpus (15)

• The Cosmic Linear Anisotropy Solving System (CLASS) I: Overview
• A search for ultra-light axions using precision cosmological data
• Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro
• Large-scale dark matter simulations
• On Neural Differential Equations
• Dark energy with non-adiabatic sound speed: initial conditions and detectability
• MGCAMB with massive neutrinos and dynamical dark energy
• JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library
• CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators
• Capse.jl: efficient and auto-differentiable CMB power spectra emulation
• Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows
• pmwd: A Differentiable Cosmological Particle-Mesh $N$-body Library
• Fisher Matrix Stability
• A Differentiable Perturbation-based Weak Lensing Shear Estimator
• Fast emulation of two-point angular statistics for photometric galaxy surveys

Cited by in corpus (7)

• The CosmoVerse White Paper: Addressing observational tensions in cosmology with systematics and fundamental physics
• syren-new: Precise formulae for the linear and nonlinear matter power spectra with massive neutrinos and dynamical dark energy
• LINX: A Fast, Differentiable, and Extensible Big Bang Nucleosynthesis Package
• BullFrog: Multi-step perturbation theory as a time integrator for cosmological simulations
• Differentiable Cosmological Hydrodynamics for Field-Level Inference and High Dimensional Parameter Constraints
• Diagnosing Systematic Effects Using the Inferred Initial Power Spectrum
• ABCMB: A Python+JAX Package for the Cosmic Microwave Background Power Spectrum