Large, fast and accurate HI intensity maps with latent overlap diffusion

TL;DR

This paper presents a machine learning pipeline that efficiently generates 21 cm hydrogen emission maps from dark matter simulations, significantly reducing computational costs while maintaining accuracy.

Contribution

It introduces a novel scalable method combining an attention-based ResUNet and a diffusion model with latent overlap for fast, accurate 21 cm map prediction.

Findings

Predicts 21 cm power spectrum within 10% accuracy for k <= 10 h Mpc^-1

Generates maps in about two minutes after training

Scalable to larger simulation volumes

Abstract

The distribution of 21 cm emission from neutral hydrogen is a powerful cosmological and astrophysical probe, as it traces the underlying dark matter and cold gas distributions throughout cosmic times. However, the prediction of observable signals is hindered by the large computational costs of the required hydrodynamic simulations. We introduce a novel machine learning pipeline that, once trained on a hydrodynamical simulation, is able to generate both halo mass density maps and the three-dimensional 21 cm brightness temperature signal, starting from a dark matter-only simulation. We use an attention-based ResUNet (HALO) to predict dark matter halo maps, which are then processed through a trained conditional variational diffusion model (LODI) to produce 21 cm brightness temperature maps. LODI is trained on smaller sub-volumes that are then seamlessly combined in 512-times larger volume using a new method, called `latent overlap'. We demonstrate that, once trained on 25^3 (Mpc/h)^3 volume simulations, we are able to predict the 21 cm power spectrum on an unseen dark matter map (with the same cosmology) to within 10% for wavenumbers k <= 10 h Mpc^-1, deep inside the non-linear regime, with a computational effort of the order of two minutes. While demonstrated on this specific volume, our approach is designed to be scalable to arbitrarily large simulations.