---
title: 'DISCO-DJ: Differentiable Cosmology in Jax'
url: https://www.emergentmind.com/topics/disco-dj-differentiable-simulations-for-cosmology-done-with-jax
type: topic
---

# DISCO-DJ: Differentiable Cosmology in Jax

DISCO-DJ (DIfferentiable Simulations for COsmology – Done with Jax) is a software ecosystem for cosmological simulation and inference, providing a comprehensive, fully differentiable framework for the forward modeling of both linear and non-linear cosmic structure formation. Written in Python using the Jax library, DISCO-DJ leverages automatic differentiation (autodiff), just-in-time (JIT) GPU acceleration, and adjoint methods to deliver memory-efficient, high-fidelity, and fast computations for cosmological observables, ranging from the linear transfer function to the non-linear matter density field. The system is designed for full-field level inference, robust parameter estimation, and seamless integration into modern Bayesian inference pipelines, enabling the extraction of maximal information from large-scale structure surveys and facilitating the development of hybrid simulation–machine learning workflows.

## 1. Architecture and Design Principles

DISCO-DJ comprises modular components centered on differentiable solutions of cosmological evolution equations: a linear Einstein–Boltzmann solver and a non-linear particle–mesh (PM) $N$-body module. The code is implemented in Jax, a high-performance array library that enables:

- **Full automatic differentiation** for all numerical computations
- **Efficient JIT compilation** for GPU/TPU acceleration
- **Forward- and reverse-mode differentiation** (crucial for field-level inference)
- **Custom adjoint evolution** for memory efficiency in high-dimensional simulations

This design marks a departure from legacy cosmological codes (e.g., CAMB, CLASS, Gadget), which are non-differentiable “black boxes” with limited applicability to advanced inference schemes requiring derivatives of output observables with respect to high-dimensional parameter spaces (e.g., initial conditions, cosmological parameters).

## 2. Differentiable Einstein–Boltzmann Solver (DISCO-DJ I)

The Einstein–Boltzmann module in DISCO-DJ numerically integrates the linearized Einstein–Boltzmann equations for cosmological perturbations in a fully differentiable manner [2311.03291]. In synchronous gauge, for example, the evolution of scalar metric perturbations $\eta$, baryon and cold dark matter perturbations, and massive neutrinos is given by:

\[
\begin{align*}
k^2 \eta - \frac{1}{2}\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 \\
\delta_b' &= -\theta_b - \frac{1}{2}h' \\
\theta_b' &= -\mathcal{H}\theta_b + c_s^2 k^2 \delta_b + \frac{4}{3} \frac{\bar{\rho}_\gamma}{\bar{\rho}_b} a n_e \sigma_T (\theta_\gamma - \theta_b)
\end{align*}
\]

and similar hierarchies for photon and neutrino perturbations.

- **Automatic Differentiation**: By using the Diffrax package (within Jax), both discretize-then-optimize (forward-mode) and adjoint (reverse-mode) derivatives of the solution with respect to all input cosmological parameters are accessible.
- **Validation**: The module produces matter power spectra, transfer functions, and other observables in per-mille agreement with CAMB and CLASS—including for massive neutrinos and general dark energy parameterizations.
- **Fisher Forecasting**: Exact Jacobians of power spectra with respect to parameters are used to construct Fisher matrices for rigorous survey forecasts.
- **Modularity**: The structure readily allows for extensibility to non-standard physics (e.g., modified gravity, additional neutrino species).

## 3. Differentiable Particle–Mesh N-body Simulations (DISCO-DJ II)

The PM module provides a fully differentiable simulation of mildly non-linear cosmic structure formation [2510.05206]. Key attributes include:

- **Theory-informed time integrators**—notably the BullFrog method—which reproduce 2LPT trajectories in the pre-shell-crossing regime:

  \[
  \begin{align*}
  X^{n+1/2} &= X^n + \tau_1 V^n \\
  V^{n+1} &= \alpha V^n + \beta A(X^{n+1/2}) \\
  X^{n+1} &= X^{n+1/2} + \tau_2 V^{n+1}
  \end{align*}
  \]
  with drift and kick coefficients $\alpha$, $\beta$ designed to match analytic perturbative growth.

- **Non-uniform FFT (NUFFT)** for direct Fourier transforms from particle positions, suppressing aliasing and controlling discreteness and spectral leakage.
- **Custom autodiff routines** for particle–mesh gridding and force assignment, supporting forward-mode, reverse-mode, and full adjoint time-stepping.
- **Adjoint formulation**: Reverse-mode differentiation is performed with memory cost independent of the number of time steps.
- **Scalability**: Simulations with $N = 512^3$ particles reach percent-level accuracy for the $z=0$ power spectrum at $k \approx 0.2\,h/\mathrm{Mpc}$ using only $\sim$6 BullFrog steps (runtime: a few seconds on modern GPUs).
- **Field-level inference**: Enables direct optimization or sampling of high-dimensional initial condition fields and cosmological parameters by comparing full simulated density fields with data, e.g., via Hamiltonian Monte Carlo.

## 4. Numerical and Algorithmic Innovations

DISCO-DJ achieves high performance and accuracy by integrating several advanced numerical strategies:

- **Custom kernels** for mass assignment and grid interpolation—supporting higher-order schemes (e.g., TSC, PCS), de-aliasing via interlacing, and Lagrangian sheet-based resampling.
- **Adjoint time integration** for memory-efficient backpropagation through hundreds of time steps—critical for reverse-mode autodiff in field-level inference.
- **GPU-accelerated computation** throughout, including FFTs, gridding, and vectorized numerical operations, yielding order-of-magnitude speed gains over comparable codes [2510.05206].
- **Detailed accuracy studies**: Extensive benchmarks quantify the impact of time-stepping, grid density, force resolution, and discreteness corrections, ensuring reliability for cosmological analyses.

## 5. Applications: Field-level Bayesian Inference and Self-consistent Cosmological Pipelines

A primary motivation for DISCO-DJ is field-level Bayesian inference, an approach that utilizes the full spatial information in density fields and galaxy distribution for cosmological parameter estimation. The differentiable nature of the DISCO-DJ pipeline is essential for:

- **Efficient gradient-based sampling** (e.g., HMC) in high-dimensional parameter spaces, where typical likelihoods depend on hundreds of thousands of parameters (initial condition amplitudes plus cosmological variables).
- **Loss functions** that can involve arbitrary field-level summary statistics or the full three-dimensional field, with gradients propagated through all simulation components.
- **Seamless coupling with differentiable Einstein–Boltzmann solvers**, yielding a fully self-consistent mapping from primordial parameters to late-time observables.
- **Training and emulation**: The framework generates large, self-consistent simulation suites for neural network–based emulators, or directly supports hybrid pipelines (e.g., with PyBird-JAX [2507.20990], CosmoPower-JAX [2305.06347]).

## 6. Comparison with Other Differentiable and Non-differentiable Codes

- **Legacy codes (CAMB, CLASS, standard $N$-body):** Cannot supply exact derivatives of output observables with respect to input parameters, limiting their use in modern gradient-based inference and making survey optimization or parameter dependence expensive and numerically fragile.
- **Related Jax-based efforts (pmwd, CosmoPower-JAX, JAX-cosmo, GODMAX, SwiftC$_\ell$, halox):** DISCO-DJ distinguishes itself by providing both a differentiable linear solver and a high-fidelity, memory-efficient, and differentiable PM code with end-to-end pipeline integration, adjoint differentiation, and scale-appropriate numerical accuracy.
- **Innovative features**: The BullFrog integrator, NUFFT force computation, and robust handling of forward- and reverse-mode autodiff set DISCO-DJ apart in terms of both speed and theoretical soundness.

## 7. Open-source Availability and Future Directions

DISCO-DJ is open-source and available at https://github.com/cosmo-sims/DISCO-DJ, facilitating transparent cross-comparison, extension, and integration into broader cosmological data analysis infrastructures [2510.05206]. Its modular Python/Jax codebase simplifies the development of new physical modules (e.g., baryonic effects, alternative gravity, biased tracer models) and enables the rapid prototyping of differentiable scientific simulators. Planned extensions include:

- Incorporation of higher-order $N$-body solvers and baryonic subgrid physics with differentiable surrogates
- Efficient hybridization with machine learning approaches for improved inference in high-dimensional observational data spaces
- End-to-end pipelines covering Einstein–Boltzmann, mildly non-linear, and strongly non-linear regimes within a consistent autodiff environment

DISCO-DJ thus provides a technically rigorous and extensible foundation for next-generation, gradient-based cosmological inference, emphasizing physical fidelity, scalability, and full differentiability at every stage in the modeling process.

Source: https://www.emergentmind.com/topics/disco-dj-differentiable-simulations-for-cosmology-done-with-jax