PINION: Physics-informed neural network for accelerating radiative transfer simulations for cosmic reionization

With the advent of the Square Kilometre Array Observatory (SKAO), scientists will be able to directly observe the Epoch of Reionization by mapping the distribution of neutral hydrogen at different redshifts. While physically motivated results can be simulated with radiative transfer codes, these simulations are computationally expensive and can not readily produce the required scale and resolution simultaneously. Here we introduce the Physics-Informed neural Network for reIONization (PINION), which can accurately and swiftly predict the complete 4-D hydrogen fraction evolution from the smoothed gas and mass density fields from pre-computed N-body simulation. We trained PINION on the C$^2$-Ray simulation outputs and a physics constraint on the reionization chemistry equation is enforced. With only five redshift snapshots and a propagation mask as a simplistic approximation of the ionizing photon mean free path, PINION can accurately predict the entire reionization history between $z=6$ and $12$. We evaluate the accuracy of our predictions by analysing the dimensionless power spectra and morphology statistics estimations against C$^2$-Ray results. We show that while the network's predictions are in good agreement with simulation to redshift $z>7$, the network's accuracy suffers for $z<7$ primarily due to the oversimplified propagation mask. We motivate how PINION performance can be drastically improved and potentially generalized to large-scale simulations.


INTRODUCTION
Immediately after the Epoch of Recombination, no significant radiation sources existed in the Universe.This period is called the Dark Ages.However, the evolution of the matter density field led to the generation of collapsed structures, leading to the formation of the first stars and galaxies.At  ∼ 30, these first luminous sources started to emit ultraviolet light, which successively ionized the intergalactic medium (IGM) (Furlanetto et al. 2006;Ferrara & Pandolfi 2014).Indirect constraints such as high-redshift quasar spectra (Fan et al. 2006;McGreer et al. 2011McGreer et al. , 2014)), the decline of Lyman  emitting galaxies (Schenker et al. 2011;Stark et al. 2011;Pentericci et al. 2014;Tilvi et al. 2014) suggest that reionization was complete for  < 6.This change from a dark universe filled with neutral hydrogen to the fully ionized IGM is known as the Epoch of Reionization (EoR) (Gorbunov & Rubakov 2011;Dayal & Ferrara 2018).The study of the EoR is essential for cosmology.The evolution of reionization is driven by the matter-energy content and geometry of the Universe, and the distribution of ionized regions reflects how galaxies and black holes formed.
The presence of neutral hydrogen during the EoR will be evident by measuring the emerging brightness temperature from a gas cloud at a given redshift against the CMB temperature (Zaroubi 2013).Modern telescopes such as LOFAR 1 are able to measure the global ★ Corresponding author: Damien Korber damien.korber@unige.ch 1 http://lofar.orgstatistics of the 21-cm signal (Paciga et al. 2013;Yatawatta et al. 2013;Parsons et al. 2014;Jelić et al. 2014;Jacobs et al. 2015;Patil et al. 2017;Mertens et al. 2020;Trott et al. 2020).However, during reionization the signal fluctuations are expected to be non-Gaussian (Mondal et al. 2017;Shimabukuro et al. 2017;Majumdar et al. 2018;Hutter et al. 2019;Gorce & Pritchard 2019;Watkinson et al. 2022) and therefore can not be fully depicted by the power spectra only.Next-generation telescopes such as the SKA 2 will be able not only to detect the 21-cm signal but also to produce maps of its fluctuations on the sky (Mellema et al. 2013;Iliev et al. 2014;Dillon et al. 2015;Wyithe et al. 2015;Koopmans et al. 2015;Bacon et al. 2020).These observations at multiple frequencies will allow us to directly measure the neutral hydrogen density distribution throughout the Universe evolution.The expected precision of SKA-Low will allow a groundbreaking comparison between observations and simulations.Future high-fidelity observations of the EoR demand equivalently high-fidelity simulations.A physically motivated simulation of reionization needs to include three main components: gravitational evolution of dark matter particles via N-Body simulations, gas dynamics via hydrodynamical simulations, and radiative transfer (RT) simulations to compute the propagation of photons and the evolution of the ionized atoms.Although it is possible to do on-the-fly simulations that combine these three steps (Rosdahl et al. 2013;Ocvirk et al. 2016Ocvirk et al. , 2020Ocvirk et al. , 2021;;Lewis et al. 2022), these simulations are very computationally expensive as they require small time steps and large 2 https://skatelescope.org number of simulated particles (Abel et al. 1999;Bolton et al. 2004;Shapiro et al. 2004).It is also possible to post-process the RT step onto existing N-body and hydrodynamical simulations in order to lighten its computational footprint.The state-of-the art post-process RT simulation is the C 2 -Ray code (Mellema et al. 2006).This method uses a 3-D ray tracing method that relaxes the constraints on the time step size.While C 2 -Ray can produce high-accuracy EoR simulations at large scale, is extremely computationally expensive to run, requiring several millions of core-hours to run across thousands of nodes.
Another post-process RT code is Grizzly (Ghara et al. 2018), which assumes spherical density profiles around each source, thus simplifying the problem to a 1-D RT process and speeding up the computation by several orders of magnitude.There is no simulation which is able to produce RT on large volumes while still accounting for small scale physics.
Deep neural networks are well-known as universal function approximators (Hornik 1991), and have been successfully used to accelerate and improve simulations of physical processes in many domains (He et al. 2019;Liu et al. 2021).Deep learning has previously been applied in the context of reionization simulations.In particular, (Chardin et al. 2019) obtained some encouraging results for the use of neural networks in predicting the 3-D ionization map at  = 6.
Despite their success, neural networks suffer from interpretability and generalization problems (Nguyen et al. 2015;Ribeiro et al. 2016).While physical phenomenons are often described by some theoretical or empirical model, neural networks attempt to learn these laws from statistically small and biased data.However, in recent years, a new class of neural networks with physics constrains has emerged.These networks are called Physics-Informed Neural Networks (PINNs) and encapsulate additional physics laws into the behaviour of the network (Karniadakis et al. 2021).
In this paper, we present PINION, a physics-informed neural network which can accurately predict the reionization map at any redshift and volume.This paper is organized as follows.In § 2 we present the numerical codes employed to simulate the reionization process and create the network training dataset.In § 3 we present the model used to study the evolution of ionization during the EoR and data preprocessing by smoothing the source distribution field to encapsulate the expanding behaviour of the ionising front.Then in § 4 we present PINION, the neural network that we developed for this project, with its physics constraint.§ 5 gives the results and their analysis, and § 6 concludes with a summary of the results and the future prospects for this project.

REIONIZATION SIMULATIONS
The simulation used in this paper is briefly described here below, but we refer to the reader the recent works (e.g.Dixon et al. 2016;Ross et al. 2017;Bianco et al. 2021) for further details.The cosmological parameters used in the simulation are based on the WMPA threeyear data observation and the results are consistent for a ΛCDM cosmology with the following parameters, Ω Λ = 0.724, Ω  = 0.276, Ω  = 0.454,  0 = 70.1 km s −1 Mpc −1 ,  8 = 0.784 and   = 0.946 (Spergel et al. 2007).The CUBEP 3 M code (Harnois-Déraps et al. 2013) is employed for the N-body simulation and it is run on a box of physical size 714 Mpc containing 6912 3 particles, with a spatial resolution of 5.17 kpc and particles mass of 5.12×10 3  .The initial conditions are generated using the Zel'dovich approximation, while the power spectrum of the linear fluctuations is given by the CAMB code (Lewis et al. 2000).The simulation then starts at redshift  = 150, which gives enough time to significantly reduce the non-linear decaying modes (Crocce et al. 2006).An on-the-fly spherical overdensity halo finder algorithm (Harnois-Déraps et al. 2013;Watson et al. 2013) with over-density parameter Δ = 130 creates a halo list catalogues at each redshift step, while the remaining particles are considered to be part of the IGM.In this study, we only consider halos with mass  ℎ ≥ 10 9  as these sources are not affected by radiative feedback and are considered to be the primary driver of reionization (Iliev et al. 2012;Dixon et al. 2016).
The N-body particle list and the halo catalogues are then interpolated into a 300 3 grid with a Smoothed-Particle-Hydrodynamic-like method (Shapiro et al. 1996;Mao et al. 2012), with a corresponding cell size of 2.381 Mpc.This dark matter density field and cumulative source mass field are used as inputs to the C 2 -Ray code, which computes the post-processed RT to generate the photoionization rates and ionization fractions.A fixed time step Δ = 11.54Myr is considered for redshifts  = 6 -50.We refer the reader to, e.g.(Dixon et al. 2016;Majumdar et al. 2016;Bianco et al. 2021) for a more general overview of the RT method and N-Body method employed in this paper.In our case, snapshots of C 2 -Ray for the redshifts range  = 6 -12 were kept for training the network, which corresponds to the time when at least one source is present at each pixel of the smoothed grid and therefore where the reionization process is rapidly evolving.
The matter density field and source mass field used by C 2 -Ray are also provided as inputs to the neural network.

EVOLUTION OF REIONIZATION
In the EoR we are interested in studying how the neutral hydrogen density  HI is converted into the ionized hydrogen density  HII .Reionization is driven by two different processes: ionization, by which high-energy photons ionize neutral hydrogen, and recombination, by which ionized hydrogen recombines with the local electron density  e into neutral hydrogen.The evolution of ionized hydrogen is given by (Choudhury 2009;Pritchard & Loeb 2012): where Γ is the reionization rate, C is the clamping factor which accounts for the inhomogeneity in the simulation, and  B ≡  B () = 2.59 × 10 −13 (/10 4 K) −0.7 cm 3 s −1 is the temperature dependent case B recombination coefficient (Furlanetto et al. 2006).In this work, the clumping factor is set to C = 1.This scenario corresponds to the case where sub-grid inhomogeneities in the IGM are ignored (Mao et al. 2019;Bianco et al. 2021).Moreover, we set the temperature to  = 10 4 K, which assumes that gas in halo is able to radiatively cool through hydrogen and helium atomic lines (Wise et al. 2014).
For a fixed density of hydrogen  H ≡  HI +  HII , we can define the neutral hydrogen fraction  HI ≡  HI / H and ionized hydrogen fraction  HII ≡  HII / H .For simplicity, we assume pure hydrogen gas,  e ∼  HII .We can then rewrite equation (1) as: Note that in the following of this paper, we will refer to  HII as the volume-average fraction of ionized hydrogen.

PHYSICS INFORMED NEURAL NETWORK FOR REIONIZATION
PINION 3 (Physics Informed Neural network for reIONization) is a 3-D CNN (Convolutional Neural Network) that reduces the information from two 7 3 input cubes of gas density and mass of sources to predict  HII at the cube centre for a single value of redshift.The input size corresponds to a volume of (16.67 Mpc) 3 .We initially chose a size scale based on the physical hard limit imposed by Lyman limit system (Shukla et al. 2016), corresponding to an input mesh-grid size of 49 3 .However, we did not observe significant improvement in the result for this larger subvolume, and decided to reduce the input size to reduce the training and predicting computation time.We considered periodic boundary conditions on the full dataset to be able to predict the border pixels.This is justified by the fact that the C 2 -Ray simulation we are training on also uses periodic boundary conditions.

PINION pipeline and architecture
PINION has 2 input channels: the source mass field and the gas density.An additional third input is derived internally by smoothing the source distribution field, described in § 4.2.These three fields are then internally given as inputs to the network, which then returns the ionized fraction  HII of the subvolume central pixel as the output.The network is composed of three convolutional layers with 3 3 kernel and doubling number of feature maps: 64, 128, 256.The convolution results in 256 feature maps with channel size 1, to which the normalized look-back time is concatenated.This tensor is then used as input to a fully connected network with three hidden layers of 64, 16, and 4 nodes, and a single output node.Each layer uses the PReLU activation function (He et al. 2015) followed by a 3-D batch normalization and a 10% dropout.The output node uses a sigmoid activation function to truncate the output values between 0 and 1.In total, this network has 1,130,252 trainable parameters.In Figure 1, we show a simplified overview of PINION pipeline, the pre-process step and the neural network architecture.We use the Adam optimization algorithm (Kingma & Ba 2014) during the network training to minimize our data loss function L data .In PINION, the data loss is the adjusted  2 loss, also referred to the 3 https://github.com/epfl-radio-astro/PINIONcoefficient of determination in statistics.This quantity is defined as follows: where  pred is the prediction from the network,  true is the ground truth and xtrue is its mean value (Gillet et al. 2019).
We train with the ODE loss as defined in § 4.3.The results from this network will be discussed in § 5 and show promising results.

Data pre-processing
During initial studies, we found that it was difficult for the network to learn the relationship between the input fields from the N-body simulation and the reionization across all redshifts.This is because, as shown in Figure 2, while the gas density and source mass fields do not change drastically over the EoR, the reionization rate and fraction rapidly change as ionized regions progressively grow with time.Without any other constraints, the network must learn to predict extremely different outputs from very similar inputs across all redshifts.
To ameliorate this problem, we provide the network with the ionizing front propagation as an additional input.This approach was inspired by (He et al. 2019), which used a neural network to predict cosmological structure formation starting from the Zel'dovich Approximation (Zel'Dovich 1970) and in our case this input should give a crude approximation of the photoionization rate.We can calculate the mean free path of ionizing photons at a given redshift   HI () using the approximated form of equation ( 23) in (Choudhury 2009): This equation is the empirical fit from simulations valid for redshifts  > 3.While the source model in (Choudhury 2009) does not exactly correspond to our simulation, it provides a good starting point for the network.We can then approximate the size of the ionized bubbles using   HI (), and use the smoothed sources mass field to infer the photoionization rate.The algorithm is the following: (i) Calculate   HI () for a given redshift .
(ii) Construct a spherical kernel of size   HI () by converting   HI () from Mpc to number of pixels given a resolution of  2.381 Mpc per pixel, setting every pixel to 0 when further than distance   HI () to the centre and scale the other pixel value between 0 and 1, where 1 is the centre pixel and 0 the borders.
(iii) Finally, convolve cube of mass of sources with the spherical kernel.
This algorithm creates a smoothed source distribution field that behaves approximately like the photoionization rate.Knowing that   HI () increases with redshift, the ionization front expand as well.
The purpose of this smoothed field is to provide an additional input to the network, and outside this machine learning context, it should not be interpreted as an estimate of the photoionization rate.

Physics constraint
One of the most popular ways to encode physics constraints in a deep neural network is through learning bias (Karniadakis et al. 2021).
We have incorporated this in our network by using the ODE from Equation 2 as an additional loss: This ODE depends on the time derivative of  HII .While automatic differentiation can be used to obtain the exact value of the derivative (Baydin et al. 2018), it is unfortunately not feasible for a network with millions of parameters.Instead, we encode the ODE loss as a finite-difference, using the Runge-Kutta 4 (RK4) method (Scherer 2013) to calculate the derivative.RK4 requires an intermediate step to perform derivations, so we simply take a step size two times larger than the time increment of the data.Let   ≡  HII (  ) the ionization fraction at the -th time step and Γ  ≡ Γ(  ) the photoionization rate at the -th time step.Then, defining ℎ  = Δ  + Δ +1 as the new time step, where Δ  is the time between snapshots  and  + 1, and   =  B  H, , we can define the RK4 terms as: where   is the predicted mean fraction of ionized hydrogen, obtained as the output of the neural network,  * +2 is the step obtained by computing the evolution of the ODE from   and Γ  is the true photoionization rate from C 2 -Ray.The large time step makes it difficult to keep the solution stable.Indeed, at higher redshifts, most of the Universe is neutral, but at lower redshifts the ionization quickly increases.With large time steps, RK4 has no issues tracking the evolution of the ODE at earlier times, as changes are slow.However, at later times, the photoionization rate is rapidly changing and the RK4 solution might drastically increase.To fix this issue, the Equation 6e is clipped between 0 and 1 to guarantee that the solution does not diverge.One can now write the physics loss as follows.
Where MinMax(, 0, 1) ensures that  is encapsulated between 0 and 1.In Figure 3, we compute the residual to check that the RK4 loss accurately describe the behaviour of the C 2 -Ray simulation.We see good agreement between the data and the ODE as the standard deviation is within 2% for all redshifts.When the reionization fraction is not rapidly evolving for  > 9, we see excellent agreement within 0.5%.
Although we use 3-D convolutional layers, we account for the time evolution constraint by using the batch dimension of the network input.For a given spatial position of the input fields, its evolution for the 46 redshift snapshots is added as adjacent inputs along the batch dimension.The truth information is fed to the network in the same manner.This way, we can use spatial convolution and study the evolution of the predicted quantity with the time evolution constraint from the ODE loss.This additional loss constrains the network's behaviour and improves performance.

PINION and approximate RT solvers
Compared to other fast approximations of RT, PINION offers a more physically motivated but less interpretable solution.
For example, radiative transfer codes such as GRIZZLY assume an uniform density profile around each isolated source (Ghara et al. 2018).This approximation reframes the reionization process as a 1D problem using rotational symmetries, which considerably simplifies the computations.Therefore, this approach computes the ionization profile around isolated sources, thus quantifying the radius of the spherical ionisation front as a function of time and optical depth (Shapiro et al. 2006) and then redistributes photons in overlapping regions in a second step (Ghara et al. 2015).Recent works use 1D RT codes to study global statistics of the 21-cm signal for EoR and Cosmic Dawn (Datta et al. 2022;Kamran et al. 2022;Ghara et al. 2020Ghara et al. , 2021)).
PINION does not make any assumptions or approximations about the source or density field.However, PINION must infer the dynamics of RT during training, learning from a subvolume of size (16.67Mpc) 3 which we consider the region of influence.We assume that this region contains all the nearby sources that contribute to the ionization state of the central pixel.During training, the ODE constraint Equation 7 enforces realistic behavior of the learned dynamics.However, as in all deep learning solutions, PINION results will have less interpretability compared to approximate 1D RT codes.

RESULTS AND ANALYSIS
In order to characterize the impact of the ODE loss, we consider three different network training scenarios: (i) NP (No Physics): The training contains the entire evolution of pixels, but does not have a physics constraint: L total = L data .This scenario is noted by dashdot lines (•−) in figures.
(ii) PFD (Physics and Full Data): The training uses the ODE loss and the data loss for the entire time evolution of pixels: L total = L data + L ODE .This scenario is noted by dashed lines (−−) in figures.
(iii) LD (Low Data): The training uses the ODE loss, but the data loss is only evaluated at 5 different snapshots (limited) of the simulation instead of the full reionization history: L total = L limited + L ODE .This scenario is noted by dotted lines (••) in figures.
During the training, the three scenarios do not differ greatly in convergence time.Training each scenario took 1.5 GPU-hours on one node for 400 epochs.The training was performed on a node from CSCS's GPU cluster Piz Daint4 that is equipped with NVIDIA Tesla P100 16GB GPU.In each scenario, only the network weights with the lowest validation loss is kept.
To predict the reionization fraction at a single pixel location, the network needs the surrounding 7 3 pixel subvolume, corresponding to a volume of (16.667Mpc) 3 , for the three internal input fields.These subvolumes are sliced out of the full cubes and are concatenated to form a higher dimensional tensor for the training dataset.For the PFD and NP scenarios, 4000 subvolumes over 46 redshifts are provided to the network (184,000 training subvolumes in total) with a batch size of 3680.For the LD training, 4000 subvolumes over 46 redshifts with a batch size of 3680 are also considered, but the data loss only had the ground truth for 5 fixed redshifts evenly spaced over the EoR evolution: ( = 12.048, 9.938, 8.515, 7.480, 6.686) in the evolution, meaning that it compares the prediction to the truth data on 20,000 subvolumes instead of 184,000 for L data .To validate the training, 500 validation subvolumes are considered, and the final network weights are selected from the training snapshot with the smallest total validation loss.Note that all the subvolumes are randomly chosen from the full dataset.

Network output
The network results for each PINION training scenario are compared to the C 2 -Ray ground truth in Figure 4 for a single redshift slice.In general, all the scenarios are able to accurately reproduce the ionization map.However, small differences can be observed at the borders of large and small ionized regions.
The  HII redshift evolution in each scenario is shown next to C 2 -Ray in Figure 5.The three scenarios evolve similarly and show good agreement with the ground truth.All the PINION scenarios are able to find the first ionizing bubbles ( > 8).See for example the bubble at (, ) = (9, 610 Mpc) and (8.1, 610 Mpc) in Figure 5.However, for redshift  < 7, we start to see large differences between C 2 -Ray and the PINION predictions.For instance, the neutral island at (, ) = (6.4,600 Mpc) or (6.6, 490 Mpc) on the same figure.Between the different training scenarios, the NP scenario exhibits much more noise at low redshift ( ≈ 6.5) compared to the physics-constrained scenarios.Furthermore, we observe that neutral islands are smaller at this redshift in all predictions.Nevertheless, NP scenario appears to predict the most accurate neutral island size.

Evolution of the volume-averaged xHII
We then evaluate the evolution of the mean fraction of ionized hydrogen over several distinct subvolumes.Figure 6, top panel, shows the evolution of the ionized fraction for four different subvolumes, which each have different ionization history.Generally, we see very good agreement for  > 8.For early ionizing subvolumes (vol.3 and 4), the PINION reionization fraction predictions are always postponed for the NP scenario.The physically constrained scenarios (PFD and LD) ionize earlier for  > 7.5, and later otherwise.For late ionizing subvolumes (vol. 1 and 2), as well as the global average (black line), when  7 we observe good agreement for the LD scenario, overprediction for the PFD scenario, and under-prediction for the NP scenario.For  < 7 all scenarios ionize earlier than the simulation.
In general, we observe that for redshifts  7 the LD scenario gives better results.For  7, the NP scenario is better at reproducing later ionizations, while LD remains better for earlier ionizations.In the residual plot, we observe a higher standard deviation at lower redshift ( 7.5) in general.For vol. 1 and 2 the deviation peaks at  ≈ 6.6 while vol. 3 and 4 have a negative peak earlier at  ≈ 7.1.The global average shows both behaviours, with a flex point at  ≈ 7, but less residual error.This is correlated with what observed in § 5.1 for the same redshift range.In Figure 6 bottom panel, we have the residual error of our prediction.We can see for  7 and all subvolumes that the LD scenario follow well the simulation, while the NP scenario under-predicts with residual under 2.5 % for  7.5 while the PFD over-predicts it with less than 1 % difference for  7.5.In Table 1, we show the corresponding redshift of four reionization milestones, when xHII = 0.3, 0.5, 0.7 and toward reionization completion at xHII = 0.95, interpolated from the results in Figure 6.In § C we show a comparison between mass-averaged and volume-averaged global ionization fraction.

Morphology of the ionization bubbles
In the study of the EoR, of particular interest are the size and volume distributions of the ionized regions in the simulation (Furlanetto et al. 2004).There are several methods to measure these statistics (Giri et al. 2018;Lin et al. 2016;Friedrich et al. 2011), but in this work we consider the mean free path method and the friend-offriend method.The former indicates the fraction of ionized bubbles in a given spherical-averaged size range, while the latter indicates the fraction of the total ionized volume that is contained in regions of a given volume.For both methods, we employ the Tools21cm5 python package for EoR simulations analysis (Giri et al. 2020).
Figure 7 shows the mean free path  of the different PINION scenarios and the simulation plotted against the probability distribution ().We chose to study redshifts for which xHII ≈ 0.3, 0.5, 0.7 which approximately corresponded to redshifts  = 7.480, 7.059, 6.830.All three PINION scenarios exhibit similar morphology.In particular, in each of them, we observe an excess around the peaks, meaning that we have a more uniform distribution of bubble size.The fact that we observe this for all three redshift means that the universe evolves more uniformly with the PINION prediction than with the C 2 -Ray simulation.
For redshifts  = 7.480 (blue) and  = 7.059 (orange), PINION predicts an excess of smaller scale bubbles below the median, and a lack of larger scale bubbles above the median compared to C 2 -Ray.For redshift  = 7.480 (blue), only the NP scenario lacks large scale bubbles ( 20 Mpc), which means that the NP scenario produces smaller bubbles on average than the PFD and LD scenarios.For redshift  = 7.059 (orange), however, we observe that all three scenarios have a turning point above which the prediction lacks large scale bubbles.For large and small scale bubbles, the LD scenario gives a more accurate distribution, but around each peak, the PFD scenario tend to be closer to the distribution in the simulation.
For redshift  = 6.830 (green), we observe a deficit of small-scale bubbles for all scenarios.In general, we observe that the predicted distributions are shifted to the larger scale, meaning that the bubbles are larger on average than in the simulation.
We also study the volume distribution with the friend-of-friend method.Contrary to the mean free path, which measures the size of the bubbles, the friend-of-friend algorithm evaluates the volume  The x-axis is common to both plot and was cropped at redshift  = 10 for readability.Some characteristic redshifts can be found in Table 1.Bottom panel: The residual between the PINION scenarios and the C 2 -Ray simulation for the above evolutions are displayed with the same colours.In blue, the standard deviation of the residual for at each redshift for 1,000 different subvolumes of the same size is shown.
distribution of the bubbles using an island-finding algorithm.Pixels are recursively grouped into islands based on a threshold criterion, in our case of  HII > 0.5.The average volume of each island  is then plotted against the cumulative probability distribution  ().
Figure 8 shows the result for the friend-of-friend analysis.First, for redshift  = 7.480 (blue), one can see that both the PFD and LD scenarios are able to accurately reconstruct the bubble volume distribution and the percolation cluster6 .On the other hand, the NP scenario tends to over-predict the number of non-overlapping bubbles, as one can observe an over-prediction of the bubbles distribution and an underestimation of the total volume of the percolation cluster.
For redshifts  = 7.059 (orange) and  = 6.830 (green) all three scenarios are able to reproduce the volume distribution of both the bubbles and the percolation cluster.In general, we observe that the volume size distribution is well reproduced by PINION regardless of the scenario, aside from NP at higher redshifts.

HI dimensionless power spectrum
We also evaluate the dimensionless power spectrum of the neutral hydrogen field Δ HI = ()  3 2  2 for both C 2 -Ray and PINION, shown in Figure 9.We observe almost perfect agreement in large and small scale structure for  = 7.480 (blue).However, as ionization progresses, we begin to see tension between PINION and C 2 -Ray.At larger scales,  < 10 −1 Mpc −1 , the power spectrum is underestimated by all PINION predictions at all redshifts, meaning a statistically more ionized field then the ground truth.We observe that at these scales the discrepancy between the C 2 -Ray simulation and the PINION prediction increases with decreasing redshift, from almost perfect agreement at  = 7.480 (blue) to up to a factor of 2 difference for  = 6.830 (green).differences can be observed in the x projection.We observe a larger residual at higher redshifts, with (−4.3 ± 0.5) %, (−3.1 ± 0.6) %, (−3.3 ± 0.6) % for  = 11.834 in the PFD, NP and LD scenarios respectively.As shown in §5.4,this indicates that PINION ionizes earlier than C 2 -Ray at high redshift.However, this difference reduces with redshift and in the particular case of NP, will have a positive peak of (1.4 ± 1.5) % at  = 8.133 which this time indicates that at this ionization redshift, PINION tends to ionize later compared to C 2 -Ray.For the y projection, we observe a large positive bias in both NP and LD with (4.1 ± 3.5) % for NP and (2.2 ± 3.1) % at  = 9.414 for LD.Near the end of the EoR,  ≤ 7, we observe a sudden change in the residuals of every scenario.We believe this results from small variations on the borders of neutral islands, which are difficult for the network to predict precisely.
We can also use Figure 10 to compare PINION to (Chardin et al. 2019), a network which predicts the 3D  reion map.PINION shows much better agreement at early and mid redshift, likely due to the addition of the physics constraint.Additionally, because PINION outputs  HII , while PINION can reproduce (Chardin et al. 2019), the inverse is not possible.

Interpretation
We observe a few consistent differences between C 2 -Ray and the PINION predictions in our analysis.PINION predictions are more uniform, having more uniform bubble sizing and less variety in mean ionization fraction.The NP scenario ionizes later at high redshift, whereas the physics-constrained scenarios PDF and LD ionize sooner.For redshifts with rapidly changing ionization fraction, at  7, there is tension between the PINION predictions and C 2 -Ray at small and particularly large scales.There are several possible explanations for these systematic errors.

Data bias
The C 2 -Ray simulation used to train PINION is an imbalanced dataset depending on the redshift.For instance, the vast majority of pixels at high redshift, early reionisation history, are neutral rather than ionized.As such, it is possible that this biases the network predictions towards more neutral predictions.Any such bias would be most evident in the NP scenario, and we see for  > 7 in Figure 6 that the NP scenario always predicts delayed ionization, thus more neutral volumes compared to truth.Additionally, in comparison to the PFD and LD scenarios, the NP has larger neutral islands and smaller ionized bubbles.

Finite-difference instabilities
Instabilities coming from the finite-difference approximation used to calculate the derivative can be observed as a direct issue for the predictions.Although PINION uses a Runge-Kutta 4 scheme for differentiation, which has an error of the order O (ℎ 5 ), the time steps between each snapshot are large compared to what is required by the scheme.We expect the tension introduced by this bias to affect the LD and PFD scenarios, and we indeed observe that the NP scenario has the best agreement with C 2 -Ray for  < 7. The finite-difference instability issue is observed, as shown in Figure 3, but the maximal observed difference lies within 2 %.This is lower than the residual that we observe in Figure 6, where a standard deviation within 12 % is observed for  < 7.While instabilities of the derivative evaluation do introduce a slight bias, this does not account for the majority of the differences observed between C 2 -Ray and PINION.

Inaccuracies from smoothed source distribution field
The last main issue encountered in this project is the use of the smoothed source distribution field.As described in § 4.2, it is a way to provide additional information to the network, given that the two outputs from the N-Body simulation do not provide enough information about the growing structure of the ionized regions.However, as the smoothed field is generated by convolving uniform spheres with the input fields, it is much more uniform in structure compared to the behaviour of the ionization front and therefore the photoionization rate.The morphology of the ionization rate map and the smoothed source distribution field are compared in appendix § A, and in § B we test the influence of the smoothed source distribution field by adjusting the scale length of Equation 4 on the reionization history.

CONCLUSIONS
PINION successfully predicts the reionization history from RT simulation standard inputs with high accuracy.PINION could be applied to even larger volume simulations to create radiative transfer simulations of unprecedented volume and resolution.Additionally, the physics-informed L ODE drastically reduces the amount of data required to train the network, with almost identical performance between the LD and PFD scenarios.
PINION is relatively fast for a reionization simulation.Training the network for each scenario takes about 1.5 GPU-hours and the prediction for each scenario takes less than 84 GPU-hours.However, the prediction is an already highly parallelizable task that can be further improved.The data I/O is a major bottleneck in network training and prediction, and optimization was not explored as part of this work.For example, the prediction for a 300 3 volume is done in parallel on 1000 GPU nodes, each working on a smaller subvolume.In the current implementation, each node loads the entire dataset and then predicts all the pixels in the subvolume in approximately five minutes.This could easily be reduced to less than two minutes by only reading the subvolume.
While PINION is already a viable candidate for reproducing the C 2 -Ray simulations, there is still room for improvement.First, no hyperparameter optimization was performed for the training, so tweaks to the optimization algorithm or network structure may improve performance.Second, the smoothed field likely introduces strong systematic biases into the network's predictions.Starting from a more accurate proxy for the ionization front or the photoionization rate would result in more accurate predictions.Using an approximate but more physically-motivated RT simulation such as Grizzly as an input field may drastically improve performance while keeping computation times low.Finally, the data bias could be reduced by modifying the data loss to account for imbalanced data.
PINION is a promising candidate for almost unsupervised learning.The network requires very little input data to train, and the ODE loss only requires a realistic estimation of the photoionization rate for the calculation of Equation 5, which in this work we took instead the Γ from C 2 -Ray simulations.Moreover, we demonstrated that by providing a simplistic approach in Equation 4 for calculating the smoothed field, the network can retrieve information on the redshift evolution observed by the photoionization rate maps, Figure 2c.In principle, any mean free path model for the UV radiation could be used as a proxy in the smoothed field.xHII = 0.308 for PFD-NT(  = 0.5) and xHII = 0.370 for PFD-T(  = 0.5).These plots show that the network outputs are agnostic to a change in the mean free path scale as long as the network is retrained using the modified smoothed field.Indeed, the PSD (  = 1) case, presented in §5.1, and the corresponding retrained case results, PFD-T (  = 0.5), are very similar in these figures, indicating that the network can retrieve a reasonable ionization map as long as it is trained with some smoothed pre-processed field.For the predicted results without proper training, PFD-NT (  = 0.5), the reionization is, as expected, generally slower.The opposite behaviour was also observed for a factor  = 2 but is omitted for brevity.This indicates the importance of proper training on the pre-processed inputs for PINION and the robustness of PINION to particular choices of pre-processing strategies.

APPENDIX C: MASS-AVERAGED ANALYSIS
To complement the volume-averaged quantities analysis from § 5, we compared the main results from Figure 6 with their mass-averaged counterparts.The mass-averaged ionization fraction is defined as where  HII, is the volume-averaged ionization fraction and (1 +   ) is the gas overdensity.The summation term  is computed on the entire simulated volume.The comparison between the mass and volume-averaged quantities is shown in Figure C1.One can notice that the general behaviour of large volume simulation, where the mass-averaged tend to ionize earlier than the volume-averaged counterparts (Iliev et al. 2014).MNRAS 000, 1-13 (2022)

Figure 1 .
Figure1.Diagram of PINION's pipeline.The input data consist in two input subvolumes, the gas density and the source field that were produced by the N-Body simulation.The source field is then processed to obtain a smoothed field.The network outputs one value for the subvolume central pixel  HI .During the training steps, the output is then compared to the data to obtain the data loss L data and the physics constraint to obtain the the ODE loss L ODE .
(a) Evolution of the gas density: light blue are under-dense regions and dark blue are over-dense regions.(b) Evolution of the mass of sources: light green are regions with few particles, and dark green are the sources.(c) Evolution of the photoionization rate: light orange are regions which receives less ionizing photons and dark orange are the ionizing sources.(d) Evolution of the ionized fraction: red are ionized regions and dark blue are neutral regions.

Figure 2 .
Figure 2. Evolution of the gas density, mass of sources, photoionization rate, and ionized fraction for four different redshifts.Note that the scale is different for every plot but the evolution shows that Figure 2a and 2b are not expanding like Figure 2c and 2d.

Figure 3 .
Figure 3. Residual between the RK4 ODE evolution of  HII and the simulated  HII .This plot is obtained by averaging the evaluated ODE over 1,000 random subvolumes of size 47.62 Mpc.The redshift is cut at  = 10 to highlight the lower redshifts, where the residual is the biggest.

Figure 4 .
Figure 4. Visualization of a small cutout of the ionized hydrogen map at redshift  = 7.305 simulated by C 2 -Ray and the three PINION scenarios.The mean fraction of ionized hydrogen over the whole cube at this redshift is xHII = 0.389 for C 2 -Ray, xHII = 0.398 for PFD, xHII = 0.364 for NP and xHII = 0.391 for LD.The residual is computed as  =  PINION HII −  C 2 -Ray HII .Axes and colour bars are shared for readability.

Figure 5 .Figure 6 .
Figure 5. Slices through the redshift axis of the ionized hydrogen light-cone, simulated by C 2 -Ray and the three PINION scenarios.The colour bar indicates the ionization fraction.For readability, we cut at redshift  = 10 and position  ∈ [400, 600] Mpc.

Figure 8 .
Figure 8. Study of the ionization bubbles volume distribution using the friendof-friend method.Solid lines are C 2 -Ray, dashed lines are PFD training, dashdot lines are NP training and dotted lines are LD training.Both lines are picked at the same redshift.The blue line corresponds to approximately xHII = 0.3, the orange to xHII = 0.5 and the green to xHII = 0.7.

Figure 9 .
Figure 9. Top panel: Dimensionless power spectrum of  HI .Solid lines are C 2 -Ray, dashed lines are PFD training, dash-dot lines are NP training and dotted lines are LD training.These lines are picked at the same redshifts.The blue line corresponds to approximately xHI = 0.7, the orange to xHI = 0.5 and the green to xHI = 0.3.Bottom panel: Ratio plot for the above dimensionless power spectrum.The ratio is obtained by dividing the predicted power spectrum with C 2 -Ray's power spectrum.

Figure B1 .
Figure B1.Visualization of a small cutout of the ionized hydrogen map at redshift  = 7.305.The first two columns are the same as in Figure 4,  = 1 indicating no modification to the mean free path.The third and fourth columns correspond to the prediction using the  = 0.5 smoothed field without modifying the training and with a retrained network.Axes and colour bars are shared for readability.

Figure B2 .
Figure B2.Slices through the redshift axis of the ionized hydrogen light-cone.The colour bar indicates the ionization fraction.For readability, we cut at redshift  = 10 and position  ∈ [400, 600] Mpc.