Differentiable PM N-body Simulations: DISCO-DJ II

• The paper introduces DISCO-DJ II, a GPU-accelerated, fully differentiable particle–mesh code that accurately simulates cosmic structure formation for end-to-end parameter inference.
• It employs advanced time integration methods like BullFrog to achieve per-cent level accuracy with minimal time steps, leveraging FFT-based force computations and NUFFT for precision.
• The framework utilizes JAX’s automatic differentiation with an adjoint strategy, ensuring efficient gradient computation and seamless integration with modern inference and machine learning pipelines.

Differentiable Particle–Mesh (PM) N-body simulations represent a convergence of modern high-performance computational astrophysics, cosmological inference, and automatic differentiation frameworks. DISCO-DJ II is a GPU-accelerated, fully differentiable particle–mesh code tailored for the rapid and accurate simulation of the large-scale structure in the mildly non-linear regime of cosmic structure formation. Its design philosophy is to combine field-level accuracy, computational performance, and full differentiability—both in forward and reverse (adjoint) mode—thus enabling end-to-end, gradient-based inference of cosmological parameters and initial conditions directly from observational data (List et al., 6 Oct 2025).

1. Mathematical and Algorithmic Foundation

DISCO-DJ II simulates the evolution of cosmic matter using the Vlasov–Poisson system in comoving coordinates $(\mathbf{x},\mathbf{p})$, with background expansion parameterized by the scale factor $a$: $$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$ where $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$ (List et al., 6 Oct 2025).

In the $N$-body discretization, the phase-space density $f$ is represented by $N$ tracer particles, whose masses are deposited onto an $N_g^3$ mesh using mass-assignment kernels $W$, most commonly using CIC (linear), TSC (quadratic), or PCS (cubic spline). The deposited density $\rho(\mathbf{x}_g)$ is then used to solve Poisson’s equation via FFTs on the mesh, enabling rapid determination of the gravitational potential and subsequent acceleration computation. The resulting mesh-based gravitational force is interpolated back to the particle positions, closing the evolution loop (List et al., 6 Oct 2025, Li et al., 2022, Li et al., 2022).

2. Time Integration Schemes: BullFrog and Beyond

Accurate time integration is essential for minimizing systematic biases in cosmological predictions. DISCO-DJ II implements a flexible drift–kick–drift (DKD) integrator,

$a$0

with step size $a$1 and $a$2 representing acceleration. The parameters $a$3 specialize the scheme:

• FastPM: Enforces 1LPT-consistency and symplecticity by matching the Zel’dovich approximation (ZA), with $a$4, $a$5, where $a$6.
• BullFrog: Implements exact 2LPT consistency at each step:

$a$7

where $a$8 is the linear growth factor, $a$9 (List et al., 6 Oct 2025).

BullFrog yields per-cent-level accuracy in the power spectrum with as few as $$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$0–$$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$1 time steps at $$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$2, significantly enhancing efficiency in parameter inference pipelines.

3. Force Computation: Particle–Mesh and NUFFT Approaches

Force evaluation in DISCO-DJ II employs advanced mass assignment and interpolation:

• Mass assignment: Standard (CIC, TSC, PCS) kernels $$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$3 are used, possessing analytically known real- and Fourier-space forms; for instance, $$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$4 for the order-$$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$5 kernel.
• FFT-based PM: Poisson's equation is solved on the mesh as $$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$6, and the acceleration as $$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$7.
• Interpolation: The force is interpolated to particles via the same kernel or with a custom Non-Uniform FFT (NUFFT), employing an exponential-of-semicircle spreading kernel of tunable support, grid upsampling ($$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$81.25$$\partial_t f + \frac{\mathbf{p}}{a^2}\cdot\nabla_{\mathbf{x}} f - \frac{\nabla_{\mathbf{x}}\phi}{a}\cdot\nabla_{\mathbf{p}} f = 0;\quad \nabla^2_{\mathbf{x}} \phi = \frac{3\Omega_m H_0^2}{2}\delta,$$9), and frequency truncation. Aliasing can be controlled down to arbitrary accuracy by increasing the kernel size.
• Anti-aliasing and deconvolution: DISCO-DJ II implements interlaced meshes, explicit deconvolution (divide by $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$0), and advanced resampling to suppress artifacts (List et al., 6 Oct 2025).

4. Differentiable Architecture and Adjoint Gradient Computation

Every operation in DISCO-DJ II is written in JAX, supporting both forward- and reverse-mode automatic differentiation. While forward-mode gradient evaluation (Jacobian-vector product, JVP) can be straightforwardly parallelized, reverse-mode (vector-Jacobian product, VJP) naïvely incurs $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$1 memory overhead.

To address this, DISCO-DJ II implements a continuous adjoint (discretize-then-optimize) strategy. The adjoint drift–kick–drift scheme propagates cotangent variables $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$2 backward in time using the same integrator as the forward pass but requires only $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$3 memory with respect to the number of time steps. All gather/scatter operations exploit their linearity for custom JVP/VJP rules, ensuring end-to-end differentiability at minimal memory footprint. This mirrors regular adjoint-based PM codes such as pmwd (Li et al., 2022, Li et al., 2022).

5. Numerical Accuracy and Performance Benchmarks

Default settings for high-fidelity cosmological inference are as follows:

• $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$4 particles, box $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$5, initial conditions at $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$6 from 2LPT; PM grid $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$7 per dimension; CIC mass assignment, no deconvolution; BullFrog integrator, uniform stepping in $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$8 (List et al., 6 Oct 2025).
• With BullFrog:
  • $$\delta(\mathbf{x}) = \int f(\mathbf{x},\mathbf{p})\, d^3p - 1$$9 time steps: per-cent accuracy in $N$0 for $N$1Mpc.
  • $N$2 steps: $N$3 error up to $N$4Mpc.
  • $N$5 steps: $N$6 error at $N$7Mpc.

Timing benchmarks (single NVIDIA A100):

Configuration	Time per step (s)
$N$8 (CIC + $N$9 grid, no deconv.)	0.4
$f$0 (CIC + $f$1 grid, deconv.)	0.2
$f$2 (NUFFT, kernel $f$3)	0.3
$f$4	0.1
$f$5	0.02

This suggests near-linear scaling with $f$6 in practice. (List et al., 6 Oct 2025)

6. Field-Level Inference: Direct Application

DISCO-DJ II enables fully differentiable field-level Bayesian inference. In a prototypical application, a synthetic Gadget-4 $f$7 matter field (box $f$8 Mpc/$f$9) is degraded by Gaussian white noise. The latent variables are the initial Gaussian random field $N$0 (for $N$1) and the cosmological amplitude $N$2.

The likelihood in (Fourier) data space, up to $N$3Mpc, is

$N$4

A forward model with $N$5 particles and BullFrog (16 steps) yields a gradient evaluation in less than $N$6 ms (A100). Posterior sampling via Hamiltonian Monte Carlo (BlackJAX) successfully recovers $N$7 (input 0.8102, cosmic-variance shift $N$80.8062), and reconstructs large-scale modes of the initial and final density fields with high accuracy (List et al., 6 Oct 2025).

7. Ecosystem Integration and Future Directions

DISCO-DJ II is designed to function in tandem with a differentiable Einstein–Boltzmann solver, forming a pipeline that is differentiable from primordial initial conditions and cosmological parameters through to the non-linear matter field at late times. Its technical backbone—JAX and the adjoint PM methodology—ensures compatibility with modern scientific machine learning platforms, simulation-based inference, and rapid prototyping workflows (List et al., 6 Oct 2025, Li et al., 2022, Li et al., 2022).

Further development directions include nonparametric time-stepping, coupling to semi-analytic or deep-learning subgrid models, distributed multi-GPU capabilities, and incorporation of additional physics such as short-range (PP) forces and baryonic effects. A plausible implication is that such pipelines will become standard tools for next-generation, inference-driven cosmological survey analyses.