A fast GPU-based bispectrum estimator implemented using JAX.

DISCO-DJ performs large-scale structure (LSS) cosmological simulations with full support for automatic differentiation via the JAX (ascl:2111.002) ecosystem. It integrates a suite of modern cosmological modelling tools, including a fast particle-mesh (PM) N-body solver enhanced by theory-informed time integrators, such as BullFrog and FastPM (ascl:1905.010), and discreteness-suppression techniques (e.g., de-aliasing, sheet-based resampling, higher-order mass assignment kernels). The code also provides an interface to a differentiable Einstein–Boltzmann solver, high-order Lagrangian perturbation theory (nLPT) recursion relations, and semi-classical propagator perturbation theory (PPT). By utilizing JAX’s forward- and reverse-mode automatic differentiation (including memory-efficient adjoint methods), DISCO-DJ enables cosmological inference scenarios where differentiable pipelines from initial conditions to late-time observables are required.

The DensityFieldTools toolset manipulates density fields and measures power spectra and bispectra using a very simple interface. After loading a density field, it computes the power spectrum and the bispectrum for a desired binning. The bispectrum estimator also automatically computes the power spectrum for the chosen binning, to facilitate, for example, shot-noise subtraction. DensityFieldTools also provides a quick way to measure (cross-)power spectra directly from density fields.

URecon reconstructs the initial conditions of N-body simulations from late time (e.g., z=0) density fields. This simple UNET architecture is implemented in TensorFlow and requires Pylians3 (ascl:2403.012) for measuring power spectrum of density fields. The package includes weights trained on Quijote fiducial cosmology simulations.