Emulating the interstellar medium chemistry with neural operators

TL;DR

This work introduces a DeepONet-based neural operator to emulated non-equilibrium ISM photo-chemistry for a nine-species network up to H$_2$ formation, enabling time evolution across varied densities, temperatures, and radiation fields while preserving temperature and species abundances. The model learns the operator mapping initial conditions and the radiation field to the solution, using a tensor-product branch-trunk architecture trained on a large KROME-generated dataset, achieving about $10^{-2}$ precision and a $128\\times$ speed-up over stiff ODE solvers. It outperforms prior neural approaches in accuracy and radiation-field flexibility, and demonstrates robustness to out-of-sample initial conditions and a physically relevant PDR benchmark. The emulator enables efficient coupling with hydrodynamic simulations and radiative transfer, offering a practical path toward on-the-fly, chemically detailed ISM modeling, while highlighting challenges in extrapolation over time and sharp transition regions for future refinement.

Abstract

Galaxy formation and evolution critically depend on understanding the complex photo-chemical processes that govern the evolution and thermodynamics of the InterStellar Medium (ISM). Computationally, solving chemistry is among the most heavy tasks in cosmological and astrophysical simulations. The evolution of such non-equilibrium photo-chemical network relies on implicit, precise, computationally costly, ordinary differential equations (ODE) solvers. Here, we aim at substituting such procedural solvers with fast, pre-trained, emulators based on neural operators. We emulate a non-equilibrium chemical network up to H$_2$ formation (9 species, 52 reactions) by adopting the DeepONet formalism, i.e. by splitting the ODE solver operator that maps the initial conditions and time evolution into a tensor product of two neural networks. We use $\texttt{KROME}$ to generate a training set spanning $-2\leq \log(n/\mathrm{cm}^{-3}) \leq 3.5$, $\log(20) \leq\log(T/\mathrm{K}) \leq 5.5$, $-6 \leq \log(n_i/n) < 0$, and by adopting an incident radiation field $\textbf{F}$ sampled in 10 energy bins with a continuity prior. We separately train the solver for $T$ and each $n_i$ for $\simeq 4.34\,\rm GPUhrs$. Compared with the reference solutions obtained by $\texttt{KROME}$ for single zone models, the typical precision obtained is of order $10^{-2}$, i.e. the $10 \times$ better with a training that is $40 \times$ less costly with respect to previous emulators which however considered only a fixed $\mathbf{F}$. The present model achieves a speed-up of a factor of $128 \times$ with respect to stiff ODE solvers. Our neural emulator represents a significant leap forward in the modeling of ISM chemistry, offering a good balance of precision, versatility, and computational efficiency.

Figures (7)

• Figure 1: Scheme of the emulator implemented in this work. The system Ordinary Differential Equations (ODE) describing the InterStellar Medium (ISM) chemical network (Sec. \ref{['sec:ISM_chemistry']}) is emulated via the DeepONet formalism (Sec. \ref{['sec:deeponet']}) by splitting the dependence i) from the initial conditions ($T$, and densities $n$ of each chemical species, i.e. $\mathrm{e}^-$, $\mathrm{H}^-$, $\rm H$, $\mathrm{H}^+$, He, $\mathrm{He}^+$, $\mathrm{He}^{++}$, $\mathrm{H}_2$, and $\mathrm{H}_2^+$) and the radiation flux ($\mathbf{F}$) with the branch network ($g$) and ii) from the temporal evolution in the time ($t$) domain with a trunk network ($f$); $f$ and $g$ are feed-forward neural networks, each consisting of 6 dense layers with 128 neurons: the tensor product ($\otimes$) of the branch and the trunk is adopted to compute the loss function (eq. \ref{['eq:UATfO']}). We individually train networks for the temperature and each of the chemical species; the data-set adopted to train DeepONet is described in Sec. \ref{['sec:dataset']} and its main properties are summarized in Tab. \ref{['tab:data_set_structure']}.
• Figure 2: Predicted vs true test for the DeepONet model. Logarithmic densities for each ion (H, H$^{+}$, ...) and the temperature (T) are normalized in the full data-set range (Tab. \ref{['tab:data_set_structure']}) and summed ($y$, see eq. \ref{['eq:def_sum_ion_and_t']}) for both the true value from KROME and the predicted value from DeepONet. The image show the 2D probability distribution function (PDF) of the summed dataset, and it is normalized such that the maximum is 1 to better appreciate the dynamical range. To guide the eye, we have added a dashed black line to mark the $\rm KROME= DeepONet$ region. See Fig. \ref{['fig:pred_vs_true_individual']} for the same diagnostic for individual ions.
• Figure 3: PDF of the relative error ($\Delta_r$) for the testing set. Each output from the emulator ($T$ and density for each species) is shown independently, with the color-code indicated in the legend. The PDFs are computed using a testing set of $\simeq 2 \times 10^7$ of points.
• Figure 4: Examples of the time ($t$) evolution of temperature ($T$) and the density ($n$, see the legend) of all the species in the chemical network. The solid lines represent the solutions computed using KROME, while the dashed lines depict the predictions of our models, with each line being a single solution from the emulator. In the left (right) panel the gas number density is $n=10^4\mathrm{cm}^{-3}$ (initial temperature is $T=10^6\mathrm{K}$), i.e. outside the range of the training dataset (see Tab. \ref{['tab:data_set_structure']}).
• Figure 5: Photo-dissociation region (PDR) benchmark. Temperature ($T$, right axis) and density of main hydrogen species (H, H$^{+}$, and H$_{2}$, left axis) profiles as a function of column density ($N$) obtained for an impinging radiation flux of $10\,G_0$ propagating in a $Z={\rm Z}_{\odot}$ slab (see Sec. \ref{['sec:pdr_test']} for details of the model). Solid lines represent the numerical solutions computed with KROME at equilibrium ($t\sim 1\,\mathrm{Kyr}$), dashed lines represent the solution predicted by the emulator. Note that each point in the profile at a given $N$ is treated independently by the emulator.