Simulation based inference of the ionization history from the 2D 21 cm power spectrum

1. Problem Statement

The 21 cm signal from neutral hydrogen is a primary probe for understanding the Cosmic Dawn (CD) and the Epoch of Reionization (EoR). While upcoming experiments like the Square Kilometre Array (SKA) aim to detect the 2D power spectrum (2DPS) of this signal, extracting astrophysical parameters and the ionization history remains challenging due to:

• Computational Cost: Semi-numerical simulation codes (e.g., 21cmFAST) are too slow to generate the $\sim 10^5$ samples required for traditional Markov Chain Monte Carlo (MCMC) convergence.
• Likelihood Complexity: The likelihood function for 21 cm data is non-Gaussian due to non-linear evolution and mode coupling, making standard Gaussian approximations potentially biased.
• Stochasticity: The 2DPS exhibits significant sample variance (SV), particularly at low $k$-modes, which is difficult to model deterministically.
• Goal: The authors aim to infer three key astrophysical parameters—virial temperature ($T_{vir}$), ionizing efficiency ($\zeta$), and photon mean free path ($R_{mfp}$)—and the full ionization history ($z \sim 5-12$) using Simulation Based Inference (SBI). They specifically investigate whether training SBI on emulated data (generated by a neural network surrogate) improves constraints compared to training on direct simulations, given the need for higher sampling density.

2. Methodology

A. Data Generation and Simulation

• Simulator: The authors used 21cmFAST to generate lightcones between redshifts $z \approx 5-12$.
• Parameters: A Latin Hypercube sampling of 39,000 parameter sets was created for:
  • $\log_{10} T_{vir}$ (3.3 – 5.0)
  • $\zeta$ (10 – 50)
  • $R_{mfp}$ (3 – 80 Mpc)
• Noise Modeling: Realistic SKA noise was added using the Tools21cm package (based on the 2016 SKA antenna configuration, 1000 hrs observation). The noise model accounts for system temperature and sky temperature scaling.
• Data Structure: The lightcones were split into 15 MHz chunks (8 redshift bins) to calculate the 2DPS.

B. Emulator Development

Two distinct emulators were developed:

1. 2DPS Emulator:
   • Input: 3 astrophysical parameters.
   • Output: 2D power spectrum (flattened to 1D array).
   • Architecture: A 7-layer fully connected dense neural network (DNN) with SILU activation functions.
   • Preprocessing: Data was scaled by noise standard deviation to focus on high Signal-to-Noise Ratio (SNR) regions. Crucially, the three lowest $k$-mode pixels (bottom-left corner) were masked to avoid high sample variance where the deterministic emulator fails.
   • Stochasticity Handling: Since the emulator is deterministic, the authors approximated sample variance by adding Gaussian noise (based on McQuinn et al. 2006) to the emulator output during SBI training.
2. Ionization History Emulator:
   • Input: 3 astrophysical parameters.
   • Output: Neutral fraction ($x_{HI}$) across 41 redshift bins.
   • Architecture: A 9-layer DNN with GELU activation functions.

C. Simulation Based Inference (SBI) Pipeline

• Method: Neural Posterior Estimation (NPE) was used to directly model the posterior distribution $P(\theta|x)$ without an explicit likelihood function.
• Architectures: An ensemble of two neural density estimators was used: a Masked Autoregressive Flow (MAF) and a Neural Spline Flow (NSF).
• Comparison: Two inference pipelines were tested:
  1. Simulation-trained: Trained on 21cmFAST simulation data + noise.
  2. Emulator-trained: Trained on data generated by the 2DPS emulator + noise + approximated sample variance.
• Validation: The test set consisted of unseen 21cmFAST simulations (ground truth) to ensure the emulator pipeline was not overfitting to its own training distribution.

3. Key Contributions

• First 2DPS Emulator for SBI: This work presents one of the first attempts to build and utilize a neural emulator specifically for the 2D power spectrum (as opposed to the more common 1DPS) within an SBI framework.
• Handling Sample Variance: The paper explicitly addresses the challenge of sample variance in 2DPS by masking low-$k$ modes and approximating stochasticity via Gaussian noise injection during training.
• End-to-End Ionization History Inference: The pipeline successfully infers the full ionization history ($z=5-12$) by combining parameter inference with a secondary ionization history emulator.
• Benchmarking Emulators vs. Simulations: A rigorous comparison is made between SBI trained on raw simulations versus SBI trained on emulated data, utilizing Simulation Based Calibration (SBC) and TARP tests.

4. Results

A. Parameter Inference

• Both pipelines successfully recovered the virial temperature ($T_{vir}$) and ionizing efficiency ($\zeta$) with high accuracy, correctly identifying the known degeneracy between them.
• The mean free path ($R_{mfp}$) was recovered with wider constraints, consistent with its lower impact on the 2DPS signal.
• Coverage Tests (SBC):
  • The Simulation-trained pipeline showed slight underconfidence (conservative), which is preferred.
  • The Emulator-trained pipeline showed slight overconfidence (overconfident constraints).
  • Conclusion: Increasing sampling density via emulation did not improve prediction accuracy. The emulator's inability to perfectly capture the stochastic nature of sample variance introduced model errors that degraded performance compared to direct simulation training.

B. Ionization History

• Both methods recovered the ionization history with extremely high accuracy (errors $< \pm 0.01$ in neutral fraction).
• The emulator-trained pipeline performed comparably to the simulation-trained pipeline for the ionization history, suggesting the emulator is sufficient for this specific output despite the 2DPS inference limitations.

C. Redshift Dependence

• Late Reionization: Tight constraints were achieved for scenarios where reionization occurs later ($z > 7$).
• Early Reionization: Constraints were significantly weaker due to lower SNR at high redshifts and the lack of signal in bins where the universe is fully ionized.

5. Significance and Conclusions

• Emulator Limitations: The study concludes that for the 2DPS, simple dense neural network emulators may not yet be robust enough to replace full simulations for parameter inference. The stochastic nature of the 2DPS (sample variance) is difficult to capture with deterministic models, leading to overconfident posteriors when using emulators.
• Recommendation: For 2DPS inference, training directly on simulations is currently more robust. However, the ionization history emulator is a valuable tool that allows for flexible, fast inference of the neutral fraction once parameters are constrained.
• Future Directions:
  • More complex architectures (e.g., non-deterministic networks) are needed to better model sample variance.
  • Dimensionality reduction (e.g., PCA) could improve emulator performance.
  • Future work must incorporate foreground residuals and instrumental systematics, which are currently assumed to be perfectly removed.
• Context: This work serves as a proof-of-concept for the SKA Science Data Challenge 3b, highlighting the trade-offs between computational speed (emulators) and statistical fidelity (simulations) in the era of 21 cm cosmology.