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

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 ( DI fferentiable S imulations for CO smology – D one with J ax ), 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 ≈ 0.2 h / 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 https://github.com/cosmo-sims/DISCO-DJ .