DISCO-EB

Installation • Quickstart • Documentation • Paper • Citation

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.

The documentation covers the API and how to use it. The paper describes the method and its accuracy against CLASS.

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.

Quickstart

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

Development

The repository uses uv:

git clone https://github.com/cosmo-sims/DISCO-EB.git && cd DISCO-EB
uv sync --extra cuda13  # or --extra cuda12

uv run pytest tests -rs   # needs a GPU, and -rs shows skipped CLASS comparisons
uv run ruff format && uv run ruff check --fix
uv run ty check
uv run pre-commit install   # the checks above, and a commit-message check

The documentation (Sphinx with MyST and autodoc) builds in an environment of its own, the package's CPU install plus the docs group:

export UV_PROJECT_ENVIRONMENT=.venv-docs
uv sync --no-default-groups --group docs
uv run --no-sync sphinx-autobuild docs docs/_build/html   # live preview
uv run --no-sync sphinx-build -W docs docs/_build/html    # what CI runs

Citation

If you use DISCO-EB in your research, please cite the JCAP article:

@article{DISCO-EB:2024,
    author = "Hahn, Oliver and List, Florian and Porqueres, Natalia",
    title = "{DISCO-DJ I: a differentiable Einstein-Boltzmann solver for cosmology}",
    eprint = "2311.03291",
    archivePrefix = "arXiv",
    primaryClass = "astro-ph.CO",
    doi = "10.1088/1475-7516/2024/06/063",
    journal = "JCAP",
    volume = "06",
    pages = "063",
    year = "2024"
}

License

DISCO-EB uses the MIT License. The recombination module adapts code from Recfast, whose UBC license its file header reproduces.