A neural network emulator for the Lyman-$α$ 1D flux power spectrum
The Lyman-$\alpha$ forest offers a unique avenue for studying the
distribution of matter in the high redshift universe and extracting precise
constraints on the nature of dark matter, neutrino masses, and other
$\Lambda$CDM extensions. However, interpreting this observable requires
accurate modelling of the thermal and ionisation state of the intergalactic
medium, and therefore resorting to computationally expensive hydrodynamical
simulations. In this work, we build a neural network that serves as a surrogate
model for rapid predictions of the one-dimensional \lya flux power spectrum
($P_{\rm 1D}$), thereby making Bayesian inference feasible for this observable.
Our emulation technique is based on modelling $P_{\rm 1D}$ as a function of the
slope and amplitude of the linear matter power spectrum rather than as a
function of cosmological parameters. We show that our emulator achieves
sub-percent precision across the full range of scales ($k_{\parallel }=0.1$ to
4Mpc$^{-1}$) and redshifts ($z=2$ to 4.5) considered, and also for three
$\Lambda$CDM extensions not included in the training set: massive neutrinos,
running of the spectral index, and curvature. Furthermore, we show that it
performs at the 1% level for ionisation and thermal histories not present in
the training set and performs at the percent level when emulating down to
$k_{\parallel}$=8Mpc$^{-1}$. These results affirm the efficacy of our emulation
strategy in providing accurate predictions even for cosmologies and
reionisation histories that were not explicitly incorporated during the
training phase, and we expect it to play a critical role in the cosmological
analysis of the DESI survey.
Laura Cabayol-Garcia, Jonás Chaves-Montero, Andreu Font-Ribera, Christian Pedersen
# A neural network emulator for the Lyman-\(\alpha\) forest 1D flux power spectrum

###### Abstract

The Lyman-\(\alpha\) forest offers a unique avenue for studying the distribution of matter in the high redshift universe and extracting precise constraints on the nature of dark matter, neutrino masses, and other \(\Lambda\)CDM extensions. However, interpreting this observable requires accurate modelling of the thermal and ionisation state of the intergalactic medium, and therefore resorting to computationally expensive hydrodynamical simulations. In this work, we build a neural network that serves as a surrogate model for rapid predictions of the one-dimensional Lyman-\(\alpha\) flux power spectrum (\(P_{\rm 1D}\)), thereby making Bayesian inference feasible for this observable. Our emulation technique is based on modelling \(P_{\rm 1D}\) as a function of the slope and amplitude of the linear matter power spectrum rather than as a function of cosmological parameters. We show that our emulator achieves sub-percent precision across the full range of scales (\(k_{\parallel}=0.1\) to \(4\,{\rm Mpc}^{-1}\)) and redshifts (\(z=2\) to \(4.5\)) considered, and also for three \(\Lambda\)CDM extensions not included in the training set: massive neutrinos, running of the spectral index, and curvature. Furthermore, we show that it performs at the \(1\%\) level for ionisation and thermal histories not present in the training set and performs at the percent level when emulating down to \(k_{\parallel}=8\,{\rm Mpc}^{-1}\). These results affirm the efficacy of our emulation strategy in providing accurate predictions even for cosmologies and reionisation histories that were not explicitly incorporated during the training phase, and we expect it to play a critical role in the cosmological analysis of the DESI survey.

keywords: quasars: absorption lines - cosmology: large-scale structure of Universe - methods: statistical

## 1 Introduction

The Lyman-\(\alpha\) forest refers to a series of absorption features in the spectra of high-redshift quasars caused by the scatter of quasar light at 1216 A of neutral hydrogen in the intergalactic medium (IGM, for a review, see McQuinn, 2016). Consequently, the Lyman-\(\alpha\) forest is sensitive to density fluctuations as well as the thermal and ionisation state of the IGM, thereby containing precise cosmological and astrophysical information about redshifts well above the reach of large-scale galaxy surveys.

Cosmological analyses of the Lyman-\(\alpha\) forest rely on either three-dimensional correlations of the Lyman-\(\alpha\) transmission field to measure baryonic acoustic oscillations (BAO, e.g., Busca et al., 2013; Slosar et al., 2013; Delubac et al., 2015; Bautista et al., 2017; de Sainte Agathe et al., 2019; da Mas des Bourboux et al., 2020), or correlations along the line-of-sight of each individual quasar (one-dimensional flux power spectrum, \(P_{\rm 1D}\), e.g., McDonald et al., 2006; Palanque-Delabrouille et al., 2013; Chabanier et al., 2019) to set constraints on the sum of neutrino masses (Palanque-Delabrouille et al., 2015, 2015; Yeche et al., 2017; Palanque-Delabrouille et al., 2020), the nature of dark matter (Baur et al., 2016, 2017; Yeche et al., 2017; Armengaud et al., 2017; Irsic et al., 2017; Palanque-Delabrouille et al., 2020; Rogers & Peiris, 2021), and even the reionisation (Zaldarriaga et al., 2001; Meiksin, 2009; Lee et al., 2015; McQuinn, 2016) and thermal history (Viel & Haehnelt, 2006; Bolton et al., 2008) of the Universe.

Over the last decade, first the Baryon Oscillation Spectroscopic Survey (BOSS; Dawson et al., 2013) and then the Extended Baryon Oscillation Spectroscopic Survey (eBOSS; Dawson et al., 2016) have dramatically increased the number of Lyman-\(\alpha\) forest measurements available, thereby significantly increasing the precision of cosmological and IGM constraints (du Mas des Bourboux et al., 2020). The ongoing Dark Energy Spectroscopic Instrument survey (DESI, DESI Collaboration et al., 2016) will quadruple the number of line-of-sights, which will increase even more the constraining power of Lyman-\(\alpha\) forest analyses (Font-Ribera et al., 2014).

BAO measurements of the Lyman-\(\alpha\) forest only use the correlations on large scales that can be modelled with linear theory. On the other hand, extracting constraints from the \(P_{\rm 1D}\) measurements is quite challenging because they are sensitive to the complex physical processes that affect the distribution of neutral hydrogen in the IGM, and thus a precise interpretation of this observable requires resorting
to time-consuming hydrodynamical simulations (Cen et al., 1994; Miralda-Escude et al., 1996; Meiksin et al., 2001; Lukic et al., 2015; Walther et al., 2021; Chabanier et al., 2023).

Bayesian inference techniques require of the order of \(10^{6}\) evaluations to set robust constraints on cosmological parameters; consequently, a traditional analysis would require running the same number of hydrodynamical simulations, each of these taking more than \(\sim 10^{5}\) CPU hours, making Bayesian inference unfeasible. One solution to this problem is constructing fast surrogate models that interpolate predictions to regions of the parameter space not sampled by simulations (for the first application of this technique in cosmology, see Heitmann et al., 2006; Habib et al., 2007). As a result, the number of simulations required for Bayesian inference decreases dramatically from millions to dozens or hundreds. We will refer to these models as emulators hereafter.

The two main approaches to building emulators are to use architectures based on a Gaussian process (GP; Sacks et al., 1989; MacKay et al., 1998) or a neural network (NN; McCulloch and Pitts, 1943). These two techniques are different types of supervised learning algorithms, with the first and second typically considered as "non-parametric" and "over-parameterized", respectively. In general, GPs require less training data than NNs and produce more robust predictions, but this comes at the cost of presenting runtimes that scale with the cube of the number of data points instead of linearly like NNs. As a result, GPs and NNs are more appropriate for small and large datasets, respectively. These two architectures have been used for multiple applications, e.g., GPs for emulating the matter power spectrum (e.g., Heitmann et al., 2009; Lawrence et al., 2010; Heitmann et al., 2016; Lawrence et al., 2017), the halo mass function (e.g., Bocquet et al., 2020), and the one-dimensional Lyman-\(\sigma\) flux power spectrum (e.g., Bird et al., 2019; Rogers et al., 2019; Walther et al., 2019; Pedersen et al., 2021; Takhtaganov et al., 2021; Rogers and Peiris, 2021) or NNs for also emulating the matter power spectrum (e.g., Angulo et al., 2021; Arico et al., 2021), galaxy clustering and lensing (e.g., Nishimichi et al., 2019; Chaves-Montero et al., 2023; Contreras et al., 2023, 2023, 2023), and accelerating the predictions of Einstein-Boltzmann Solvers (Gunther et al., 2022; Nygaard et al., 2022).

In this paper, we build the first NN-based emulator for the one-dimensional power spectrum of the Lyman-\(\alpha\) forest. We create it following the approach devised by Pedersen et al. (2021, P21 hereafter), which relies on emulating \(P_{\rm 1D}\) as a function of the amplitude and slope of the linear matter power spectrum on small scales rather than as a function of cosmological parameters. The main advantage of this approach is that it reduces the dimensionality of the problem, and it enables precise predictions for redshifts, cosmological parameters, and \(\Lambda\)CDM extensions not considered in the training set (Pedersen et al., 2021; Pedersen et al., 2023). Our main motivation for building an NN- instead of a GP-based emulator like P21 is that they can handle larger training datasets, like the one that will be needed to accurately interpret \(P_{\rm 1D}\) measurements from DESI. The emulator developed in this paper will be publicly accessible at [https://github.com/igmhub/LaCE](https://github.com/igmhub/LaCE) at the time of publication.

The outline of this paper is as follows. Section 2 presents the hydrodynamical simulations, their post-processing, and the \(P_{\rm 1D}\) parametrisation. In SS3 and SS4, we present the neural-network emulator developed in this paper and characterise its training procedure and training sample. Section 5 contains the main results obtained with the emulator. This includes testing the training simulations, validating the emulator at arbitrary redshifts, and testing it on simulations with different cosmology and astrophysics. Section 6 presents an extended version of the emulator to smaller scales. Finally, SS7 concludes the paper.

## 2 Methods

In this section, we describe the simulations from which we extract \(P_{\rm 1D}\) measurements for calibrating and testing our emulator in SS2.1, how we extract these measurements in SS2.2, and the input parameters for our emulator in SS2.3.

### Simulations

We train our emulator using \(P_{\rm 1D}\) measurements from a suite of 60 flat \(\Lambda\)CDM cosmological hydrodynamical simulations described in detail in P21; throughout this work, we refer to these simulations as training. The simulations were ran employing mp-gadget1(Feng et al., 2018; Bird et al., 2019), a massively scalable version of the cosmological structure formation code gadget-3 (last described in Springel, 2005). Each simulation tracks the evolution of \(768^{3}\) dark matter and baryon particles from \(z=99\) to \(z=2\) inside a simulation box of \(L=67.5\) Mpc on a side and generates 11 output snapshots uniformly spaced in redshift between \(z=4.5\) and \(z=2\). To increase computational efficiency, star formation is included using a simplified prescription that turns regions of baryon overdensity \(\Delta_{\rm b}>1000\) and temperature \(T<10^{5}\) K into collisionless stars, which is justified by the negligible contribution of high-density regions to the Ly\(\alpha\) forest (e.g., Viel et al., 2004). Also for efficiency purposes, the simulations use a spatially uniform ultraviolet background implementation from Haardt and Madau (2012) and do not include AGN feedback; the first and second approximations may lead to up to \(\sim 10\%\) errors on \(P_{\rm 1D}\) predictions especially at high redshift (\(z=4\), e.g., Pontzen, 2014; Gontcho A. Gontcho et al., 2014; Suarez and Pontzen, 2017) and low redshift (\(z\simeq 2\), e.g., Chabanier et al., 2020), respectively. Nevertheless, the impact of AGN feedback on \(P_{\rm 1D}\) measurements is cosmology-independent and it can be accounted for at the post-processing stage (Chabanier et al., 2020).

Footnote 1: [https://github.com/MP-Gadget/MP-Gadget/](https://github.com/MP-Gadget/MP-Gadget/)

The training simulations adopt 30 different sets of cosmological and astrophysical parameters selected according to a Latin hypercube design so the space of interest is sampled efficiently (McKay et al., 1979). Two realisations were run for each combination using

Figure 1: Redshift evolution of the gas temperature at mean density. The gray lines show the results for training simulations, while the red line does so for the reimisation simulation. The thermal histories of the training simulations, which are used to train the \(P_{\rm 1D}\) emulator, are significantly different from that of the reimisation simulation, which is used for testing.
an initial mode amplitude fixed to the ensemble mean and opposite Fourier phases (Angulo and Pontzen, 2016; Pontzen et al., 2016). These initial conditions, commonly known as "fixed-and-paired", significantly reduce cosmic variance in the Ly\(\alpha\) forest power spectrum (Villaescusa-Navarro et al., 2018; Anderson et al., 2019); throughout the remainder of this work, we refer to measurements from the simulations of a pair as from different phases. The impact of cosmic variance on training simulations is nevertheless considerable due to their limited size. To separate this source of uncertainty from others, all simulation pairs were run using the same distribution of Fourier phases.

Motivated by the emulation strategy, the simulations explore different values of the amplitude and slope of the linear power spectrum,

\[\Delta_{\rm p}^{2}(z) = k^{3}P_{\rm lin}(k_{\rm p},z), \tag{1}\] \[n_{\rm p}(z) = (\mathrm{d}\log P_{\rm lin}/\mathrm{d}\log k)\mid_{k=k_{\rm p}}, \tag{2}\]

where \(k_{\rm p}\) is the pivot scale at which these are computed and \(P_{\rm lin}\) is the linear power spectrum of cold dark matter and baryons2. Specifically, the simulations use values within the ranges \(\Delta_{\rm p}^{2}(z=z_{\star})\in[0.25,\,0.45]\) and \(n_{\rm p}(z=z_{\star})\in[-2.35,\,-2.25]\), which are defined at \(z_{\star}=3\) and \(k_{\rm p}=0.7\,\mathrm{Mpc}^{-1}\) because this redshift and scale are approximately at the centre of the ranges of interest for DESI Lyman-\(\alpha\) studies (Ravoux et al. in prep, Karacayli et al. in prep.). As for the other cosmological parameters, the simulations use the same value for the Hubble parameter (\(H_{0}=67\,\mathrm{km}\,\mathrm{s}^{-1}\mathrm{Mpc}^{-1}\)), physical cold dark matter density (\(\omega_{\rm c}\equiv\Omega_{\rm c}h^{2}=0.12\)), and physical baryon density (\(\omega_{\rm b}\equiv\Omega_{\rm b}h^{2}=0.022\)), where \(h=0.67\) is the dimensionless Hubble parameter. Note that P21 and Pedersen et al. (2023) showed that our emulation strategy produces precise results for simulations with cosmological parameters outside of the training set and \(\Lambda\)CDM extensions; we test this further in SS5.5.

Footnote 2: Note that \(P_{\rm lin}\) does not include the contribution of neutrinos for cosmologies with massive neutrinos.

The training simulations consider three astrophysical parameters to account for uncertainties in the reionisation and thermal history of the Universe. We consider as a fiducial model the histories from Haardt and Madau (2012), and when vary the value of the previous 3 parameters to perturb these histories following prescriptions from Onorbe et al. (2017). Specifically, the simulations explore \(z_{\rm H}\in[5.5,\,15]\), which indicates the midpoint of hydrogen reionisation, while holding fixed the redshift of the second helium ionisation to \(z_{\rm He\,H}=3.5\). In addition, these use \(H_{\rm A}\in[0.5,\,1.5]\) and \(H_{\rm B}\in[0.5,\,1.5]\), which account for the uncertain effect of helium reionisation on the IGM temperature by rescaling the He\(\,\)ii photo-heating rate, \(\epsilon_{0}\), such that \(\epsilon=H_{\rm A}\Delta_{\rm b}^{H_{\rm B}}\epsilon_{0}\). In this way, the thermal state of the IGM is correctly coupled to the gas pressure.

In addition, we use seven pairs of simulations with cosmological and astrophysical parameters not considered in the training simulations for evaluating different aspects of the emulation strategy3. The first is the central simulation, which uses parameters at the centre of the training parameter space and serves to evaluate the performance of the emulator in optimal conditions. The second is the seed simulation, with the same parameters as the central simulation but different initial conditions, which serves to characterise the impact of cosmic variance on the results.

Footnote 3: Some of these simulations were already described in P21 and Pedersen et al. (2023) but named differently. The simulations growth and neutrinos were referred to as \(h\) and \(\nu\) in P21 and the seed simulation as _diff seed_ in Pedersen et al. (2023).

We also use the growth, neutrinos, running, and curved simulation pairs, which present the same amplitude and slope of the matter power spectrum at \(z=3\), physical CDM and baryonic densities, and astrophysical parameters as the central simulation, but the growth simulation uses 9% larger Hubble constant and 18% smaller \(\Omega_{\rm M}\), the neutrinos simulation includes massive neutrinos (\(\sum m_{\nu}=0.3\) eV) implemented using the linear response approximation Ali-Haimoud and Bird (2013), the running simulation uses a non-zero running of the primordial power spectrum slope (\(\mathrm{d}\,n_{\rm s}/\mathrm{d}\log k=0.015\)), and the curved simulation considers an open universe (\(\Omega_{k}=0.03\)). These simulations serve to test the precision of the emulation strategy for cosmologies not included in the training set.

Furthermore, we consider the reionisation simulation, with the same cosmological parameters as the central simulation but implementing an ionisation history from Puchuin et al. (2019). The main difference between this and the training simulations is that the helium ionisation history of the first peaks at a later time, which translates into different IGM thermal histories. In Fig. 1, the grey lines display the thermal histories of all training simulations, while the red line does so for the reionisation simulation. As expected, we can see that the thermal history of the reionisation simulation peaks at a later time relative to those of the training simulations due to the also lower \(z_{\rm He\,H}\) for the first. This simulation serves to test the performance of the \(P_{\rm 1D}\) emulator for ionisation and thermal histories different from those used in the training set.

### Post-processing

We extract \(P_{\rm 1D}\) measurements from the simulations described in SS2.1 as follows. For each simulation, we first consider one of the simulation axes as the line of sight, and then we displace particles from real to redshift space along this axis. We continue by computing the transmitted flux fraction along \(768^{2}\) uniformly-distributed line of sights along this axis using FSFE4(Bird, 2017); these lines of sights are commonly known as skewers. The line-of-sight resolution of the skewers is set to 0.05 Mpc, which is enough to resolve the thermal broadening scale. Then, we compute the Fourier transform of the transmitted flux fraction for each skewer, and we estimate \(P_{\rm 1D}\) by taking the average of the Fourier transform of all skewers. Finally, we iterate over the two remainder simulation axes. By doing so, we sample different directions of the velocity field, extracting further information from the simulations.

Footnote 4: [https://github.com/sbird/fake_spectra](https://github.com/sbird/fake_spectra)

We repeat the previous procedure for each simulation and snapshot, ending up with 30 (cosmologies) \(\times\) 2 (opposite Fourier phases) \(\times\) 3 (simulation axes) \(\times\) 11 (snapshots) = 1980 \(P_{\rm 1D}\) measurements. Additionally, in post-processing, we vary the mean flux of the snapshots by scaling the effective optical depth of the skewers to 0.90, 0.95, 1.05, and 1.10 times its original value (see Lukic et al., 2015, for more details about this approach), and then we recompute the power spectrum, ending up with 9900 \(P_{\rm 1D}\) measurements in total. As a result, this post-processing presents a few improvements compared to that carried out by P21: three simulation axes instead of one, \(768^{2}\) skewers instead of \(500^{2}\), and mean flux rescalings.

### Emulator parameterisation

We use the amplitude and slope of the linear matter power spectrum at \(k_{\rm p}=0.7\,\mathrm{Mpc}^{-1}\) to capture the cosmological dependence of
measurements. This is justified because, when rescaling the linear matter power spectra of cosmologies within _Planck_ priors so these match the amplitude and slope of the best-fitting _Planck_ solution, the amplitude of the variations is smaller than 1% (see fig. 1 of P21). It is also important to note that the Lyman-\(\alpha\) forest probes cosmic times during which the universe is practically Einstein de-Sitter, and for such universe, the growth rate of velocities (and therefore redshift space distortions) has no cosmological dependence.

We use four parameters to describe the astrophysical dependence of the measurements. The first is the mean transmitted flux fraction \(\bar{F}\), which encodes information about the ionisation state of the gas and is related to the effective optical depth as \(\tau_{\rm eff}=-\log\bar{F}\). The next two inform us about the thermal state of the gas probed by the Lyman-\(\alpha\) forest. The first, \(\gamma\), describes the slope of the temperature-density relation (e.g., Lukic et al., 2015), \(T=T_{0}\Delta_{\rm b}^{-1}\), where \(T_{0}\) is the gas temperature at mean density. The second, \(\sigma_{\rm T}\), captures the thermal broadening of absorption features due to the thermal motion of gas,

\[\sigma_{\rm T}=\sigma_{\rm T,0}\sqrt{\frac{T_{0}[K]}{10^{4}}\frac{1+z}{H(z)}}\, \tag{3}\]

where \(\sigma_{\rm T,0}=9.1\) km s\({}^{-1}\). Finally, the fourth parameter, \(k_{\rm F}\), captures the pressure smoothing scale (Kulkarni et al., 2015). Note that we use \(\sigma_{\rm T}\) and \(k_{\rm F}\) in inverse comoving units (for more details, see P21).

## 3 Emulators

In this section, we describe the two \(P_{\rm 1D}\) emulators that we use throughout the paper: a GP-based emulator (LaCE-GP) in SS3.1 and an NN-based emulator (LaCE-NN) in SS3.2. We provide a brief description of the first, which was already discussed in P21, and we focus on the second, which is the main contribution of this work.

### Gaussian process emulator (LaCE-GP)

The LaCE-GP emulator maps the parameter space \(\vec{\theta}=\{\Delta_{\rm p}^{2},n_{\rm p},\bar{F},\sigma_{\rm T},\gamma,k_{ \rm F}\}\) to the \(P_{\rm 1D}\) measurements (see SS2.3) based on a Python-based implementation using the package Gyr. Specifically, it normalises input parameters \(\vec{\theta}\) so these vary within a unit volume and \(P_{\rm 1D}\) by the median of all \(P_{\rm 1D}\) used. The original version of this emulator (P21) predicts the value of \(P_{\rm 1D}\) for each input scale \(k_{\parallel}\). The emulator was later updated to predict the best-fitting coefficients of a fourth-degree polynomial describing \(P_{\rm 1D}\) measurements because Pedersen et al. (2023) found that the original implementation was significantly affected by cosmic variance.

In this work, we used a modified version of the Pedersen et al. (2023) emulator: we train it using the new post-processing of the LaCE simulations (SS2.2), extend the fitting range from \(k_{\parallel}=3\) to \(k_{\parallel}=4\) Mpc\({}^{-1}\), and predict the coefficients of a fifth-order polynomial (to account for the extended range of scales, see Appendix A). In SS5.3, we provide a detail comparison between the performance of this and the LaCE-NN emulator.

### Neural network emulator (LaCE-NN)

Neural networks are mathematical models composed of interconnected layers of nodes, each of which possesses trainable weights and biases (\(w\), \(b\)). The initial and final layers are commonly referred to as the input and output layers, respectively, and they have the same number of nodes as the dimensions of the input and output data. The intermediate layers are known as hidden layers and their number of nodes is a flexible hyperparameter of the network. Increasing the number of layers enhances the network's complexity, enabling it to handle more challenging tasks. However, networks with a higher number of layers are more susceptible to overfitting, which occurs when an excessively intricate model captures irrelevant patterns in the training data, leading to impaired generalization on new data.

In practice, the input data to the layer \(i\) (\(\vec{x_{i}}\)) is multiplied by the set of trainable weights associated with the nodes in such layer, and the resulting values are passed through a non-linear function

\[\vec{y_{i}}=g(\vec{x_{i}}\cdot\vec{w_{i}}+\vec{b_{i}}), \tag{4}\]

where \(g(\cdot)\) is commonly known as the activation function and \(\vec{y_{i}}\) represents the output of this layer. The training process involves finding the optimal weights and biases that minimise the difference between the network's output and the target "true" values. This discrepancy is quantified using a loss function (\(\mathcal{L}\)) and the optimisation is achieved using an optimisation algorithm, such as gradient descent. This algorithm calculates the gradient of the loss function with respect to the weights and biases and adjusts them in a manner that minimises the loss.

In this paper, we use a mixture-density network (MDN, Bishop, 1994), a type of neural network that can model complex probability distributions. MDNs use a combination of neural network architecture and mixture models to estimate the parameters of a probability distribution from the input data. Each mixture component is associated with the mean and variance of a (typically) Gaussian distribution, and the network optimises such parameters to best describe the probability distribution of the target prediction.

Fig. 2 shows the MDN used in this work, which is written in PyTorch. It maps the six parameters \(\vec{\theta}=[\Delta_{\rm p}^{2},n_{\rm p},\bar{F},\sigma_{T},\gamma,k_{F}]\) described in SS2.3 (red circles) to the probability distribution of six polynomial coefficients best fitting the \(P_{\rm 1D}\), where we assume the polynomial coefficients to follow a Gaussian around the true value. The input layer of the emulator maps the six input parameters to ten hidden variables (6:10) and it is then followed by a hidden space with four hidden layers mapping ten input to 100 output parameters (blue circles). The 100 output parameters of the hidden layers are the input of two independent sets of layers with architecture 100:50:5 (soft-purple circles) mapping the output of the previous hidden space to the mean and standard deviation of the Gaussian distribution best describing the six polynomial coefficients (yellow circles). The MDN is trained with an Adam optimiser (Kingma & Ba, 2015) with an initial learning rate is \(10^{-3}\) for 100 epochs.

The neural network fits the polynomial form

\[P^{\prime}_{\rm 1D}=\sum_{i=0}^{n-1}\alpha_{i}\cdot\left(\log_{10}k_{\parallel} \right)^{i}, \tag{5}\]

where \(n\) is the order of the polynomial and \(P^{\prime}_{\rm 1D}\) is the \(\log_{10}P_{\rm 1D}\) scaled by the median of all \(P_{\rm 1D}\) measurements in the training sample (\(A_{\rm scalings}\)), i.e.

\[P^{\prime}_{\rm 1D}=\log_{10}\left(P_{\rm 1D}/A_{\rm scalings}\right). \tag{6}\]

The network also predicts an uncertainty for each polynomial coefficient (\(\sigma_{i}\) ) that propagates to an uncertainty for \(\log_{10}P^{\prime}_{\rm 1D}\)

\[\sigma_{\log}^{2}\,p_{\rm 1D}=\sum_{i=0}^{n-1}\sigma_{i}^{2}\cdot\left(\log_{10}k_{ \parallel}\right)^{2i}. \tag{7}\]

In App. B, we also explore MonteCarlo methods to obtain a covariance matrix from the distribution of the polynomial coefficients.
To fit the \(P_{\rm 1D}\) in the \(k_{\parallel}\) range \((0,4]\) Mpc\({}^{-1}\), we use a fifth-order polynomial (\(n=5\)), which is justified in App. A. To further back up this decision, Fig. 3 shows the contribution of the optimal polynomial terms once they have been optimised by the emulator up to sixth degree. The zero- and first-order coefficients broadly contribute to the \(P_{\rm 1D}\). Then, second- and third-order contributions compensate for each other and a similar thing happens with order four and five. Finally, the sixth order term is already centred at zero, and therefore its contribution to the \(P_{\rm 1D}\) measurements is smaller.

After each training iteration, we compare the emulator predictions (Eqs. 5,7) with the true values using a log-likelihood loss function \(\mathcal{L}\)

\[\mathcal{L}=\sum_{i}^{N}\left[\sum_{k}^{N_{A}}\left(\frac{\log_{ 10}P_{\rm 1D}^{\prime\ \rm pred}-\log_{10}P_{\rm 1D}^{\prime\ \rm true}}{\sigma_{\log_{10}}P_{\rm 1D}^{ \prime}}\right)^{2}+2\log\sigma_{\log_{10}}P_{\rm 1D}^{\prime}\right]. \frac{1}{N}\,, \tag{8}\]

where \(N\) corresponds to the number of training samples considered in the loss and \(N_{k}\) is the number of wave-number \(k_{\parallel}\) bins. Evaluating the loss function in the \(P_{\rm 1D}\) space rather than the space of polynomial coefficients enables the neural network to learn the importance of each coefficient (e.g. the zero- and first- order polynomial degrees contribute more to the \(P_{\rm 1D}\) than the fourth-order one) and weight the learning accordingly. In Appendix C, we also explore emulating the \(P_{\rm 1D}\) in principal components space, but the polyfit emulator is more accurate and robust.

Figure 3: Contribution of each polynomial term to the \(P_{\rm 1D}\) reconstruction. Each colour corresponds to the histogram of the contribution of each polynomial order computed for all snapshots in the training simulations (§2.1). The zero- and first-order coefficients (\(i=0\) and \(i=1\), blue and purple dotted lines) gently contribute to the \(P_{\rm 1D}\). The second- and third-order coefficients compensate for each other (\(j=2\) and \(j=3\), red and yellow lines). Similarly, the fourth and fifth-order coefficients (\(i=4\), \(i=5\), green and orange dashed lines) also compensate for each other. Finally, the sixth-order contribution (\(i=6\), solid brown) is centred at zero.

Figure 2: Architecture of the mixture density network used in this work. The red circles correspond to the six input parameters (§2.1), and the blue circles and their connections represent a set of shared hidden layers. The symbol within the blue circles represent the activation function, which in our case is a Leaky-ReLU. The output of these hidden layers is the input of two independent sets of hidden layers, each of them mapping to six output parameters. These output parameters parametrise the mean and standard deviation of six Gaussian distributions corresponding to the six polynomial coefficients of the \(P_{\rm 1D}\) fit (Eq.5.)
## 4 Emulator Characterisation

In this section, we optimise the emulator's training strategy (SS4.1) and the training sample by studying whether the precision of the emulator increases after including optical-depth rescalings in the training sample (SS4.2) and exploring different ways of doing data augmentation with the three-axes data (SS4.3).

### Neural network characterisation

In order to optimise the performance of the emulator, we have tested several actions on the network architecture, the training procedure, the input data, and the target \(P_{\rm 1D}\). Figure 4 shows the impact of some of these actions on the emulator's performance. In the \(x\)-axis, the plot indicates the action performed on the network, while the \(y\)-axis displays the percent error in the emulated \(P_{\rm 1D}\). Comparing the error after each modification, we can quantify the impact that each action has on the \(P_{\rm 1D}\). In all cases, the emulator is trained on the training simulation pairs (SS2.1) and tested on central. To reduce the sources of variability between runs, the neural network is initialised with the same set of initial weights in all runs, which are also randomly generated.

We first implemented one modification at a time to test their effect independently of the others, and then tested the impact of cumulative modifications. The starting point is a neural network fitting a polynomial in \(k_{\parallel}\) and predicting \(P^{\prime}_{\rm 1D}\) (Eq.6). We first study the impact of fitting the \(P^{\prime}_{\rm 1D}\) in \(\log_{10}k_{\parallel}\) space (A). This drastically reduces the error by a factor of 7. The second action (B), consists in fitting the \(\log_{10}(P^{\prime}_{\rm 1D})\) instead of \(P^{\prime}_{\rm 1D}\) directly, which clearly benefits the emulator.

Actions C and D modify the activation function and the input's normalisation, respectively. Initially, the input parameter space (\(\theta\)) was normalised with a min-max normalisation

\[\vec{\theta}^{\prime}=\frac{\vec{\theta}-\min(\vec{\theta})}{\max(\vec{ \theta})-\min(\vec{\theta})}\,, \tag{9}\]

which rescales the parameter space to the [0,1] range. Instead, we modify this scaling by also subtracting 0.5, in such a way that our input parameter space is the [-0.5,0.5] range. Since the ReLU activation function sets to zero all negative values, we include the Leaky-ReLU activation function (Xu et al., 2015) to avoid vanishing gradients. In Fig.4, we can see that the Leaky-ReLU alone does not have an impact on the performance (A+C), but it does when it is combined with the prediction in \(\log_{10}k_{\parallel}\) (A+B+C). Similarly, the input parameter shift does (A+D) does not improve the emulator's performance but combined with the other changes there is a small improvement.

Based on these results, the configuration of our fiducial emulator includes the fitting in \(\log_{10}k_{\parallel}\), the fitting of the \(\log_{10}P_{\rm 1D}\), the Leaky-ReLU activation function and the parameter shift (A+B+C+D).

### Optical-depth rescalings

Generating hydrodynamical simulations and extracting the \(P_{\rm 1D}\) is computationally expensive and time consuming. However, once the \(P_{\rm 1D}\) measurement per snapshot has been measured, we can generate several realisations of such \(P_{\rm 1D}\) rescaling the mean-transmitted flux fraction (Lukic et al., 2015; Walther et al., 2021). This is an effective way of increasing the training sample of the emulator. Our simulations include four mean-flux rescalings per snapshot, which augment the training sample from 330 to 1650 training points. Such rescalings correspond to \(\pm\)(5 and 10)% in optical depth, which does not propagate into \(\pm\)(5 and 10)% mean flux. This is a novelty of the LaCE-NN emulator that can benefit from the training sample augmentation with optical-depth rescalings.

Figure 5 compares the percent error in the \(P_{\rm 1D}\) predictions with emulators trained on samples without optical-depth rescalings (red), with 5% rescalings (blue), 10% (green), and combining both 5% and 10% optical-depth rescalings (yellow). For this plot, we have run four independent leave-one-out tests (see SS5.1), on per training set, obtaining the error in the \(P_{\rm 1D}\) emulation of each simulation and snapshot for each of the training samples. The error is then averaged over simulations and scales to have an error estimate per redshift. The emulator is always tested on non-scaled \(P_{\rm 1D}\).

There is a significant improvement in the performance of the emulator when rescalings are incorporated into the training sample. For instance, the inclusion of 5% optical-depth rescalings resulted in a reduction of the percentage error in the emulated \(P_{\rm 1D}\) by a factor of two. The emulator benefited more from the inclusion of 10% optical

Figure 4: Optimisation of the neural network emulator. The \(x\)-axis corresponds to different modifications implemented on the network’s architecture, training data, and training procedure. The \(y\)-axis indicates the percent error in the \(P_{\rm 1D}\) emulation for each of these modifications.

Figure 5: Impact of the optical-depth rescalings on the \(P_{\rm 1D}\) emulation. The plot shows the percent error in the \(P_{\rm 1D}\) emulation for the 30 hypercube simulations pairs averaged over simulations and scales. The solid red and solid blue lines correspond to samples without and with all rescalings, respectively. In turn, the soft dashed lines indicate the samples including 5% (yellow) and 10% (green) optical-depth rescalings. The yellow and green lines are artificially shifted in the \(x\) axis for a better visualisation.
depth rescalings as compared to the 5% ones (green vs yellow). This may be due to the fact that the 10% rescalings increase the parameter space to a greater extent than the 5% rescalings, leading to the new parameter space encompassing the 5% mean-flux rescaling points as well.

Nevertheless, the optical-depth rescalings included in this post-processings mostly populate the parameter space and barely affect the limits of its convex hull. Future emulators aiming to predict universes with different IGM history will also require temperature and mean-flux rescalings increasing the covered parameter space.

### Training sample characterisation

With the new post-processing of the hydrodynamical simulations, we obtain one \(P_{\rm 1D}\) measurement along each of the three axes for every simulation box. Given that our simulations are fixed-and-paired (SS2.1), we have a total of six \(P_{\rm 1D}\) measurements per simulation snapshot. These measurements can be averaged to obtain a single, less noisy \(P_{\rm 1D}\) measurement.

However, in this section, we explore data augmentation techniques by considering different combinations of these six \(P_{\rm 1D}\) measurements. While the most common approach involves training the emulator with the average of the six measurements, we also investigate the possibility of utilising each measurement independently and combining only \(P_{\rm 1D}\) along axes or phases.

Figure 6 explores the emulator's performance by considering different combinations of the six \(P_{\rm 1D}\) measurements. The plot depicts the percent error in \(P_{\rm 1D}\) emulation for various training sets. These errors were estimated using all training simulations and are compared to the most accurate \(P_{\rm 1D}\) measurement available, which results from averaging all axes and phases.

The red line represents the percent error in \(P_{\rm 1D}\) emulation for different training samples. The first red point, labeled "Avg. all," corresponds to a training sample that averages the \(P_{\rm 1D}\) measurements from all axes and phases, resulting in a total of 1650 training points (30 simulations \(\times\) 11 snapshots \(\times\) 5 optical-depth rescalings).

The next red points are "Avg. axes" and "Avg. phases", and reduce the percent error compared to the previous result. In the first case, "Avg. axes", we have averaged the \(P_{\rm 1D}\) from over the three axes, but not over phases. Therefore, this corresponds to two \(P_{\rm 1D}\) measurements per snapshot. Additionally, we also include the \(P_{\rm 1D}\) averaged along axes and phases, which is the most accurate measurement we have. Therefore, this corresponds to 4950 training points (30 simulations \(\times\) 11 snapshots \(\times\) 5 optical-depth rescalings \(\times\) (2+1) \(P_{\rm 1D}\)). In the second case, "Avg. phases", we average over phases, but not over axes. For each snapshot, we have three \(P_{\rm 1D}\) measurements, and as before, we additionally include the average over axes and phases. Therefore, the training sample contains 6600 training points (30 simulations \(\times\) 11 snapshots \(\times\) 5 optical-depth rescalings \(\times\) (3+1) \(P_{\rm 1D}\)).

The following red point, "Avg. axes + Avg. phases" combines the two previous ones, which makes a total of 9900 training points (30 simulations \(\times\) 11 snapshots \(\times\) 5 optical-depth rescalings \(\times\) (3+2+1) \(P_{\rm 1D}\)) and clearly reduces the percent error compared to using only one of the data sets. Finally, the last point, "Avg. axes + Avg. phases + inde" adds up the average over axes, the average over phases and the six \(P_{\rm 1D}\) measurements, which correspond to 19800 training points (30 simulations \(\times\) 11 snapshots \(\times\) 5 optical-depth rescalings \(\times\) (3+2+1+6) \(P_{\rm 1D}\)). In this case, the emulator does not seem to benefit from the additional training points.

Based on these results, the fiducial training sample for our emulator corresponds to "Avg. axes + Avg. phases" and contains 9900 training points.

## 5 Testing emulation strategy

In this section, we evaluate the performance of the proposed emulation approach for the surrogate models discussed in the preceding sections. We specifically focus on assessing the precision of the emulator when applied to training simulations (SS5.1) and redshifts (SS5.2) that were not part of the training sample. In SS5.3, we compare the LaCE-GP and the LaCE-NN emulators. Additionally, we investigate the effect that cosmic variance has on the emulator (SS5.4) and the emulator's performance on cosmologies (SS5.5) and IGM thermal histories (SS5.6) that were not included in the training set.

### Leave-one-out tests

The leave-one-out test is a commonly used validation test to evaluate the accuracy of an emulator. The emulator is trained using the set of training simulations, holding one out as a validation set, which is then evaluated. The leave-one-out test is repeated for each simulation in the training set, determining how well the emulator is able to generalise to unseen simulations and the overall accuracy it can achieve.

Figure 7 presents the leave-one-out test results for training simulations. To generate the plot, we have optimised 30 independent emulators, each one training on 29 training simulation pairs and evaluated on the left-out simulation.

To produce the Fig. 7 we group all snapshots with the same redshift and take the mean and the standard deviation of the percent error in the emulation across simulations. The red line indicates the mean percent error and the red shaded region corresponds to the standard deviation of the percent error. The shaded grey area in the figure indicates the 1% error requirement for the emulator. For most cases, the emulator reaches the \(<\) 1% error requirement at all redshifts, although some measurements at \(z=2\) are slightly over 1% error. This is not unexpected since \(z=2\) lies in the extreme of the Lyman-\(\alpha\) forest detection and therefore the emulator is more likely to degrade.

Figure 6: Characterising the training sample of the emulator. The figure shows the percent error obtained in the \(P_{\rm 1D}\) emulation for different training sample constructed from several possible combinations of \(P_{\rm 1D}\) measurements over the three axes and two paired phases. The horizontal lines correspond to the 1% error requirement (dotted-black line) and the percent error of the one-axis emulator (blue lines).
Finally, the blue-shaded region represents the observational uncertainties from the latest analysis of BOSS/eBOSS data (Chabanier et al., 2019).

### Redshifts outside the training set

The LaCE emulators assume that the \(P_{\rm 1D}\) can be reconstructed without explicitly providing redshift information. This is crucial to analyse DESI \(P_{\rm 1D}\) data with any redshift binning of the data. To test if the assumption holds, we have evaluated the emulator's performance at a given redshift \(z=z_{0}\), when all snapshots at that redshift are removed from the training sample. In the top, middle, and bottom panels of Fig. 8, we show the percent error in the emulated \(P_{\rm 1D}\) at \(z=2.25\), \(z=3\), and \(z=4\), respectively, when the corresponding testing redshift is removed from the training sample. The shaded grey area indicates the 1% error requirement, which is achieved in all cases. This proves that the emulator can make \(P_{\rm 1D}\) predictions at arbitrary redshifts. We specifically choose to test the emulator at \(z=2.25\) instead of \(z=2\) due to the latter being on the boundary of the convex hull. Consequently, if we were to exclude \(z=2\) from the training sample, it would no longer be within the training space. Therefore, to ensure a comprehensive evaluation, we opt to test the emulator at \(z=2.25\) instead.

### Emulator comparison

The success of this initial version demonstrated that we could accurately emulate the \(P_{\rm 1D}\) to sub-percent errors from the amplitude and slope of the linear matter power spectrum. However, this first emulator version was only tested on simulations with the same IGM history as the training set.

LaCE-NN incorporates several modifications compared to LaCE-GP. The most notable difference is the replacement of Gaussian processes with neural networks. Additionally, LaCE-NN incorporates a new post-processing step that measures the \(P_{\rm 1D}\) along all three simulation axes (SS2.2) using 768 skewers. This increases the resolution of the measured \(P_{\rm 1D}\) and provides three \(P_{\rm 1D}\) measurements per snapshot. Furthermore, since our simulations are fixed-paired simulations, we have six \(P_{\rm 1D}\) measurements per snapshot. The new post-processing step also includes applying five optical-depth rescalings to each snapshot (SS4.2), resulting in a total of 9900 \(P_{\rm 1D}\) measurements. This represents a substantial improvement over the LaCE-GP data processing, which yielded 330 \(P_{\rm 1D}\) measurements, and already produces datasets that are difficult to handle with GPs.

### Cosmic variance

Even though the training simulations use "fixed-and-paired" conditions, their volume is very limited and thus cosmic variance should affect emulator predictions significantly. To estimate the impact of cosmic variance on the results, we compare the performance of the emulator for the emtral and seeb simulation pairs (see SS2.1). These pairs present cosmological parameters at the centre of the training parameter space; therefore, the precision of the emulator should be close to optimal for them. On the other hand, the central pair use the same set of Fourier phases as the training simulations, while the seeb simulations use different initial conditions. As a result, any difference in the performance of the emulator for these simulations isolates the impact of cosmic variance.

In the first two panels of Fig. 9, we display the precision of the LaCE-GP and LaCE-NN emulators for these two simulation pairs. As we can see, the two emulators present better than 1% performance for both cases, confirming that the impact of cosmic variance on emulator predictions is negligible. This is the consequence of

Figure 8: Emulator’s performance at arbitrary redshifts. The plot shows the percent error in the \(P_{\rm 1D}\) emulation at redshifts \(z=2.25\), \(z=3\), \(z=4\), when all snapshots at these redshifts are dropped from the training sample. The dotted line indicates the mean percent error, while the shaded-coloured area corresponds to the standard deviation across simulations. The shaded-gray area indicates the 1% error requirement. We specifically choose to test the emulator at \(z=2.25\) instead of \(z=2\) due to the latter being on the boundary of the course hull. Consequently, if we were to exclude \(z=2\) from the training sample, it would no longer be within the training space.

Figure 7: Leave-one-out test. Percent error in the emulator \(P_{\rm 1D}\) at different redshifts. The shaded grey area indicates the 1% error requirement for the emulator. The red line show the mean percent error and the red shaded region corresponds to the standard deviation of the percent error. The blue shaded region represents the observational uncertainties from the latest analysis of BOSS/eBOSS data.
using a polynomial fit to smooth \(P_{\rm 1D}\) before training and testing the emulator (see Appendix A), which is justified by the smoothness of this observable and greatly reduces the impact of cosmic variance on large scales. As shown in Pedersen et al. (2023), the impact of cosmic variance is much larger if no smoothing is applied (see also Appendix C for another smoothing strategy).

### Cosmologies and \(\Lambda\)CDM extensions not considered in the training set

The training simulations adopt a standard \(\Lambda\)CDM parameterisation and explore different values of the amplitude and slope of the primordial power spectrum while considering the same expansion and growth histories (see SS2.1). Given that \(P_{\rm 1D}\) is sensitive to the velocity field, it is important to check the precision of the emulator for other growth rates. To do so, we analyse the growth simulation pair (see SS2.1), which uses a different value of \(h\) relative to the training simulations. In the third panel of Fig. 9, we show that our emulators present sub-percent precision for this pair, confirming that our emulation strategy enables producing precise predictions for cosmologies outside the training set (Pedersen et al., 2021; Pedersen et al., 2023).

We now proceed to study if the emulation strategy also works for three \(\Lambda\)CDM extensions: massive neutrinos, running of the spectral index, and curvature. We test the previous scenarios using the neutrinos, running, and curved simulation pairs (see SS2.1), which present extreme values for the previous ingredients already ruled out by observations (Planck Collaboration et al., 2020): \(\sum m_{\nu}=0.3\) eV neutrino mass, \(\Delta n_{\rm s}/\mathrm{d}\log k=0.015\), and \(\Omega_{k}=0.03\), respectively. Remarkably, our emulators also present sub-percent precision for these simulations, as shown in the fourth, fifth, and sixth panels of Fig. 9. Consequently, we can use this emulation strategy to set accurate constraints on \(\Lambda\)CDM extensions.

### Thermal histories not included in the training set

In SS2.1, we introduced the reionisation pair, which implements a different He ii reionisation history than the training simulations, which result in a very different thermal history (see Fig. 1). In the last panel of Fig. 9, we show the performance of the LaCE-GP and the LaCE-NN emulators for this pair. As we can see, the precision of both emulators is similar and on average better than 1.5%, letting us conclude that the emulation strategy works even for He ii reionisation histories not considered in the training set. This is the consequence of emulating in the \(\sigma_{T}\)-\(\gamma\) space, which captures the state of the IGM at a particular cosmic time, instead of as a function of parameters encoding information about the entire ionisation or thermal history.

We find that the overall precision of the emulator for this pair is on average four times worse than for the central pair, especially for low redshift. To further understand this decrease in performance, in Fig. 10 we display the results of the LaCE-NN emulator for the reionisation pair without collapsing the wavelength information. Lines indicate the results for different redshifts, while shaded areas denote the level of uncertainty predicted by the emulator. As we can see, the emulator produces biased results for \(z=2\). To trace the origin of this issue, in Fig. 11 we display how the training, central, and reionisation simulations sample the emulator parameter space. Although the central and reionisation simulations present the same cosmology, we can readily see that the relation between \(\tilde{F}\) and \(T_{0}\) is very different for these. Furthermore, we find that the value of the

Figure 10: Percent error in the reionisation emulation for snapshots at different redshift. The solid lines indicate the mean percent error while the shaded areas correspond to the emulator’s predicted uncertainty. The prediction at \(z=2\) shows worse performance than the rest of the snapshots, which is potentially because at such redshift, \(\tilde{F}\) and \(T_{0}\) are outside the convex hull of the training data.

Figure 9: Emulator performance on simulations with different cosmologies as the one sampled by the training simulations. The \(\nu\)-sim (top panel) contains massive neutrinos, while the \(h\)-sim (bottom panel) modifies the growth rates with respect to the training sample. The dotted line corresponds to the mean percent error in the \(P_{\rm 1D}\) prediction, while the shaded area indicates the standard deviation across \(k_{\parallel}\).
reionisation parameters for \(z=2\) lies outside of the parameter space covered by the training simulations, which explains the decrease in accuracy.

An important observation is that the emulator assigns a higher uncertainty to the \(P_{\rm 1D}\) measurement at \(z=2\). This increased uncertainty arises from the awareness that the parameters associated with \(z=2\) lie outside the boundaries of the convex hull. By acknowledging this discrepancy, the emulator correctly accounts for the greater uncertainty in predicting \(P_{\rm 1D}\) values at this particular redshift.

Interestingly, the overall performance of the emulator for the reionisation pair is worse than for cosmologies and \(\Lambda\)CDM extensions not included in the training set. This result emphasises the significance of exploring various IGM models instead of only concentrating on densely sampling the cosmological part of the emulator parameter space. While some of this exploration can be carried out in a post-processing stage via mean flux and temperature rescalings (Lukic et al., 2015), the gas pressure smoothing scale cannot be rescaled so it is necessary to run simulations with different reionisation histories.

## 6 Extended Emulator

The LaCE-GP emulator has been developed and tested on scales ranging from \(k_{\parallel}\in(0,3]\) Mpc\({}^{-1}\). However, in order to cover the entire range of scales that DESI is designed to observe (Ravoux et al. in prep, Karacayli et al. in prep.), the LaCE-NN emulator has been developed to extend this range to \(k_{\parallel}\in(0,4]\) Mpc\({}^{-1}\). In this section, we present the results of an extended version of the LaCE-NN emulator, which predicts the \(P_{\rm 1D}\) on scales up to \(k_{\parallel}=8\) Mpc\({}^{-1}\). It is worth noting that while DESI cannot observe scales beyond \(k_{\parallel}=4\) Mpc\({}^{-1}\), this extension has been motivated by the \(P_{\rm 1D}\) measurements obtained from high-resolution quasars in Karacayli et al. (2022).

Figure 12 shows the performance of the extended emulator in the test simulations (left panel) and the training simulations (right panel). The figure includes the LaCE-GP emulator's performance for comparison purposes. In both cases, we employ the same emulator as used for \(k_{\parallel}<4\) Mpc\({}^{-1}\) without any adjustments to the hyperparameters. However, we modify the order of the polynomial used to fit the \(P_{\rm 1D}\), which is now set to \(n=7\). Further details can be found in Appendix A.

The LaCE-NN emulator outperforms the LaCE-GP performance in the extended case. Without any tuning of the emulator, LaCE-NN reaches a 1-2% error in all test simulations but the reionisation one, which is the most challenging. On the training simulations (right panel), LaCE-NN emulates the \(P_{\rm 1D}\) of the training simulations with a mean percent error of only \(\sim 2\%\). This is lower at \(k_{\parallel}<4\) Mpc\({}^{-1}\) and higher for smaller scales. However, note that the error requirement for very small scales is not as stringent as 1% (Karacayli et al., 2022). We present our initial attempt to create an extended emulator, acknowledging that there is room for improvement in future iterations. For instance, we aim to address the degradation observed at larger scales, and we propose potential enhancements such as down-weighting the contribution of smaller scales.

The extended emulator's performance exhibits a notable decline in the case of the reionisation simulation (Fig. 13). Specifically, we observe a significant degradation in performance compared to the fiducial emulator. The errors primarily stem from inaccuracies at high \(k_{\parallel}\) values, but we also observe a deterioration in predictions across all scales, particularly at \(2<z<3\). At the redshifts in question, the parameter values for \(\tilde{F}\) and \(T_{0}\) lie at the boundaries or beyond the parameter space covered by the training simulations (see Fig. 11). While this situation did not pose significant issues for the fiducial emulator for \(z>2\), it has a pronounced impact on the extended emulator, also at larger scales. Again, an avenue for enhancing these results would involve reducing the influence of scales with \(k_{\parallel}>4\) Mpc\({}^{-1}\). However, we defer this exploration to future investigations.

The poor performance of the extended emulator in the reionisation highlights the need to expand the training sample parameter space by incorporating additional temperature and mean-flux rescalings that encompass universes with diverse IGM histories. Another noteworthy observation is that the emulator predicts larger uncertainties for the snapshot at \(z=2\) when compared to other more accurate predictions. This feature aligns well with the specific snapshot lying beyond the training parameter space.

## 7 Conclusions

In this paper, we build the first emulator of the one-dimensional Lyman-\(\alpha\) flux power spectrum (\(P_{\rm 1D}\)) using a neural network (NN) architecture. To do so, we adopt the emulation strategy devised by P21, which relies on emulating \(P_{\rm 1D}\) as a function of the amplitude and slope of the linear matter power spectrum on small scales. We summarise our main findings below:

* In SS3.2, we build an emulator that uses a Mixture Density Network (MDN) to predict the probability distribution of six polynomial coefficients describing \(P_{\rm 1D}\) measurements. Then, it combines these distributions to generate predictions for the best-fitting solution and error prediction for each combination of input parameters. In Fig. 4,

Figure 11: Distribution of training, central, and reionisation simulations in the space of emulation parameters. Small dots indicate the results for training simulations with optical depth rescalings, while pink stars and red triangles show the results for central and reionisation, respectively. Small dots are coloured by the redshift of the simulation snapshot for visual purposes, as the redshift information is not considered by the emulator. Even though the reionisation and training simulations present different astrophysical implementations, we can readily see that the first lies within the range of the parameter space covered by the latter.
we show how different decisions regarding the configuration of the emulator improve its precision. On the other hand, in Fig. 6 we show that thanks to its MDN architecture, the emulator performance improves by 20% when training it using \(P_{\rm 1D}\) measurements from different simulation axes and phases instead of just relying on the average of these.
* In Fig. 7, we show that the emulator precision is better than 1% for cosmologies within the training set (leave-one-out tests) across the full range of scales (\(0.1<k_{\parallel}[\,\mathrm{Mpc}^{-1}]<4\)) and redshifts (\(2<z<4.5\)) considered. Even though this value is similar to the precision quoted for the GP-based emulator described in P21, the actual performance of our emulator is better because we carry out a more detailed post-processing of the suite of hydrodynamical simulations from which these two studies extract \(P_{\rm 1D}\) measurements.
* In Fig. 9, we show that emulator predictions are largely insensitive to the impact of cosmic variance thanks to emulating on the space of polynomial coefficients (see also Pedersen et al., 2023). Furthermore, we show that our emulator presents sub-percent precision for growth histories not included in the training set as well as for three \(\Lambda\)CDM extensions: massive neutrinos, running of the spectral index, and curvature. These findings confirm the advantage of emulating as a function of the amplitude and slope of the linear matter power spectrum on small scales rather than as a function of cosmological parameters.
* In Fig. 10, we show that the emulator achieves on average 1.5% precision for thermal and reionisation histories not considered in the training set. This is the consequence of emulating \(P_{\rm 1D}\) as a function of the instantaneous properties of the IGM rather than parameters encoding information about the entire reionisation or thermal history of the universe.
* In Fig. 12, we show the performance of an extended version of our emulator reaching \(k_{\parallel}=8\,\mathrm{Mpc}^{-1}\). In the central simulation, we find that the extended emulator presents an overall 1% precision, a factor of \(\sim\)2 worse than the fiducial emulator. Its performance is especially bad for thermal and reionisation histories not considered in the training set, reaching on average 3.5%.

Figure 12: Performance of an extended version of the emulator to \(k_{\parallel}=8\,\mathrm{Mpc}^{-1}\) for LaCE-NN and LaCE-GP. The _left_ panel presents the leave-one-out test for the extended emulator. In the _right_, the emulator is trained using the tradinno simulations and tested on the seven test simulation pairs presented in §2.1

Figure 13: Percent error in the reionisation emulation for snapshots at different redshifts. The solid lines indicate the mean percent error while the shaded areas correspond to the emulator’s predicted uncertainty. This plot is analogous to Fig. 10 but for the extended emulator.
As shown in SS5, the overall performance of the emulator is better for cosmologies and \(\Lambda\)CDM extensions not included in the training set than for reionisation histories not considered, which emphasises the importance of running simulations adopting distinct reionisation histories. This issue can be ameliorated by carrying out mean flux and temperature rescalings in a post-processing stage (Lukic et al., 2015), leading to a significant increase in the size of the training set. Such an increase makes NN-based models more suited to emulate \(P_{\rm 1D}\) measurements than GP-based models because the runtime of the first and second increase linearly and with the cube of the number of input points, respectively.

We designed the fiducial version of our emulator aiming to analyse medium-resolution spectra from the DESI survey, which explains the range of redshifts and scales considered. In SS6, we present an extended emulator conceived for the joint analysis of DESI and high-resolution measurements (e.g., Karagayli et al., 2022) that reaches a factor of two smaller scales than the fiducial one. We found that the NN-based emulator performs significantly better than the GP-based emulator on small scales, further highlighting the advantages of NNs for \(P_{\rm 1D}\) emulation. On the other hand, we found the performance of the extended emulator is on average a factor of two times worse than the one of the fiducial emulator. In future work, we will further work on improving the precision of the extended emulator.

Throughout this paper, we have focused on the one-dimensional Lyman-\(\alpha\) flux power spectrum. However, this statistic only contains part of the cosmological and astrophysical information encoded in the Lyman-\(\alpha\) forest as it neglects correlations between different lines-of-sights. In future work, we will use the training simulations and an NN-based architecture to develop the first emulators for the Lyman-\(\alpha\) flux probability distribution (Lee et al., 2015), one-dimensional bispectrum (Viel et al., 2009), and three-dimensional power spectrum (Font-Ribera et al., 2018). A coherent analysis of \(P_{\rm 1D}\) and these complementary statistics would enable fully exploiting the constraining power of the Lyman-\(\alpha\) forest.

## Acknowledgements

LCG, JCM and AFR acknowledge support from the European Union's Horizon Europe research and innovation programme (COSMO-LYA, grant agreement 101044612). AFR acknowledges support from the Spanish Ministry of Science and Innovation through the program Ramon y Cajal (RYC-2018-025210). IFAE is partially funded by the CERCA program of the Generalitat de Catalunya. The analysis in this article have been performed at Port d'I Informacio Cientfica (PIC). We would like to acknowledge the support provided by PIC in granting us access to their computing resources. The simulations were run using the Cambridge Service for Data Driven Discovery (CSD3), part of which is operated by the University of Cambridge Research Computing on behalf of the STFC DiRAC HPC Facility (www.dirac.ac.ac.uk). The DiRAC component of CSD3 was funded by BEIS capital funding via STFC capital grants ST/P002307/1 and ST/R002452/1 and STFC operations grant ST/R00689X/1. DiRAC is part of the National e-Infrastructure

## Data Availability

The simulations utilized for the development and testing of this work are publicly accessible via the following link: [https://github.com/igmhub/LaCE](https://github.com/igmhub/LaCE).
