CosmoUiT: 3D 21-cm Emulation Framework

• CosmoUiT is a neural field-level emulator that accurately simulates 3D 21-cm intensity maps from Epoch of Reionization conditions using a hybrid Transformer–UNet architecture.
• It conditions on cosmological initial and reionization parameters to capture global dependencies via self-attention and resolve local features through convolutional networks.
• Its performance is validated with metrics like voxel-wise error, SSIM, and power spectrum analysis, ensuring high fidelity and efficiency for parameter inference.

CosmoUiT is a neural field-level emulator designed for rapid and accurate generation of three-dimensional (3D) 21-cm intensity maps from the Epoch of Reionization (EoR), conditioned on cosmological initial conditions and astrophysical reionization parameters. Its architecture combines a Vision Transformer with a UNet, leveraging the advantages of multi-head self-attention for global context modeling and convolutional networks for precise local feature reconstruction. The primary objective of CosmoUiT is to emulate the complex, highly non-Gaussian 3D 21-cm signal with high fidelity across a wide range of spatial scales while allowing efficient parameter inference relevant to upcoming EoR surveys (Posture et al., 1 Oct 2025).

1. Architecture: Vision Transformer–UNet Hybrid

CosmoUiT fuses a 3D Vision Transformer block with a UNet, both operating within a unified neural architecture to optimally handle the statistical and morphological complexity of the EoR 21-cm signal. The processing pipeline proceeds as follows:

• Tokenization: The input 3D physical fields—dark matter density and halo density—are partitioned into non-overlapping cubic patches (subcubes), each of which is flattened and linearly projected into an embedding space.
• Parameter Conditioning: Reionization parameters (e.g., minimum halo mass $M_{h, \min}$, ionizing efficiency $N_{ion}$, and the mean free path $R_{mfp}$) are simultaneously embedded and concatenated with the spatial tokens to form an augmented token sequence.
• Multi-Head Self-Attention: The concatenated tokens are processed through several transformer encoder layers, implementing for each head:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{Q K^{T}}{\sqrt{D_{H}}}\right) V,$$

where $D_{H}$ is the attention head feature dimension. This allows joint modeling of long-range spatial dependencies and parameter interactions.

• Reshaping and UNet Processing: After the transformer block, only the field tokens are retained and spatially reconstructed into downsampled 3D fields. These are then input to a UNet, whose convolutional encoder extracts multiscale features, and whose decoder restores full spatial resolution via transpose convolutions and skip connections. The reionization parameters are reintroduced at the UNet bottleneck as additional feature maps, enforcing parameter conditionality at multiple stages.

This hybrid design allows CosmoUiT to capture non-local correlations in the evolving reionization morphology (transformer), while also resolving small-scale features such as ionized bubble boundaries (UNet).

2. Inputs, Outputs, and Data Flow

Inputs:

• 3D dark matter density field, discretized on a mesh.
• 3D halo density field, generated with a Friends-of-Friends algorithm.
• Vector of reionization parameters: minimum halo mass $M_{h, \min}$, ionizing efficiency $N_{ion}$, mean free path $R_{mfp}$.

Outputs:

• Full 3D cube of the 21-cm brightness temperature $\delta T_b$, derived via the neutral hydrogen fraction $x_{HI}$ and baryonic density.

Data Flow:

1. Input fields and parameters are embedded and processed through the transformer encoder.
2. Recombined embeddings are passed to the UNet for hierarchical spatial reconstruction.
3. The output 3D field provides a physical realization of $\delta T_b$ at the requested parameter values.

This data flow permits conditional emulation—emulated signals adjust according to arbitrary user-specified reionization parameter sets, enabling fast parameter sweeps.

3. Training, Conditioning, and Parameter Generalization

CosmoUiT is trained using pairs of input fields and reionization parameter sets, with targets provided by full semi-numerical EoR simulations (e.g., ReionYuga). The conditioning mechanism operates as follows:

• At transformer input, the projected parameter tokens are concatenated with spatial tokens, allowing the attention mechanism to modulate feature extraction based on the physical modeling context.
• At the UNet bottleneck, the parameters are linearly projected and concatenated with the encoder feature maps, reinforcing their effect on the decoded spatial structures.

This dual-stage conditioning ensures that the emulator not only reproduces the map morphology for a particular input field, but also responds accurately to changes in physical parameters, yielding a smooth, interpretable parameter dependence in output statistics.

4. Performance Metrics and Comparative Analysis

Model performance is validated by direct comparison to reference simulations, using both voxel-wise and field-level summary statistics:

• Voxel-wise error: Mean Squared Error (MSE)

$$\text{MSE} = \frac{1}{n}\sum_{i} (y_i - \hat{y}_i)^2$$

and Coefficient of Determination ($R^2$).

• Structural Similarity Index Measure (SSIM) for perceptual image similarity:

$$\text{SSIM}(x, y) = \frac{(2\mu_x \mu_y + C_1)(2\sigma_{xy} + C_2)}{(\mu_x^2 + \mu_y^2 + C_1)(\sigma_x^2 + \sigma_y^2 + C_2)}$$

• Summary statistics:
  • Bubble Size Distribution (BSD), evaluated via the mean free path method, capturing the size distribution of ionized regions.
  • Power Spectrum $\Delta^2(k)$, with

  $$\delta(\mathbf{k}) = \int x_{HI}(\mathbf{r}) e^{-2\pi i \mathbf{k}\cdot\mathbf{r}} d\mathbf{r}$$

  and

  $$\Delta^2(k) = \frac{k^3}{2\pi^2} P(k), \quad \langle\delta(\mathbf{k}) \delta(\mathbf{k'})\rangle = (2\pi)^3 \delta^D(\mathbf{k}-\mathbf{k'}) P(k)$$

• Emulation accuracy is demonstrated both at the level of global power spectra and higher-order bubble/morphological metrics, with some “fuzzy boundary” smoothing observed at sharp ionization fronts.

The emulator is further tested for its ability to generalize to unseen initial conditions, establishing its utility for ensemble analyses.

5. Emulation Applications in EoR Inference and Survey Analysis

CosmoUiT enables new capabilities in cosmological inference:

• Accelerated mock generation: Emulate statistically consistent 21-cm signal cubes for arbitrary parameters orders of magnitude faster than running full radiative transfer or semi-numerical simulations.
• Bayesian parameter estimation: Integrate CosmoUiT into inference pipelines for direct likelihood evaluation, thereby enabling Bayesian reconstruction of reionization parameters (e.g., via MCMC) from 21-cm observations, with support for fast sampling over realizations and parameter space.
• Survey design and strategy: Generate statistically representative mock maps conditioned on survey-specific science requirements (e.g., for SKAO), facilitating theoretical error forecasts and sensitivity analysis.
• Physical insight: By producing full 3D fields, CosmoUiT captures complex non-Gaussian features—such as bubble topologies and long-range ionization patterns—not accessible to summary-statistic emulators.

Its parameter-conditional formulation makes CosmoUiT suitable for field-level inference tasks, a critical requirement for extracting maximal cosmological and astrophysical information from next-generation 21-cm tomographic data.

6. Mathematical Underpinnings and Algorithmic Summary

The key mathematical components are:

• Self-attention mechanism: For an input token sequence $X_D$ and parameter sequence $X_{para}$, define:

$Q = X_D W^Q,\quad K = X_D W^K,\quad V = X_D W^V$

and the multi-head attention applies

$$\text{Attention}(Q, K, V) = \text{softmax}(Q K^T / \sqrt{D_H}) V$$

across all heads, with appropriate concatenation of results.

• Loss functions: MSE and SSIM described above, together with power spectrum and BSD loss (when used in auxiliary objectives).
• Parameter conditioning: Achieved via concatenation (transformer) and projection+concatenation (UNet bottleneck).

A schematic representation:

Input Fields (DM, halo) + Parameters --> Patch Embeddings & Parameter Tokens | Transformer (multi-head attention, param conditioning) v Reconstructed spatial embedding (field tokens only) | UNet encoder (convolutions) | <parameter projection at bottleneck> v UNet decoder (transpose convolutions + skip connections) Output 3D cube (x_HI / δT_b)

This structure systematically combines the strengths of attention-based global modeling with convolutional local detail, synergized by explicit conditioning on astrophysical parameters.

7. Impact, Scope, and Comparative Context

CosmoUiT advances the state-of-the-art in EoR field-level emulation by:

• Addressing the highly non-Gaussian, multi-scale nature of the 21-cm field, which is inadequately modeled by either pure convolutional or global-statistic emulators.
• Enabling interactive, parameter-dependent exploration of reionization scenarios in full 3D, rather than restricting analysis to summary statistics.
• Providing a computational framework for rigorous, scalable Bayesian parameter inference in the context of upcoming large-volume 21-cm surveys, with direct applicability to SKAO, HERA, LOFAR, and beyond.

Its hybrid transformer–UNet design, conditioning strategy, and demonstrated accuracy set a new benchmark for EoR signal emulation, with further implications for cosmic dawn studies and related cosmological machine learning applications (Posture et al., 1 Oct 2025).