AMORE: Adaptive Multi-Output Operator Network for Stiff Chemical Kinetics

TL;DR

AMORE advances surrogate modeling for stiff chemical kinetics by learning a single neural operator that predicts multiple thermo-chemical states from initial conditions. It integrates adaptive, gradient-free loss functions, Partition of Unity constraints, and a mass-conserving coordinate transformation, plus a two-step training paradigm to improve multi-output accuracy. Demonstrations on Syngas and GRI-Mech 3.0 show substantial gains in predictive accuracy, robustness to out-of-distribution inputs, and viable recursive predictions, with a generalization example using Fourier Neural Operator. The framework lays groundwork for fast, scalable operator-CFD coupling, offering practical speedups for turbulent combustion simulations while ensuring mass conservation and controllable extrapolation behavior.

Abstract

Time integration of stiff systems is a primary source of computational cost in combustion, hypersonics, and other reactive transport systems. This stiffness can introduce time scales significantly smaller than those associated with other physical processes, requiring extremely small time steps in explicit schemes or computationally intensive implicit methods. Consequently, strategies to alleviate challenges posed by stiffness are important. While neural operators (DeepONets) can act as surrogates for stiff kinetics, a reliable operator learning strategy is required to appropriately account for differences in the error between output variables and samples. Here, we develop AMORE, Adaptive Multi-Output Operator Network, a framework comprising an operator capable of predicting multiple outputs and adaptive loss functions ensuring reliable operator learning. The operator predicts all thermochemical states from given initial conditions. We propose two adaptive loss functions within the framework, considering each state variable's and sample's error to penalize the loss function. We designed the trunk to automatically satisfy Partition of Unity. To enforce unity mass-fraction constraint exactly, we propose an invertible analytical map that transforms the $n$-dimensional species mass-fraction vector into an ($n-1$)-dimensional space, where DeepONet training is performed. We consider two-step training for DeepONet for multiple outputs and extend adaptive loss functions for trunk and branch training. We demonstrate the efficacy and applicability of our models through two examples: the syngas (12 states) and GRI-Mech 3.0 (24 active states out of 54). The proposed DeepONet will be a backbone for future CFD studies to accelerate turbulent combustion simulations. AMORE is a general framework, and here, in addition to DeepONet, we also demonstrate it for FNO.

Figures (32)

• Figure 1: AMORE: A schematic diagram of the proposed Adaptive Multi-output Operator networks (AMORE) framework. It consists of an operator network that is capable of predicting multiple outputs and adaptive loss functions. In the present study, we consider Deep Operator Network (DeepONet) with two different architectures as the choice of the operator network, the DeepONet and DeepOKAN. We propose two adaptive loss functions that are updated using a gradient-free method based on the relative error between the state variables and samples.
• Figure 2: DeepONet for stiff chemical kinetics: Schematic of the DeepONet prediction scheme considered in the present study. In the left middle, within the red dashed box, we have shown the prediction of the state variable using the pre-trained DeepONet parametrized with $\bm{\Psi}$. The input to the DeepONet is the initial condition of the state variables, $\bm{Y}_0$. The output is the predicted dynamics of the state variables, $\bm{Y}_0, \bm{Y}_1, \bm{Y}_2,\cdots\bm{Y}_{n_t}$. At the top right, we have shown the DeepONet learning process. It solves the initial value problem of the stiff system (\ref{['Eq:ODE general']} by mapping the initial condition to the dynamics of the state variables. Generally, the input for the DeepONet, the initial conditions, comes from the other associated solver, and the output of the DeepONet at a particular time instant goes to the other associated solver as input. Since we have not considered the DeepONet-CFD/PINN integration in the present study, we consider the time decomposition of the dynamic equation as shown at the bottom of the figure. In the time decomposition conditions, it is assumed that the initial conditions of the ODEs are known after every $t=t_{n_t}=n_t\Delta t$, which ideally comes from the associated solver. Thus, we will design our DeepONet to predict the dynamics considering time decomposition, assuming the initial conditions are known at every $t=t_{n_t} = n_t\Delta t$. Furthermore, when the DeepONet output is not sufficiently progressed in time to be input to the associated solver, we propose to consider a recursive prediction strategy as shown on the right within the black dashed box with the hatch. The output of the DeepONet at time instant $t=t_{n_t}=n_t\Delta t$ is considered as the input to the pretrained DeepONet, and the process continues till sufficient progress in time is achieved for the chemical kinetics to be input to the associated solver.
• Figure 3: Schematic diagram of the proposed DeepONet architecture showing branch and trunk networks and multiple outputs. We have shown the branch network in the top left in the red dashed box. The inputs to the branch are the initial conditions of the state variables $\bm{Y}_0\in \mathbb{R}^{bs\times j}$ and the outputs of the branch are $br(\bm{\Theta})\in \mathbb{R}^{bs\times P}$. Where $bs$ and $j$ are the number of samples and the number of state variables, respectively, and $P$ is the number of output neurons. The input to the trunk is time $t\in\mathbb{R}^{(n_t+1)\times 1}$ and the outputs of the trunk are $tr(\bm{\Phi})\in \mathbb{R}^{(n_t+1)\times P}$. Where $(n_t+1)$ is the number of time points where the output needs to be predicted. In the present study, we have considered ResNet (with trainable parameters $\bm{\Theta}$) as a choice of network architecture for the branch network and ResNet or KAN (with trainable parameters $\bm{\Phi}$) as a choice of network architecture for the trunk network. A brief discussion on the ResNet and KAN architectures is included in \ref{['Appedix:Network Architecture']}. The branch and trunk network each has $P = j\times p$ neurons in the output layers, where $j$ is the number of state variables and $p$ is the number of basis functions and coefficients considered to approximate each variable. The output $\bm{Y}\in\mathbb{R}^{bs\times (n_t+1)\times j}$ is given by the dot product between the output of the branch and trunk networks, and is shown on the right side in standard form and tensor form. We also implemented the partition of unity in the trunk network, and the implementation is discussed in \ref{['Subsubsection:PoU']}. Furthermore, we also consider a loss due to the conservation of mass along with data loss, which is discussed in \ref{['Subsubsection:CoM']}. As discussed in the AMORE framework, we propose two adaptive loss functions, and these are discussed in \ref{['Subsection:Adaptive loss function']}.
• Figure 4: Mass conserving map showing the species mass fraction vector ($\bm{y} \in \mathbb{R}^n$) is mapped uniquely to a mass-conserving coordinate ($\bm{z} \in \mathbb{R}^{n-1}$) and vice versa.
• Figure 5: Syngas problem, Violin plot of the relative $L_2$ error of the reconstructed prediction of the state variables of the test dataset. Here, we have shown only four methods; the violin plots for all the methods considered are shown in \ref{['Figure:Syngas:Violin plot all 9 method']} (in \ref{['Appendix:Subsection:Syngas problem']}). We observed that the predicted results using DON-NA have the highest mean of the relative $L_2$ error compared to the other methods. Furthermore, DON-NA also has a higher standard deviation of the relative $L_2$ error. The mean and standard deviation are reduced as adaptive loss functions are considered. We would particularly like to mention that the state variables corresponding to the dynamics of HCO and HO$_2$ shown in (k) and (h), respectively, have higher errors, which are reduced when we consider the adaptive loss function. We have observed that the DON-Ad-B and 2S-Ad-B have similar accuracy; however, 2S-Ad-B has smaller standard deviations. The relative $L_2$ errors are calculated as discussed in \ref{['Subsection:Numerical:Syngas:Result analysis']}.