Using Neural Emulators and Hamiltonian Monte Carlo to constrain the Epoch of Reionization's History with the Ly$α$ Forest Power Spectrum

Technical Summary: Using Neural Emulators and Hamiltonian Monte Carlo to Constrain the Epoch of Reionization's History

Problem Statement
The Epoch of Reionization (EoR) represents a critical transition in cosmic history, yet its detailed timeline and inhomogeneous structure remain poorly constrained. While Cosmic Microwave Background (CMB) measurements provide an integrated optical depth constraint, they are insensitive to the specific duration and topology of reionization. Current constraints from high-redshift quasar spectra (e.g., damping wings, dark pixel fractions) offer limited precision and often rely on model-dependent interpretations. The Lyman-$\alpha$ (Ly$\alpha$) forest at $z \sim 5$ offers a promising probe of the intergalactic medium (IGM) thermal and ionization states, but extracting detailed reionization histories requires leveraging the full distribution of Ly$\alpha$ opacities.

A significant barrier to quantitative inference is computational cost. Generating hydrodynamical simulations with sufficient resolution to capture small-scale structures relevant to the EoR requires thousands of GPU-hours per model. Consequently, previous studies have been limited to qualitative comparisons or restricted parameter spaces. Furthermore, robust parameter estimation requires accurate covariance matrices, which are difficult to derive empirically from the sparse high-$z$ quasar samples available. Traditional approaches often assume fixed covariance structures or rely on prohibitively expensive ensembles of simulations to compute model-dependent covariances.

Methodology
This work presents a new inference framework based on JAX that integrates forward-modeled Ly$\alpha$ forest observables with differentiable neural emulators and Hamiltonian Monte Carlo (HMC). The methodology proceeds through the following steps:

1. Simulation Dataset Construction: The authors generated a dataset of 501 low-resolution ($256^3$ voxels, $L_{box} = 20 \, h^{-1}$ Mpc) hydrodynamical simulations using the Nyx code. These simulations were driven by reionization histories generated by AMBER, a semi-numerical code that maps user-defined parameters to ionization fields. The parameter space includes the reionization midpoint ($z_{mid}$), duration ($\Delta z$), asymmetry ($A_z$), maximum heating temperature ($\Delta T_{re}$), and mean flux ($\langle F \rangle$).
2. Forward Modeling: From these simulations, the authors generated 1D Ly$\alpha$ flux skewers, applying observational effects such as spectrograph resolution ($R=30,000$) and noise ($SNR=30$). They computed the mean 1D power spectrum ($\langle P_{Ly\alpha} \rangle$) and model-dependent covariance matrices ($\Sigma_{P_{Ly\alpha}}$) for each model.
3. Neural Emulation: Two independent neural emulators were trained:
   • Power Spectrum Emulator: A fully connected feed-forward network predicting $\log_{10}\langle P_{Ly\alpha} \rangle$ from the input parameters.
   • Covariance Matrix Emulator: To ensure the output remains symmetric and positive-definite, the authors employed Cholesky decomposition ($\Sigma = LL^T$). A neural network predicts the flattened lower-triangular Cholesky factor ($L_{flat}$), with an exponential activation applied to diagonal elements to enforce positivity.
4. Inference Pipeline: The emulators were combined with Hamiltonian Monte Carlo (HMC) via the Numpyro package. The differentiability of the JAX-based emulators allows for the calculation of gradients required by HMC, significantly accelerating parameter space exploration compared to standard Metropolis-Hastings algorithms.

Key Contributions

• Differentiable Covariance Emulation: This work presents the first attempt to emulate Ly$\alpha$ forest covariance matrices using neural networks. By predicting the Cholesky factor, the method ensures physical validity (symmetry and positive-definiteness) while capturing parameter-dependent correlations.
• End-to-End Differentiable Framework: The integration of forward modeling, neural emulation, and HMC within the JAX ecosystem creates a fully differentiable pipeline for cosmological parameter inference.
• Quantitative Validation: The framework is validated against a suite of 100 mock observations, demonstrating the ability to recover true parameters and assessing the statistical validity of the resulting posterior distributions.

Results

• Emulator Accuracy: The power spectrum emulator achieved a median Mean Absolute Percentage Error (MAPE) of 0.15% on the test set, with 95% of errors remaining below 2%. The covariance matrix emulator achieved a median error of 0.42% (relative to the product of diagonal elements), with errors generally remaining below 5% across all scales. Both emulators performed with accuracy well below the expected observational noise levels assumed in the forward modeling.
• Parameter Recovery: When applied to mock observations, the HMC framework successfully recovered the true input parameters. The posterior distributions were visually consistent with the ground truth.
• Statistical Validity: A coverage test performed on 100 mock observations revealed that the posterior distributions are slightly overconfident (the estimated coverage probability lies marginally below the ideal line). The authors attribute this primarily to the assumption of a multivariate Gaussian likelihood function, noting that the Ly$\alpha$ forest is inherently a non-Gaussian random field.
• Data Efficiency: The authors tested the method with a reduced training set of only 100 simulations. While emulation errors increased (power spectrum MAPE rose to ~1.2%, covariance errors increased), the resulting posterior distributions and statistical coverage remained comparable to those derived from the full dataset. This suggests the method is robust even with limited training data, a crucial factor for future high-resolution simulations.

Significance and Limitations
The paper claims that this framework provides a viable path toward quantitatively constraining the history of the EoR using Ly$\alpha$ forest power spectra. The primary significance lies in demonstrating that neural emulators can efficiently replace expensive forward models and that model-dependent covariance matrices can be emulated, enabling faster and more rigorous Bayesian inference.

However, the authors explicitly state that the current results cannot yet be applied to real observational data due to the low resolution of the simulations used, which suppresses power at small scales. Additionally, the reionization model relies on simplified heat injection prescriptions and lacks spatial fluctuations in the ultraviolet background (UVB). The authors conclude that future work must incorporate higher-resolution simulations, more physically motivated thermal treatments, and fluctuating UVB models. They also suggest that the observed slight overconfidence in the posteriors may be addressed by moving toward simulation-based inference (SBI) methods that do not rely on Gaussian likelihood assumptions.