Neural-Network Chemical Emulator for First-Star Formation: Robust Iterative Predictions over a Wide Density Range

TL;DR

This paper introduces a neural-emulator for the thermochemical evolution during Population III star formation that remains accurate across a wide density range from $n_{ m H}=10^{-3}$ to $10^{18}\ \mathrm{cm^{-3}}$. It uses region-specific DeepONet models to predict temperature and six primordial species and couples this with a novel timescale-based update to ensure robust iterative predictions, aided by Deep Galerkin Method–based estimates of the characteristic timescale $\tau_\epsilon$. The emulator achieves sub-10% relative errors for over 90% of cases (except very rare species like $\mathrm{H_2^+}$) and delivers substantial speedups: about $\sim$10× on CPU and $>10^3$× for batched GPU predictions, enabling efficient integration into hydrodynamic simulations. Validation includes single-step tests, one-zone collapse, and short-timestep iterative runs, showing reliable performance even under challenging, highly stiff conditions. This work lays the groundwork for GPU-accelerated, NN-based chemical evolution in star-formation simulations and outlines directions for incorporating metals and radiation in future extensions.

Abstract

We present a neural-network emulator for the thermal and chemical evolution in Population III star formation. The emulator accurately reproduces the thermochemical evolution over a wide density range spanning 21 orders of magnitude (10$^{-3}$-10$^{18}$ cm$^{-3}$), tracking six primordial species: H, H$_2$, e$^{-}$, H$^{+}$, H$^{-}$, and H$_2^{+}$. To handle the broad dynamic range, we partition the density range into five subregions and train separate deep operator networks (DeepONets) in each region. When applied to randomly sampled thermochemical states, the emulator achieves relative errors below 10% in over 90% of cases for both temperature and chemical abundances (except for the rare species H$_2^{+}$). The emulator is roughly ten times faster on a CPU and more than 1000 times faster for batched predictions on a GPU, compared with conventional numerical integration. Furthermore, to ensure robust predictions under many iterations, we introduce a novel timescale-based update method, where a short-timestep update of each variable is computed by rescaling the predicted change over a longer timestep equal to its characteristic variation timescale. In one-zone collapse calculations, the results from the timescale-based method agree well with traditional numerical integration even with many iterations at a timestep as short as 10$^{-4}$ of the free-fall time. This proof-of-concept study suggests the potential for neural network-based chemical emulators to accelerate hydrodynamic simulations of star formation.

Figures (7)

• Figure 1: Schematic diagram of DeepONet, which consists of two networks: the branch net and the trunk net. The branch net receives the hydrogen nucleus number density $n_\mathrm{H}$, the initial temperature $T_{0}$, and the initial chemical abundances $\bm{y}_{0}$ as input, in the form of $\log (n_\mathrm{H}/\mathrm{cm}^{-3})$, $\log (T_{0}/\mathrm{K})$, and $\log (\bm{y}_{0})$, respectively. The trunk net takes the timestep $\Delta t$ as input in the form of $\log\bigl(1 + \Delta t/\mathrm{sec}\bigr)$. The final output is computed as the inner product of the branch net output $\bm{B}_\mathrm{net}$ and the trunk net output $\bm{T}_\mathrm{net}$, plus a bias term $b$: $\bm{B}_\mathrm{net}\cdot\bm{T}_\mathrm{net}+b$. A separate DeepONet model is constructed for each target variable, which is either the updated temperature $T$ in the form of $\log (T/\mathrm{K})$ or the updated chemical abundance of the $i$-th species $y(i)$ in the form of $\log (y(i))$.
• Figure 2: Single-step DeepONet predictions of the thermal and chemical evolution for a fixed hydrogen number density of $n_\mathrm{H} = 10^{12}\,\mathrm{cm}^{-3}$. The top panel compares the DeepONet predictions (dashed) with the ground-truth results from numerical integrations (solid). The bottom panel plots the corresponding residual errors normalized by the ground-truth values, i.e., $\left(T_\mathrm{pred}-T_\mathrm{true}\right)/T_\mathrm{true}$ or $\left(y_\mathrm{pred}(i)-y_\mathrm{true}(i)\right)/y_\mathrm{true}(i)$, where the subscripts "pred" and "true" denote the DeepONet and ground-truth values, respectively. The horizontal axis represents the transformed timestep, $\log(1 + \Delta t/\mathrm{sec})$. The DeepONet predictions almost perfectly reproduce the ground-truth results. The variables are color-coded as follows: $T$ (magenta); $y(\mathrm{H})$ (green); $y(\mathrm{H}_2)$ (cyan); $y(\mathrm{e}^-)$ (orange); $y(\mathrm{H}^+)$ (yellow); $y(\mathrm{H}^-)$ (blue); and $y(\mathrm{H}_2^+)$ (red). Note that the curves for $y(\mathrm{e}^-)$ are nearly invisible because they almost completely coincide with those for $y(\mathrm{H}^+)$.
• Figure 3: Cumulative distribution functions of the relative prediction errors of the DeepONet emulator for each density region. The horizontal axis shows the logarithm of relative error $\Delta_Y$, defined in Eq. (\ref{['eq:delta_Y']}), and the vertical axis indicates the cumulative relative frequency. The two blue dotted vertical lines show $\Delta_Y=0.05$ (left) and $0.1$ (right). Each panel corresponds to a different density range: Region 0 ($10^{-3}\leq n_\mathrm{H}/\mathrm{cm^{-3}}\leq10^4$, top left); Region 1 ($10^4\leq n_\mathrm{H}/\mathrm{cm^{-3}}\leq10^8$, top right); Region 2 ($10^8\leq n_\mathrm{H}/\mathrm{cm^{-3}}\leq10^{12}$, middle left); Region 3 ($10^{12}\leq n_\mathrm{H}/\mathrm{cm^{-3}}\leq10^{15}$, middle right); and Region 4 ($10^{15}\leq n_\mathrm{H}/\mathrm{cm^{-3}}\leq10^{18}$, bottom left). Colors are the same as in Figure \ref{['fig:single_cell_nocomp']}.
• Figure 4: One-zone evolution with timestep $\Delta t = 10^{-1}\,t_\mathrm{ff}$. The horizontal axis is the hydrogen number density $n_\mathrm{H}$, corresponding to the time evolution during gravitational collapse. The top panel shows the temperature evolution, the middle panel shows the chemical evolution, and the bottom panel shows the residual errors. Blue dotted vertical lines indicate the boundaries between different density regions. Colors are the same as in Figure \ref{['fig:single_cell_nocomp']}. The solid lines show the ground-truth numerical integration results, while the dashed lines show the DeepONet predictions for $n_{\mathrm{H}} \leq 10^{18}\,\mathrm{cm^{-3}}$ and the results obtained by interpolating a precalculated table for $n_{\mathrm{H}} \ge 10^{18}\,\mathrm{cm^{-3}}$ (see Appendix \ref{['sec:app-table-equilibrium']}). Note that the apparent jump in the abundances of e$^-$, H$^+$, H$^-$, and H$_2^+$ at $n_{\mathrm{H}} = 10^{18}\,\mathrm{cm^{-3}}$, arising due to the absence of inverse reactions in our chemical network, does not affect the thermal evolution Nakauchi:2019. The abundances of $\mathrm{e}^-$ and $\mathrm{H}^+$ overlap almost completely with each other.
• Figure 5: Same as Figure \ref{['fig:single_cell_CDF']}, but for the characteristic timescales $\tau_\epsilon$ with $\epsilon=0.1$. The horizontal axis shows the logarithm of $\Delta_\tau$ as defined in Eq. \ref{['eq:delta_chemcool']}.