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

TL;DR

The paper tackles the computational cost of predicting the 21 cm HI signal by learning from hydrodynamical simulations and generating large-volume HI maps from dark matter-only runs. It introduces HALOgen, a 3D U‑Net with attention that maps DM density to four halo-mass channels, and LODI, a conditional variational diffusion model that paints the 21 cm brightness temperature using latent overlap to stitch subvolumes. The latent overlap scheme enables seamless assembly of 32^3 sub-volumes into a $25^3$ volume, reducing boundary discontinuities relative to naive tiling. On the CAMELS/IllustrisTNG CV data, the pipeline reproduces the dimensionless power spectrum within $\le 10\%$ for $k \le 10\,h\,\mathrm{Mpc}^{-1}$ and requires only about 2 minutes of compute per field, demonstrating scalability to arbitrary volumes. This approach supports fast HI mock generation for cross-correlation studies and cosmological parameter inference in next-generation intensity mapping surveys.

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.

Figures (11)

• Figure 1: An overview of our generative pipeline, starting from dark matter particle distribution as produced by an N-body code to the final 21cm intensity map: in the first step, a ResNet with attention bottleneck (HALOgen) is used to predict the dark matter halo mass density in four mass channels; subsequently, a variational diffusion model with latent overlap (LODI) generates the 21 cm brightness temperature map.
• Figure 2: Overview of the HALOgen (Halo Assignment and Generation with U-Net) architecture used for halo assignment from an input DM-only map, employing grouped 3D convolutions (gray blocks) and an attention mechanism at the bottleneck. Upsampling is done via trilinear interpolation to avoid artifacts in the upsampled map. Horizontal lines represent skip connections.
• Figure 3: Illustration of the halo density channels masking prescription. The first row of 4 images shows a 2D slice of the training halo maps $\rho_\text{h}^{(j)}$ for each halo mass channel $j$. The second row shows the corresponding 2D weighted mask maps, $m_w^{j,k}$. In each masking map, the contribution from the halo $j$ is shown in beige, and the contribution from the neighboring channel's halo location is shown in golden.
• Figure 4: The architecture of the denoising model: a 3D U-net with residual network blocks and an attention layer at the bottleneck. The input to the network is the noisy temperature map at time step $t$, the time step embedding and the conditional halo maps. The output is a less noisy version of the temperature map.
• Figure 5: A comparison of different overlapping methods are shown, where the line of discontinuity is vertical and passes through the 2 marked crosses in the images. The left image is the ground truth. The latent overlap method is shown in the left center, compared with two other cases. The center right is the case where for the overlap region, we take the average pixel values of the two boxes and the right image shows the other case when we generate 2 separate boxes and just concatenate them, while keeping the the overlapping pixel region from the first box. The images shown are averaged in the third axis over a distance of about 0.9 $\text{Mpc} h^{-1}$