--- title: 'Differentiable PM N-body Simulations: DISCO-DJ II' url: https://www.emergentmind.com/topics/differentiable-particle-mesh-n-body-simulations-disco-dj-ii type: topic --- # Differentiable PM N-body Simulations: DISCO-DJ II 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 [2510.05206]. ## 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$ [2510.05206]. 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 [2510.05206, 2211.09815, 2211.09958]. ## 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, \[ \begin{align*} \mathbf{X}_i^{n+1/2} &= \mathbf{X}_i^n + \tau_1\mathbf{V}_i^n, \\ \mathbf{V}_i^{n+1} &= \alpha\mathbf{V}_i^n + \beta \mathbf{A}(\mathbf{X}_i^{n+1/2}), \\ \mathbf{X}_i^{n+1} &= \mathbf{X}_i^{n+1/2} + \tau_2\mathbf{V}_i^{n+1}, \end{align*} \] with step size $\tau = \tau_1 + \tau_2$ and $\mathbf{A}$ representing acceleration. The parameters $(\alpha, \beta)$ specialize the scheme: - **FastPM**: Enforces 1LPT-consistency and symplecticity by matching the Zel’dovich approximation (ZA), with $\alpha = \zeta(a_n)/\zeta(a_{n+1})$, $\beta = (1-\alpha)D_{n+1/2}$, where $\zeta(a) = H a^3 D'(a)$. - **BullFrog**: Implements exact 2LPT consistency at each step: \[ \alpha = \frac{E'(D_{n+1}) - \xi_{n+1/2}}{E'(D_n) - \xi_{n+1/2}},\quad \beta = (1-\alpha) D_{n+1/2},\quad \xi_{n+1/2} = \frac{E_n + E'(D_n)(D_{n+1}-D_n)/2}{D_{n+1/2}}, \] where $D$ is the linear growth factor, $E(D) = -3/7D^2 + \ldots$ [2510.05206]. BullFrog yields per-cent-level accuracy in the power spectrum with as few as $6$–$10$ time steps at $N=512^3$, 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 $W(\mathbf{x})$ are used, possessing analytically known real- and Fourier-space forms; for instance, $\widehat{W}(k) = \mathrm{sinc}^{n}(k/2)$ for the order-$n$ kernel. - **FFT-based PM**: Poisson's equation is solved on the mesh as $\widehat{\phi}(\mathbf{k}) = -\widehat{\rho}(\mathbf{k})/(k^2)$, and the acceleration as $\widehat{\mathbf{g}}(\mathbf{k}) = -i\mathbf{k}\widehat{\phi}(\mathbf{k})$. - **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 ($\sim$1.25$\times$), 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 $|\widehat{W}(k)|^2$), and advanced resampling to suppress artifacts [2510.05206]. ## 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 $O(N_\text{steps})$ 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 $(\overline{\mathbf{X}},\overline{\mathbf{V}})$ backward in time using the same integrator as the forward pass but requires only $O(1)$ 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 [2211.09815, 2211.09958]. ## 5. Numerical Accuracy and Performance Benchmarks Default settings for high-fidelity cosmological inference are as follows: - $N = 512^3$ particles, box $L=500\,\mathrm{Mpc}/h$, initial conditions at $z_\text{ini}=50$ from 2LPT; PM grid $N_g=2N$ per dimension; CIC mass assignment, no deconvolution; BullFrog integrator, uniform stepping in $D$ [2510.05206]. - With BullFrog: - $6$ time steps: per-cent accuracy in $P(k)$ for $k \approx 0.2\,h/$Mpc. - $10$ steps: $<1\%$ error up to $k \approx 0.3\,h/$Mpc. - $25$ steps: $<1\%$ error at $k \approx 0.4\,h/$Mpc. Timing benchmarks (single NVIDIA A100): | Configuration | Time per step (s) | |------------------------------------------------|-------------------| | $N=512^3$ (CIC + $2N$ grid, no deconv.) | 0.4 | | $N=512^3$ (CIC + $N$ grid, deconv.) | 0.2 | | $N=512^3$ (NUFFT, kernel $m=5^3$) | 0.3 | | $N=256^3$ | 0.1 | | $N=128^3$ | 0.02 | *This suggests near-linear scaling with $N$ in practice.* [2510.05206] ## 6. Field-Level Inference: Direct Application DISCO-DJ II enables fully differentiable field-level Bayesian inference. In a prototypical application, a synthetic Gadget-4 $z=0$ matter field (box $L=500$ Mpc/$h$) is degraded by Gaussian white noise. The latent variables are the initial Gaussian random field $w(\mathbf{k})$ (for $|\mathbf{k}|