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.

arXiv comment32 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 (51)

• Efficient Computation of CMB anisotropies in closed FRW models
• Exploring the Expansion History of the Universe
• Accelerating Universes with Scaling Dark Matter
• The Cosmic Linear Anisotropy Solving System (CLASS) II: Approximation schemes
• A Line of Sight Approach to Cosmic Microwave Background Anisotropies
• Large-Scale Structure of the Universe and Cosmological Perturbation Theory
• Cosmological Perturbation Theory in the Synchronous and Conformal Newtonian Gauges
• Cosmological Perturbation Theory in the Synchronous vs. Conformal Newtonian Gauge
• Cosmology and fundamental physics with the Euclid satellite
• Large-Scale Galaxy Bias
• Probing Dark Energy with Baryonic Acoustic Oscillations from Future Large Galaxy Redshift Surveys
• Reconstructing the linear power spectrum of cosmological mass fluctuations
• A New Calculation of the Recombination Epoch
• The Microwave Anisotropy Probe (MAP) Mission
• Planck Early Results: The Planck mission
• CMB power spectrum parameter degeneracies in the era of precision cosmology
• The Cosmic Linear Anisotropy Solving System (CLASS) I: Overview
• A search for ultra-light axions using precision cosmological data
• Measuring cosmological parameters with galaxy surveys
• The Cosmic Linear Anisotropy Solving System (CLASS) IV: Efficient implementation of non-cold relics
• Early Dark Energy Does Not Restore Cosmological Concordance
• Uncorrelated Modes of the Nonlinear Power Spectrum
• Core Cosmology Library: Precision Cosmological Predictions for LSST
• HyRec: A fast and highly accurate primordial hydrogen and helium recombination code
• Composable Effects for Flexible and Accelerated Probabilistic Programming in NumPyro
• Effective Field Theory of Cosmic Acceleration: an implementation in CAMB
• The 2dF Galaxy Redshift Survey: Spherical Harmonics analysis of fluctuations in the final catalogue
• The General Primordial Cosmic Perturbation
• COSMOPOWER: emulating cosmological power spectra for accelerated Bayesian inference from next-generation surveys
• hi_class: Horndeski in the Cosmic Linear Anisotropy Solving System
• Large-scale dark matter simulations
• Fourier Analysis of Redshift Space Distortions and the Determination of Omega
• Towards a complete treatment of the cosmological recombination problem
• On Neural Differential Equations
• Dark energy with non-adiabatic sound speed: initial conditions and detectability
• MGCAMB with massive neutrinos and dynamical dark energy
• HYREC-2: a highly accurate sub-millisecond recombination code
• JAX-COSMO: An End-to-End Differentiable and GPU Accelerated Cosmology Library
• An $n$-th order Lagrangian Forward Model for Large-Scale Structure
• Accelerating Large-Scale-Structure data analyses by emulating Boltzmann solvers and Lagrangian Perturbation Theory
• CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators
• The effect of photometric redshift uncertainties on galaxy clustering and baryonic acoustic oscillations
• Capse.jl: efficient and auto-differentiable CMB power spectra emulation
• Discretize-Optimize vs. Optimize-Discretize for Time-Series Regression and Continuous Normalizing Flows
• $\text{matryoshka}$: Halo Model Emulator for the Galaxy Power Spectrum
• Shell-crossing in a $\Lambda$CDM Universe
• 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
• The Einstein-Boltzmann equations revisited

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