The goal is to reconstruct the initial Fourier-space white-noise field of a cosmological simulation whose evolved late-time density field matches a noisy target observation. Disco-DJ provides the differentiable gravitational evolution model in an expanding universe.
In this demo, the target contains a deliberately embedded tesseract-shaped structure in a noisy cosmological density field, so the reconstruction progress is easy to inspect visually. The optimization starts from a random white-noise
field defining the initial density perturbations on scales of hundreds of millions of comoving light years.
The notebook demonstrates:
Wrapping a JAX-based large-scale structure formation simulator from Disco-DJ inside a Tesseract container
Calling the differentiable Tesseract model from Python via tesseract-jax
Optimizing high-dimensional initial conditions with SciPy/L-BFGS
Maximum a posteriori (MAP) reconstruction using a Gaussian likelihood, Gaussian white-noise prior, Fourier-space preconditioning, and a coarse-to-fine Fourier-space annealing schedule
After optimization, the same forward simulation produces a late-time density field that closely matches the target:
Notebook link:
The included demo uses resolution 64^3; a laptop GPU should be enough. The notebook also includes saved outputs, so it should be possible to inspect the workflow without rerunning the optimization.
This is meant as a compact example of composable differentiable physics workflows: Disco-DJ supplies the cosmological forward model, Tesseract packages it as a differentiable container, and gradient-based optimization reconstructs the
initial conditions directly from the final 3D field.
This demo partly builds on the related Tesseract optimization example “Gradient-based optimization of fluid flows” in this forum.
Thank you for sharing Disco-DJ with us, and for giving Tesseracts a spin! What drew you to building with Tesseracts? Either way, a couple of things that stood out to me were
1 - You used .val_and_grad on an apply_tesseract call to have tesseract_jax work with jax to get the gradient rather than, e.g. using the Jacobian and wiring in the scalar objective into the tesseract. This def makes for a more composable workflow and it’s neat to see that it landed with you
2 - The multi-scale optimization. I have never seen this approach before of annealing from large to small scales to reproduce the multi-scale features from the desired observable. Probably going to use that in the future myself! Do you happen to have a citation you can share for that?
Do you foresee using Tesseracts for your research down the line? For performing HMC calculations on a cluster perhaps?
