Cosmo-FOLD: Fast generation and upscaling of field-level cosmological maps with overlap latent diffusion

TL;DR

Cosmo-FOLD addresses the computational bottleneck of hydrodynamical cosmological simulations by training a diffusion-based generative model on small volumes and upscaling to full cosmological boxes. It introduces Field Overlap Latent Diffusion, a sliding-window diffusion scheme that preserves periodic boundary conditions and reduces edge artifacts without fixed overlap prescriptions. The method reproduces dark matter and gas-temperature fields with high fidelity, achieving within about 10% accuracy for the matter power spectrum up to $k\le 5\,h\mathrm{Mpc}^{-1}$ and recovering higher-order statistics such as the bispectrum when explicit positional encodings are used; it also demonstrates transferability from CAMELS to the larger TNG300-2 volume. This fast, scalable, field-level generative capability enables simulation-based inference and mock survey generation at cosmological scales, with broad implications for data-model comparisons and parameter estimation.

Abstract

We demonstrate the capabilities of probabilistic diffusion models to reduce dramatically the computational cost of expensive hydrodynamical simulations to study the relationship between observable baryonic cosmological probes and dark matter at field level and well into the non-linear regime. We introduce a novel technique, Cosmo-FOLD (Cosmological Fields via Overlap Latent Diffusion) to rapidly generate accurate and arbitrarily large cosmological and astrophysical 3-dimensional fields, conditioned on a given input field. We are able to generate TNG300-2 dark matter density and gas temperature fields from a model trained only on ~1% of the volume (a process we refer to as `upscaling'), reproducing both large scale coherent dark matter filaments and power spectra to within 10% for wavenumbers k <= 5 h Mpc^-1. These results are obtained within a small fraction of the original simulation cost and produced on a single GPU. Beyond one and two points statistics, the bispectrum is also faithfully reproduced through the inclusion of positional encodings. Finally, we demonstrate Cosmo-FOLD's generalisation capabilities by upscaling a CAMELS volume of 25 (Mpc h^-1)^3 to a full TNG300-2 volume of 205 (Mpc h^-1)^3$ with no fine-tuning. Cosmo-FOLD opens the door to full field-level simulation-based inference on cosmological scale.

Figures (10)

• Figure 1: Architecture of the denoising network of Cosmo-FOLD: a 3D U-net with residual network blocks and a self-attention bottleneck. The inputs to the network are the conditioning field $\xi$, the noised target $\eta^t$ at time step $t$, the sinusoidal time step embedding $\boldsymbol{\tau}^t$ corresponding to the diffusion time step $t$, and the optional positional embedding $[\mathbf{p}_x, \mathbf{p}_y, \mathbf{p}_z]$ for each of the coordinate axes. The output $\eta^{t-1}$ is the target field denoised by one extra step.
• Figure 2: Illustration of the Cosmo-FOLD overlap latent diffusion method trained on CAMELS to generated a gas temperature map conditioned on dark matter: example generated gas temperature cube at $t=310$ (left panel), with its sub-volume grid shown in green solid lines. At a later diffusion timestep, $t=230$ (middle panel), the gas field has been denoised, with solid lines showing the current grid and dashed lines the grid location at $t=310$. The process continues until $t=0$ is reached, with the right panel showing another snapshot at $t=140$.
• Figure 3: Generated gas temperature map (top right), conditioned on the $\boldsymbol{\rho}_\text{DM}$ field (top left) for the CAMELS simulation, with the target box size equal to the training box size. The true (unseen) temperature map is shown in the top middle panel. For comparison, the bottom left panel shows the output obtained by simply tiling independent sub-boxes, while the bottom middle panel is the output of our previous LODI method. The histograms of the distribution of pixel values are shown in the bottom right, averaged over 5 realizations of the diffusion models outputs.
• Figure 4: Upscaling to larger volumes via the Cosmo-FOLD method, applied to generate the target gas temperature field $\textbf{T}_g$ of TNG300-2 (top right), conditioned on the true $\boldsymbol{\rho}_\text{DM}$ (top left) or to the $\boldsymbol{\rho}_\text{DM}$ field (middle left) generated from the stellar density field (bottom left). Orange arrows depict upscaling, i.e., large volume generation after training on a much smaller volume. The dark blue arrows show the generation direction, pointing from the conditional field to the generated field.
• Figure 5: Voxel statistics, comparing the quality of the generated $\textbf{T}_g$ fields of TNG300-2 simulations using the Cosmo-FOLD pipeline. Left: voxel distribution of the gas temperature field, generated using both the ground–truth $\boldsymbol{\rho}_\text{DM}$ field and the Cosmo-FOLD-generated $\boldsymbol{\rho}_\text{DM}$ field. Right: $95\%$ contour levels of the joint $\boldsymbol{\rho}_\text{DM}$-$\textbf{T}_g$ distribution for the same two generated cases, compared against the ground truth.