following up on my previous Tesseract + Disco-DJ optimization example, I put together a new demo showing field-level Bayesian inference through a coupled JAX + Fortran simulation pipeline using Tesseract and Enzyme.
The demo combines two Tesseracts:
a Disco-DJ forward-model Tesseract, which evolves Fourier-space initial conditions into a nonlinear 3D matter density field, and
a Fortran FGPA Tesseract, which turns that density field into synthetic Lyman-alpha forest skewers using the Fluctuating Gunn-Peterson Approximation (FGPA).
The FGPA post-processing code is written in Fortran, and its derivatives are obtained automatically using Enzyme. Through tesseract-jax, these derivatives can then be combined with a JAX-based workflow for gradient-based inference.
The notebook considers a Bayesian reconstruction problem: starting only from noisy 1D Lyman-alpha skewer observations, it samples the posterior over
the high-dimensional initial Fourier white-noise field\epsilon, and
a global FGPA amplitude parameter, \log A.
The inference is done with Hamiltonian Monte Carlo (HMC) using BlackJAX. Before HMC, the notebook runs a staged L-BFGS MAP optimization to find a good initial point.
In this run at resolution 64^3, HMC samples the posterior over the spherically UV-filtered Fourier modes of the initial white-noise field. There are 68,529 active complex representatives, so the parameter space consists of 137,058 real scalar field coordinates, plus the scalar log_A.
The chain recovers the true FGPA parameter within the posterior uncertainty. The posterior samples reproduce the observed skewers well, while the density away from the skewer locations remains much less constrained, as expected.
Hi Florian. This looks really cool, thanks for contributing! Unfortunately I was a bit short on time this week, but looking forward to digging into it properly next week.
Hey @EpicentreOZ, great chatting with you about both of your showcase demos last week. Having this example of how Tesseract can be used to combine Fortran + jax into a single pipeline will definitely be helpful to others, so thanks a lot for taking the time to build and share it. Also really interesting to see the two different demonstrations of Disco-DJ in action, hopefully this inspires others to check it out! Do you think you’ll add other examples to the repo in the future (and/or do you have ideas for demos you’d be excited to see others in the community contribute)?
Yes, I’d be happy to add further examples over time. One direction I recently started looking into is a possible cosmology extension for JAX-Fluids, with two applications in mind:
an effective-fluid model for cosmic structure formation, where the Euler equations are evolved with JAX-Fluids and supplemented by a learned closure for the stress tensor;
longer term, potentially combining Disco-DJ and JAX-Fluids into a modular hydrodynamical cosmological simulation framework.
In both cases, Tesseract could be very useful for keeping the components nicely modular, e.g. for connecting simulation codes, differentiable closures, and inference/optimisation workflows.
Again, this is really nice work! Non-differentiable scientific legacy code can be a major blocker to tackling inference problems at scale, and this is a great example of how to work around these limitations. Using Enzyme-based AD inside a Tesseract to differentiate through the Fortran code and then drop it into a differentiable JAX-based pipeline is exactly where the Tesseract abstraction shines.
I couldn’t resist also poking at the inference side a bit . I was initially skeptical whether 5k HMC transitions with just 10 leapfrog steps would be enough to sample the rather high-dimensional posterior. But looking closer, there’s actually a lot of nice structure. Sampling in the whitened basis diagonalizes the prior via the FFT, so you effectively get an exact prior preconditioner for free. The only “ill-conditioning” left is in the subspace the skewers inform, so everything orthogonal to that stays prior-dominated. So if I understand correctly, the difficult part is sampling this skewer-informed subspace, which is lower-dimensional than the full space (but still high-dimensional).
I did a quick simulation-based calibration check at a lower resolution (16^3). Reassuringly, on the smaller domain everything looks well-calibrated: the rank statistics for \log A and the lowest-k field mode are roughly uniformly distributed, and a deliberately under-tuned sampler fails the test as it should, so it’s not just an insensitive check. The obvious caveat is that it’s far easier to mix on the smaller grid. It would be interested to see what it would look like at full resolution.
I had a lot of fun digging into this. Thanks again for the really cool contribution Florian!
Thanks for going so carefully through my demo and for your thoughtful feedback. Your comments motivated me to do some further analysis at the full 64^3 resolution used in the notebook.
For a single synthetic truth, I computed a mode-wise coverage diagnostic, i.e. what fraction of Fourier-mode coordinates of the true field fall inside the corresponding posterior credible intervals. This is of course not a full SBC test, and the posterior samples are not entirely uncorrelated, but it is still a useful sanity check for this particular mock observation. The resulting posterior samples are somewhat overconfident, see the plot below, so at that resolution your concern regarding mixing is absolutely valid. For the demo I think this is probably good enough, but for a more serious science version I’d run the chain for longer and spend more time on HMC tuning; e.g. the step size, number of leapfrog steps, or possibly switching to NUTS.
Regarding the subspace informed by the skewers, I was curious to see what the posterior uncertainty looks like as a function of Fourier-mode scale and alignment with the skewer direction, using |\!\cos\theta| = |k_z|/|\mathbf{k}|. The qualitative structure looks reasonable: large-scale modes are more strongly constrained by the flux skewers, high-|\mathbf{k}| modes remain more uncertain, and modes aligned with the skewers are somewhat more constrained than orthogonal modes.
For reference, in the units used in the uncertainty plot the prior standard deviation is 1/\sqrt{2} \simeq 0.707 for the stored real/imaginary rFFT components. So the high-|\mathbf{k}| posterior standard deviations being close to 0.7 means those modes are essentially prior-dominated, while the large-scale modes are more strongly constrained.
. I was initially skeptical whether 5k HMC transitions with just 10 leapfrog steps would be enough to sample the rather high-dimensional posterior. But looking closer, there’s actually a lot of nice structure. Sampling in the whitened basis diagonalizes the prior via the FFT, so you effectively get an exact prior preconditioner for free. The only “ill-conditioning” left is in the subspace the skewers inform, so everything orthogonal to that stays prior-dominated. So if I understand correctly, the difficult part is sampling this skewer-informed subspace, which is lower-dimensional than the full space (but still high-dimensional).