Quickstart

Quickstart#

DISCO-EB is a differentiable linear Einstein-Boltzmann solver for cosmology, and the Einstein-Boltzmann module of the DISCO-DJ framework (DIfferentiable Simulations for COsmology, Done with Jax). It evolves the background, the recombination history and the linear perturbations of every wave mode of a batch of cosmologies in one GPU launch.

The solve runs as a modax Rodas5P kernel. RECFAST and the optical depth step alongside the perturbation hierarchy in the same kernel, so the solve needs no tables. You choose which quantities to keep at which times, and the kernel evaluates them as it goes. The results are JAX arrays, differentiable in the cosmological parameters by forward sensitivities. DISCO-EB supports flat or curved ΛCDM, a CPL dark-energy fluid \((w_0, w_a)\), massive neutrinos and reionization.

Installation#

pip install "disco-eb[cuda13]"   # or [cuda12], which Python 3.10 needs

You need a GPU to run a solve, but not to import the package. modax brings the CUDA compiler stack with it. The CMB spectra live in a separate package, DISCO-CMB, which builds on this one.

Your first solve#

This example keeps the comoving curvature perturbation \(\mathcal{R}\) and the CDM density contrast of 64 wave modes at 100 conformal times:

import jax

jax.config.update("jax_enable_x64", True)

import numpy as np
from discoeb import solve_eb

h = 0.6732
cosmology = {
    "Omegam": 0.1424 / h**2, "Omegab": 0.02238 / h**2, "H0": 100 * h,
    "w_DE_0": -1.0, "w_DE_a": 0.0, "cs2_DE": 1.0, "Omegak": 0.0,
    "A_s": 2.1e-9, "n_s": 0.966, "k_p": 0.05,
    "Tcmb": 2.7255, "YHe": 0.2454, "Neff": 3.046, "Nmnu": 0.0, "mnu": 0.0,
}
k = np.geomspace(1e-4, 1.0, 64)      # 1/Mpc
tau = np.geomspace(1.0, 2e4, 100)    # Mpc, the last times past today

solution = solve_eb(cosmology, k, tau, keep=("R", "clxc"))
R = solution["R"]  # (1, 64, 100): cosmologies, wave modes, times

Pass a list of dictionaries to solve a batch of cosmologies in the same launch. A time past a cosmology’s conformal age holds its values today, which solution.tau0 reports.

Where to go next?#

Keeping quantities

What keep accepts, and how to take derivatives.