DISCO-DJ I: Differentiable Einstein–Boltzmann Solver

• Differentiable Einstein–Boltzmann solvers are computational frameworks that leverage automatic differentiation to compute cosmological observables and their gradients efficiently.
• DISCO-DJ I employs a modular JAX-based integration of stiff ODEs, achieving per-mille numerical accuracy and fast Fisher forecasting for cosmological surveys.
• Its extensible architecture supports rapid incorporation of new physics modules, enabling end-to-end pipelines with GPU acceleration and precise derivative computations.

A Differentiable Einstein–Boltzmann Solver (DISCO-DJ I) is a computational framework that provides an end-to-end, fully differentiable solution to the linearized Einstein–Boltzmann equations governing cosmological perturbations. Developed within the DISCO-DJ (DIfferentiable Simulations for COsmology – Done with JAX) software package, this solver computes not only the cosmological observables themselves but also their Jacobians with respect to all input parameters via automatic differentiation, fundamentally transforming cosmological inference and forecasting workflows (Hahn et al., 2023).

1. Mathematical Formulation of the Einstein–Boltzmann System

The foundational element of DISCO-DJ I (“Disco-EB”) is the set of coupled “stiff” ordinary differential equations (ODEs) that describe the evolution of linear perturbations around the homogeneous Friedmann–Lemaître–Robertson–Walker (FLRW) background. In the synchronous gauge and adopting the notation of Ma & Bertschinger (1995), the metric perturbations—denoted by $h$ and $\eta$—satisfy the linearized Einstein equations in Fourier space:

$$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$

Here, prime denotes differentiation with respect to conformal time $\tau$, $\mathcal{H} = a'/a$, and the quantities $\theta$ and $\sigma$ represent velocity divergence and shear, respectively.

For each cosmological fluid or particle species (Cold Dark Matter, baryons, photons, massless and massive neutrinos, and clustering dark energy), the solver incorporates the respective continuity and Euler equations, Legendre multipole expansions for relativistic components, and momentum discretization for massive neutrinos (using a 5-point Gauss–Laguerre quadrature). The dark energy sector supports a Chevallier–Polarski–Linder equation of state $w(a) = w_0 + w_a(1-a)$ with a constant sound speed.

The full system is written compactly as:

$$\mathbf{y}'(\tau) = \mathbf{f}(\tau, \mathbf{y}, \boldsymbol{\theta})$$

where $\mathbf{y}$ collects all perturbation variables, and $\eta$0 specifies the cosmological and nuisance parameters, including $\eta$1, $\eta$2, $\eta$3, $\eta$4, $\eta$5, $\eta$6, $\eta$7, $\eta$8, $\eta$9, $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$0, among others.

2. Implementation in JAX and Integration Techniques

DISCO-DJ I is implemented in pure Python using the JAX ecosystem, specifically leveraging the Diffrax library. The right-hand side of the ODE system, $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$1, is formulated as a JAX-compatible function that receives arrays $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$2 and outputs $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$3. Integration is performed using the high-order, singly diagonally implicit Runge–Kutta method Kvaerno5.

Key techniques include:

• All JAX primitives allow native use of automatic differentiation, enabling on-the-fly Jacobian matrix computation without hand-written derivatives.
• Each Fourier mode $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$4 is evolved independently, resulting in block-diagonal Jacobians and facilitating parallelization via jax.vmap on GPUs.
• Modular design: each species' set of equations and the recombination solver are written as separate JAX functions, supporting rapid development of new physics modules (e.g., modified gravity, exotic fluids).
• For massive neutrinos, a 5-point momentum quadrature achieves per-mille accuracy.

The thermal history (recombination) is handled by a fast “Recfast-mini” module following Seager et al. (1999), sufficient for $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$5 accuracy in late-time matter power spectrum predictions.

3. Automatic Differentiation and the Tangent Hyperplane in Parameter Space

A central innovation of DISCO-DJ I is its full compatibility with JAX's forward- and reverse-mode automatic differentiation. Any observable output $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$6 is a differentiable function with respect to $$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$7, such that, in JAX syntax:

$$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$8

$$\begin{align*} k^2\,\eta-\tfrac12\,\mathcal{H}\,h' &= 4\pi G\,a^2\,\delta T^0{}_0 \ k^2\,\eta' &= 4\pi G\,a^2(\bar\rho+\bar P)\,\theta \ h''+2\mathcal{H}\,h'-2k^2\,\eta &= -8\pi G\,a^2\,\delta T^i{}_i \ h''+6\,\eta''+2\mathcal{H}(h'+6\,\eta')-2k^2\,\eta &= -24\pi G\,a^2(\bar\rho+\bar P)\,\sigma \end{align*}$$9

This tangent hyperplane is computed at a computational cost only marginally higher than a single function evaluation in forward-mode AD, vastly outperforming finite-difference methods which require a full solver call per parameter. As a result, DISCO-DJ I directly provides:

• Gradient information for use in Hamiltonian Monte Carlo and other sampling-based inference frameworks
• Sensitivity analysis for cosmological parameters
• Efficient, high-precision Fisher-matrix forecasting for cosmological surveys
• Optimal data compression via derivative-informed summary statistics

4. Numerical Accuracy and Validation Against Existing Solvers

The accuracy of DISCO-DJ I has been validated against established codes CAMB and CLASS, including cases with massive neutrinos and dark energy models parameterized by $\tau$0. Using a baseline consistent with Planck 2018 cosmology, the relative differences in the computed total matter power spectrum at $\tau$1 and $\tau$2 remain below several $\tau$3, with largest deviations confined to dynamically negligible regimes (e.g., neutrino power on small scales where $\tau$4). All relevant species (CDM, baryons, massive neutrinos, total matter) exhibit per-mille agreement over the parameter space relevant for cosmological inference.

5. Cosmological Applications: Fisher Forecasting and Large-Scale Structure

DISCO-DJ I enables differentiable, end-to-end pipelines for data analysis and forecasting. As a demonstration, the pipeline computes the Fisher information matrix for the spectroscopic Euclid galaxy clustering survey. The observed redshift-space power spectrum includes Alcock–Paczynski and Finger-of-God effects:

$\tau$5

The Fisher element per tomographic bin is:

$\tau$6

Since the derivatives $\tau$7 are available via automatic differentiation, no finite differencing is required. The resulting parameter forecast—the Fisher ellipses—exhibits agreement at the few-percent level with official Euclid Fisher code results.

6. Architecture, Modularity, and Extensibility

DISCO-DJ I’s codebase is markedly concise (approximately 2,000 lines) and designed for extensibility. The modular structure allows for the isolation and augmentation of individual physics components:

• Equations of Motion (EOM) for each species and the thermal history chain are independent, facilitating the incorporation of additional physics such as modified gravity, early/interacting dark energy, and alternative dark matter models.
• The implementation currently does not perform on-the-fly switching for tight-coupling or ultra-relativistic approximations, concentrating the solver on full system integration for maximum fidelity. Such optimizations are earmarked for future development.
• GPU acceleration is inherent via JAX, and the entire computational stack (ODES, recombination, quadrature, power spectrum, correlation transforms) is differentiable.

7. Prospects and Related Developments

DISCO-DJ I represents a new paradigm in cosmological simulation pipelines by integrating differentiable programming with state-of-the-art numerical cosmology. Anticipated extensions include:

• Incorporation of modified gravity models (e.g., Horndeski/EFT)
• Modeling of early or interacting dark-energy/radiation components
• Implementation of nonlinear corrections via differentiable perturbation theory or particle–mesh approaches
• Direct integration with gradient-based inference engines, such as Hamiltonian Monte Carlo and normalizing flows
• Derivative-informed emulators for compressed data analysis

The codebase is openly available at https://github.com/ohahn/DISCO-EB. Further community-driven contributions have the potential to establish DISCO-DJ I as a core tool for differentiable cosmological modeling and simulation (Hahn et al., 2023). A notable parallel effort is "SymBoltz.jl: a symbolic-numeric, approximation-free and differentiable linear Einstein-Boltzmann solver" (Sletmoen, 29 Sep 2025), which introduces a Julia-based symbolic-numeric environment, further underscoring the field-wide shift toward fully differentiable, composable, and modular Einstein–Boltzmann solvers.