Cosmo-FOLD: Fast generation and upscaling of field-level cosmological maps with overlap latent diffusion
This enables faster, cheaper field-level cosmological analysis for researchers studying dark matter and baryonic probes, though it builds incrementally on existing diffusion model techniques.
The paper tackles the high computational cost of hydrodynamical cosmological simulations by introducing Cosmo-FOLD, a diffusion model that generates accurate 3D cosmological fields from minimal training data, reproducing power spectra within 10% for specific wavenumbers and scaling volumes by over 8x without fine-tuning.
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.