---
title: 'DISCO-DJ: Differentiable Cosmology in JAX'
url: https://www.emergentmind.com/topics/disco-dj-differentiable-simulations-for-cosmology-done-with-jax-0bec186c-5c88-4add-a7fb-ca8befc66a96
type: topic
---

# DISCO-DJ: Differentiable Cosmology in JAX

DISCO-DJ (DIfferentiable Simulations for COsmology – Done with JAX) is an ecosystem of end-to-end, fully differentiable simulation and inference codes for cosmological large-scale structure, written in the JAX framework to enable GPU acceleration and automatic differentiation (autodiff) across all physical modules. The architecture rigorously enforces differentiability from the level of initial inflationary conditions, through linear and mildly non-linear evolution, to baryonic effects and observables. DISCO-DJ’s differentiable design supports field-level inference, Fisher forecasting, and high-dimensional MCMC/variational optimization schemes that would be computationally prohibitive with traditional simulation codes.

## 1. Architectural Overview and Modular Design

The core design principle of DISCO-DJ is modularity: each pipeline stage implements a fundamental cosmological solver or physics module, exposing a JAX-compliant, JIT-compiled API. The major modules include:

- **Einstein–Boltzmann Solver (DISCO-DJ I):** Differentiable integration of the linear perturbation equations for all relevant species (CDM, baryons, photons, massless/massive neutrinos, dark energy) in synchronous gauge. Provides linear transfer functions, matter power spectra, and their parameter gradients at sub-per-mille precision. Utilizes the Diffrax 7-stage Kværnø5 ESDIRK integrator for stiff ODEs with on-the-fly Jacobian evaluation [2311.03291].
- **Initial Condition Generation:** Sampling of Gaussian random fields at high redshift. Second- or higher-order Lagrangian Perturbation Theory (LPT) is applied to forward particles to starting redshift, with arbitrary LPT order exposed in the interface [2510.05206].
- **N-body Particle-Mesh Evolution (DISCO-DJ II):** Time integration via theory-informed (BullFrog) or standard symplectic schemes. Gravity is solved with custom FFT/NUFFT (non-uniform FFT) force solvers, supporting mass assignment with various kernels (CIC/TSC/PCS). All scatter/gather and FFT operations are implemented in JAX with custom adjoints for memory- and compute-efficiency [2510.05206, 2211.09958, 2211.09815].
- **Baryonic Hydrodynamics and Feedback:** Hybrid schemes solve for gas dynamics using explicit upwind finite volume TVD algorithms, with a PM solver for DM. Stochastic subgrid processes (e.g., star formation, feedback) are made differentiable through Gumbel–Softmax reparameterization, so that autodiff flows through all feedback events [2502.02294].
- **Galaxy-Halo Connection Modules:** Fully differentiable galaxy population models such as DiffstarPop provide analytical, vectorized, and differentiable surrogates for mapping halo mass assembly to galaxy star formation and photometric observables, suitable for joint cosmology–astrophysics parameter inference [2510.27604].

Each module is written as a stateless, pure JAX function, allowing composition in arbitrary order. The system is structured to enable full-pipeline autodiff (with `jax.grad`, `jax.jacrev`, or `jax.jacfwd`) with machine precision from any observable back to all parameters.

## 2. Differentiable Physical Solvers and Governing Equations

DISCO-DJ’s physical kernels are implemented to allow XLA-backed differentiation at every computational step. Key solver details include:

- **Einstein–Boltzmann Equations:** All ODEs for perturbations (including multipole hierarchies and neutrino momentum bins) are encoded such that the right-hand-sides and their Jacobians are automatically differentiated. Observable derivatives ∂ln P(k)/∂θ for all cosmological parameters are produced in a single pass, without finite differencing or manual Jacobian assembly [2311.03291].
- **N-body Gravity and Adjoint Method:** Main DM evolution employs a symplectic DKD integrator on a fixed grid, with force evaluation via FFT/NUFFT Poisson solves. The adjoint method reconstructs forward states on the fly during backpropagation, keeping memory cost O(1) in number of time steps and enabling gradients for deep networks in high-dimensional systems [2510.05206, 2211.09958, 2211.09815]. The BullFrog integrator further guarantees exact recovery of 2LPT trajectories to shell-crossing.
- **Hydrodynamical Solver Differentiability:** The Euler equations for baryons are discretized with explicit TVD finite-volume schemes and Strang splitting, using a differentiable Van Leer limiter. Stochastic subgrid models employ continuous relaxations (e.g., Gumbel–Softmax) with temperature τ≈0.1 for stable yet faithful gradient propagation. All update stencils, fluxes, and source terms propagate derivatives as efficiently as the forward pass [2502.02294].

This strictly vectorized, functional approach ensures compatibility with JAX’s transformation stack, including JIT compilation (`jax.jit`), vectorization (`jax.vmap`), and parallelization (`jax.pmap`) for both forward and reverse-mode derivatives.

## 3. Memory-Optimal Gradient Computation: The Adjoint Approach

A central technical challenge in differentiable physical simulation is the memory cost of retaining the full forward trajectory for backpropagation. Naive reverse-mode AD schemes would require storage of O(N_steps × N_particles), rapidly exceeding GPU memory at survey-scale resolutions.

DISCO-DJ’s N-body module and pmwd both solve this via an explicit adjoint method: during the backward pass, the forward pass is replayed in reverse, synchronously integrating the discrete adjoint equations for each DKD map. The method is expressed as:

\[
λ_n = (∂M/∂(q_n, p_n))^\mathsf{T} λ_{n+1}, \qquad \frac{dL}{dθ} = \sum_{n=0}^{N-1} λ_{n+1}^\mathsf{T} \frac{∂M}{∂θ}(q_n, p_n; θ) + \frac{∂L}{∂θ}|_{n=N}
\]

Memory cost is then determined only by storage of the current state and adjoint variables, typically <1 GB even for 256³+ grids and 100+ time steps [2211.09958, 2211.09815, 2510.05206]. Gradient accuracy is validated to RMS ≲10⁻⁵ versus brute-force AD.

## 4. End-to-End Bayesian Inference with Autodiff

DISCO-DJ’s full-chain differentiability enables Bayesian posterior sampling and gradient-based likelihood evaluation across parameter spaces of O(10–100+), using statistics such as the late-time power spectrum, bispectra, or even full density fields. Practically:

- **Hamiltonian Monte Carlo (HMC):** The full log-likelihood is wrapped as a JAX function. With gradients provided via autodiff, the No-U-Turn Sampler (NUTS, e.g., via BlackJAX) explores the posterior efficiently, with effective sample sizes ≳6000 in ≈500 steps for {Ω_m, σ₈, δ_c, E_SN}, and field-level optimization of initial conditions [2502.02294].
- **Field-level Losses:** For field-level inference, e.g., of initial modes s, a quadratic loss

  \[
  -\log P(s|d) = \frac12 \sum_i \frac{(d_i - F(s)_i)^2}{N_i} + \frac12 \lVert s/T(k) \rVert^2
  \]

  is minimized with L-BFGS (Optax), with gradients flowing through initial condition generation, N-body evolution, and hydrodynamics. Large-scale modes are reconstructed down to k~few h/Mpc at S/N=1 [2502.02294, 2510.05206].

- **Batch Inference and Emulator Training:** Through `jax.vmap`, derivatives with respect to cosmological and subgrid parameters can be batch-evaluated, providing direct input to emulator training and survey pipeline optimization.

## 5. Performance, Numerical Fidelity, and Scaling

DISCO-DJ achieves high performance and scaling through careful exploitation of JAX-XLA, memory-efficient adjoints, and hardware accelerators:

- **JIT Compilation:** Complete pipelines (from initial conditions through power-spectra estimation) are wrapped with `@jax.jit`, resulting in fused GPU kernels for time-stepping, FFT, and field operations [2311.03291, 2510.05206, 2211.09958].
- **GPU/TPU Acceleration:** For a 128³ box, forward+backward passes (100 steps) complete in ≈10 s on a single A100 GPU; for 256³, under a minute. On A100, PM N-body forward+adjoint pass at 256³ in ≈18 s, with memory less than 4 GB [2211.09958, 2510.05206].
- **Convergence and Accuracy:** Percent-level agreement with reference solvers (CAMB, CLASS, Gadget) is achieved (e.g., ≤1% error in P(k) at k≈0.2 h/Mpc with 6 BullFrog steps at 512³), with sub-per-mille parameter derivatives in linear regimes [2311.03291, 2510.05206].
- **Discreteness and Aliasing Control:** Force accuracy is maintained via higher-order mass assignment, deconvolution, interlacing, and NUFFT schemes. Resolution and kernel order are tunable for performance/fidelity trade-off [2510.05206].

## 6. Applications, Limitations, and Integration with Next-Generation Inference

DISCO-DJ’s structure enables a suite of scientific applications:

- **Joint Inference Pipelines:** Seamless chaining of the Einstein–Boltzmann, LPT initial conditions, N-body, hydrodynamics, galaxy–halo connection, and instrument/noise modules. Enables joint inference on parameters spanning inflation, dark energy, feedback, and stochastic initial conditions [2502.02294, 2510.05206, 2311.03291].
- **Differentiable Field-Level Forward Modeling:** Pipeline design allows differentiability through fast emulators, field-level kSZ/CMB likelihoods, and forward-modeled galaxy observables. For example, DiffstarPop provides differentiable SFH distributions for 10⁶ galaxies in 0.03 s on GPU [2510.27604].
- **Survey Forecasting & Experimental Design:** Fisher matrices are constructed with exact gradients, leveraging vectorized, machine-precision Jacobians for optimal design and error propagation [2311.03291].
- **Limitations:** Main current restrictions include fixed-grid PM and Eulerian hydrodynamics (no adaptive mesh refinement), memory/performance bottlenecks for expansion to ultra-high-resolution (>256³) boxes, and reliance on piecewise-smooth modules for differentiability. Adaptive mesh and non-smooth solvers (e.g., N-body trees) are not yet fully differentiable, though neural surrogates and quasi-PM schemes are being explored for integration [2502.02294, 2510.05206, 2211.09958].
- **Extensibility:** New physics (e.g., warm dark matter, modified gravity) can be incorporated by encoding their perturbation equations in the differential system and letting autodiff manage the Jacobians [2311.03291, 2510.05206, 2502.02294].

## 7. Code Availability, Benchmark Summary, and Scientific Outlook

DISCO-DJ and its components are released as open-source Python packages, with detailed API documentation and reproducible examples. Key public repositories include:
- DISCO-DJ (main): [https://github.com/cosmo-sims/DISCO-DJ]
- pmwd: [github.com/eelregit/pmwd]
- DiffstarPop: integrated in main releases [2510.27604].

Benchmarks consistently demonstrate:
- Forward+adjoint gradient runs for 256³–512³ boxes with 100+ time steps are feasible within ∼1 minute per forward/backward pass (A100 GPU).
- Gradient accuracy validated at ≲10⁻⁵ RMS error to standard AD at tractable sizes.
- High-dimensional MCMC reveals effective sample sizes (ESS) on par with or exceeding traditional methods by 1–3 orders of magnitude.
- Full field-level and summary-statistic pipelines scaling efficiently to O(100) parameters [2502.02294, 2311.03291, 2510.05206, 2211.09958].

DISCO-DJ provides a rigorous computational framework for next-generation cosmological inference, with extensible, autodiff-ready solvers enabling comprehensive, high-fidelity forward modeling and statistical analysis directly coupled to raw observational data. Its design pattern is informing growing developments in differentiable field modeling, galaxy formation population synthesis, and survey analysis for future cosmological investigations.

Source: https://www.emergentmind.com/topics/disco-dj-differentiable-simulations-for-cosmology-done-with-jax-0bec186c-5c88-4add-a7fb-ca8befc66a96