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

TL;DR

Cosmo-FOLD is a novel probabilistic diffusion model that rapidly generates and upscales large, accurate 3D cosmological fields from limited simulation data, significantly reducing computational costs.

Contribution

Introduces Cosmo-FOLD, a new overlap latent diffusion technique for fast, accurate generation and upscaling of cosmological fields conditioned on input data.

Findings

Reproduces large-scale structures and power spectra within 10% accuracy.

Upscales small-volume simulations to large volumes with no fine-tuning.

Faithfully reproduces bispectrum beyond two-point statistics.

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.