DISCO-DJ II: a differentiable particle-mesh code for cosmology

TL;DR

DISCO-DJ II delivers a differentiable, GPU-accelerated PM N-body framework for cosmology, combining advanced time integrators (notably BullFrog) with LPT/PPT perturbation theories and a differentiable force solver. The architecture supports forward and reverse autodiff, with an adjoint method that achieves $O(1)$ memory for backpropagation through time, enabling efficient gradient-based field-level inference and end-to-end coupling to an Einstein--Boltzmann solver. Through extensive validation, it demonstrates fast, accurate predictions on mildly non-linear scales (e.g., $P(k)$ accurate to sub-percent up to $k\sim0.4\, h/\mathrm{Mpc}$ with ~100 steps) and provides guidance on force-approximation choices, resolution, and time-stepping. The work enables explicit inference of initial conditions and cosmological parameters from three-dimensional density fields and lays groundwork for future extensions to lightcones, redshift-space distortions, bias modeling, and distributed multi-GPU runs, all within a fully differentiable pipeline.

Abstract

The mildly non-linear regime of cosmic structure formation holds much of the information that upcoming large-scale structure surveys aim to exploit, making fast and accurate predictions on these scales essential. We present the $N$-body module of DISCO-DJ (DIfferentiable Simulations for COsmology - Done with Jax), designed to deliver high-fidelity, GPU-accelerated, and differentiable particle-mesh simulations tailored for cosmological inference. Theory-informed time integrators such as the recently introduced BullFrog method allow for accurate predictions already with few time steps (e.g. $6$ steps for per-cent-level accuracy in terms of the present-day power spectrum at $k \approx 0.2 \, h / \mathrm{Mpc}$ using $N = 512^3$ particles, which takes just a few seconds). To control discreteness effects and achieve high accuracy, the code incorporates a suite of advanced techniques, for example a custom non-uniform FFT implementation for force evaluation. Both forward- and reverse-mode differentiation are supported, with memory requirements independent of the number of time steps; in the reverse case, this is achieved through an adjoint formulation. We extensively study the effect of various numerical parameters on the accuracy. As an application of DISCO-DJ, we perform field-level inference by recovering $σ_8$ and the initial conditions from a noisy Gadget matter density field. Coupled with our recently introduced Einstein--Boltzmann solver, the DISCO-DJ ecosystem provides a self-consistent, fully differentiable pipeline for modelling the large-scale structure of the universe. The code is available at https://github.com/cosmo-sims/DISCO-DJ.

Figures (20)

• Figure 1: High-level structure of Disco-Dj. Given a white noise realisation $w \in \mathbb{R}^N$ of the initial density fluctuations -- whose dimensionality $N$ typically equals the number of $N$-body particles -- and a set of cosmological parameters $\theta$, Disco-Dj enables forward modelling the non-linear gravitational collapse in an end-to-end autodifferentiable and GPU-accelerated manner using Jaxjax2018github. The linear power spectrum is computed by solving the (linearised) Einstein--Boltzmann equations as presented in Ref. Hahn2023DISCO-DJCosmology (Disco-Dj I). The non-linear structure formation module is introduced in this work (Disco-Dj II), consisting of perturbative models and a fast $N$-body PM code with theory-informed time integrators. Extensions of Disco-Dj regarding bias modelling etc. are currently in preparation.
• Figure 2: Simulated matter density slices through a box of comoving side length $L = 100 \, \mathrm{Mpc} / h$ at redshifts $z = 3$ (left half) and $z = 0$ (right half), computed with Gadget-4 (upper left) and with various methods implemented in Disco-Dj. The slices are averaged over $25 \, \mathrm{Mpc} / h$ along the third dimension. Specifically, we consider a particle-mesh (PM) simulation with 10 BullFrog time steps (upper right), second-order Lagrangian perturbation theory (2LPT, lower left), and propagator perturbation theory (PPT, lower right). The number of particles / fluid elements is $N = 1024^3$ for Gadget-4 and $N = 512^3$ for all methods in Disco-Dj. The colour mapping is logarithmic w.r.t. density, and the colour bar limits are individual for each redshift, but shared across the four panels. At $z = 3$, the predictions by the perturbative methods in the bottom row still agree fairly well with the Gadget-4 reference. In contrast, at $z = 0$ the filamentary structures with these methods are smeared out, and interference patterns are visible in the PPT panel. Judging by eye, our PM simulation reproduces the reference very accurately, although small differences can be spotted by eye, e.g. the ellipticity of the cluster slightly below the vertical centre near the right edge. For quantitative results, see Sec. \ref{['sec:validation']}.
• Figure 3: Fourier-transformed kernels used for interpolation (left) and derivatives (right), where the wave number $k$ is given in units of the Nyquist wave number $k_{\mathrm{Nyquist}}$. As for the interpolation kernels, a larger support (such as for PCS) suppresses the power more strongly. The inset plot illustrates the kernels in real space, where $x$ has units of grid cells, and the $y$-axis is such that the areas under the curves integrate to unity. The finite-difference derivative kernels closely match the exact derivatives on large scales, but fall off as $k \to k_{\mathrm{Nyquist}}$, with lower order implying an earlier drop.
• Figure 4: Thin slice of the gradient of the non-linear power spectrum w.r.t. the white noise field $\partial P/\partial w$ through a Disco-Dj simulation with 10 time steps using the BullFrog stepper, evaluated for three different $k$-bins. The results in the upper left corners have been computed using standard (discretise-then-optimise) reverse-mode autodiff and the default VJP for the gather and scatter operations for the density computation in the simulation, whereas the lower right corners show the results with the adjoint (optimise-then-discretise) method and custom VJP operations. The transition is smooth, confirming the correctness of our implementation (see Fig. \ref{['fig:dPk_dw_hist']} for quantitative results). The colour scale is normalised individually for each $k$-bin.
• Figure 5: Histogram of differences in the power spectrum gradient $\partial P/\partial w$ for a single $k$-bin, normalised by the standard deviation of the autodiff result, for the same scenario as in Fig. \ref{['fig:dPk_dw']}. We compare two independent adjoint runs (blue), two autodiff runs (grey), and a cross-comparison between adjoint and autodiff (red). The vertical dashed line marks zero difference. The backwards integration leads to somewhat noisier gradients with the adjoint method than with standard autodiff; still, the deviations are small in comparison to the magnitude of $\partial P/\partial w$ itself.