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

TL;DR

DISCO-DJ's Einstein-Boltzmann solver offers a fully differentiable, modular, and accurate tool for cosmological inference, enabling automatic Jacobian computation and easier extension to new physics.

Contribution

We introduce a differentiable Einstein-Boltzmann solver in JAX that matches existing solvers' accuracy and simplifies extensions to new cosmological models.

Findings

Solver agrees at the per-mille level with CAMB and CLASS

Enables automatic Jacobian computation for all outputs

Facilitates Fisher forecasting for cosmological surveys

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.