DISCO-DJ: Differentiable Cosmology in JAX

• DISCO-DJ is an end-to-end differentiable simulation ecosystem for cosmological large-scale structure, integrating JAX for GPU acceleration and automatic differentiation.
• The system modularly links Einstein–Boltzmann solvers, LPT initial conditions, N-body and hydrodynamics modules to support field-level inference, Fisher forecasting, and high-dimensional MCMC.
• Its innovative adjoint method and JIT-compiled modules deliver memory optimality and sub-per-mille gradient accuracy, enabling efficient Bayesian inference and survey design.

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 (Hahn et al., 2023).
• 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 (List et al., 6 Oct 2025).
• 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 (List et al., 6 Oct 2025, Li et al., 2022, Li et al., 2022).
• 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 (Horowitz et al., 4 Feb 2025).
• 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 (Alarcon et al., 31 Oct 2025).

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 (Hahn et al., 2023).
• 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 (List et al., 6 Oct 2025, Li et al., 2022, Li et al., 2022). 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 (Horowitz et al., 4 Feb 2025).

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 (Li et al., 2022, Li et al., 2022, List et al., 6 Oct 2025). 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 (Horowitz et al., 4 Feb 2025).
• 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 (Horowitz et al., 4 Feb 2025, List et al., 6 Oct 2025).

• 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 (Hahn et al., 2023, List et al., 6 Oct 2025, Li et al., 2022).
• 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 (Li et al., 2022, List et al., 6 Oct 2025).
• 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 (Hahn et al., 2023, List et al., 6 Oct 2025).
• 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 (List et al., 6 Oct 2025).

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 (Horowitz et al., 4 Feb 2025, List et al., 6 Oct 2025, Hahn et al., 2023).
• 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 (Alarcon et al., 31 Oct 2025).
• Survey Forecasting & Experimental Design: Fisher matrices are constructed with exact gradients, leveraging vectorized, machine-precision Jacobians for optimal design and error propagation (Hahn et al., 2023).
• 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 (Horowitz et al., 4 Feb 2025, List et al., 6 Oct 2025, Li et al., 2022).
• 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 (Hahn et al., 2023, List et al., 6 Oct 2025, Horowitz et al., 4 Feb 2025).

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 (Alarcon et al., 31 Oct 2025).

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 (Horowitz et al., 4 Feb 2025, Hahn et al., 2023, List et al., 6 Oct 2025, Li et al., 2022).

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.