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

Understand

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.