DISCO-DJ: DIfferentiable Simulations for COsmology - Done with Jax

Florian List (MPI for Astrophysics, Garching)

Sep 24. 2025, 09:20 — 09:40

I will give a brief overview of the DISCO-DJ framework for modelling cosmic large-scale structure, developed by our group at Uni Vienna. DISCO-DJ includes a linear Einstein-Boltzmann code, perturbation theory models, and a fast particle-mesh N-body module. Built with the Jax library, all components are GPU-accelerated, automatically differentiable, and seamlessly integrate with inference and machine learning frameworks. 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 numerical 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. DISCO-DJ provides a self-consistent, highly performant, and fully differentiable pipeline for modelling the large-scale structure of the universe.

 

Further Information
Venue:
ESI Boltzmann Lecture Hall
Recordings:
Recording
Associated Event:
Putting the Cosmic Large-scale Structure on the Map: Theory Meets Numerics (Workshop)
Organizer(s):
Oliver Hahn (U of Vienna)
Cornelius Rampf (RBI, Zagreb)
Cora Uhlemann (Bielefeld U)