Neural Network Emulator for Atmospheric Chemical ODE

TL;DR

Atmospheric chemistry modeling is computationally expensive when predicting the time evolution of many interacting species. The authors introduce ChemNNE, an attention-based neural ODE emulator that integrates sinusoidal time embeddings and Fourier Neural Operators to efficiently model chemical dynamics, trained on a large ARCA-box–derived dataset. Physics-informed losses enforce mass conservation and trajectory consistency, delivering state-of-the-art accuracy with substantial speedups over traditional solvers and baselines. The approach enables near real-time, high-resolution forecasts of multi-species chemistry, with implications for climate modelling and air-quality prediction.

Abstract

Modeling atmospheric chemistry is complex and computationally intense. Given the recent success of Deep neural networks in digital signal processing, we propose a Neural Network Emulator for fast chemical concentration modeling. We consider atmospheric chemistry as a time-dependent Ordinary Differential Equation. To extract the hidden correlations between initial states and future time evolution, we propose ChemNNE, an Attention based Neural Network Emulator (NNE) that can model the atmospheric chemistry as a neural ODE process. To efficiently simulate the chemical changes, we propose the sinusoidal time embedding to estimate the oscillating tendency over time. More importantly, we use the Fourier neural operator to model the ODE process for efficient computation. We also propose three physical-informed losses to supervise the training optimization. To evaluate our model, we propose a large-scale chemical dataset that can be used for neural network training and evaluation. The extensive experiments show that our approach achieves state-of-the-art performance in modeling accuracy and computational speed.

Figures (8)

• Figure 1: The proposed ChemNNE for chemical concentration prediction. It takes the environmental parameters and chemical initial concentration to predict the future chemical reaction process.
• Figure 2: Three chemical prediction task for model evaluation.
• Figure 3: Visualization of the time evolution of chemistry in Task 1, 2 and 3. In (a), (b), and (c), we show the ground truth as red lines and the predictions with different colors. We pick four different chemical compounds for comparison, and we also enlarge the region in red boxes to highlight the prediction errors.
• Figure 4: Visualization of the mean errors of the time evolution in Task 1, 2 and 3. We average all chemical compounds and show the mean absolute errors between ground truth and other predictions across different time steps.
• Figure 5: Training loss comparison among ours and others. We show different approaches in different colors. We enlarge the red-boxed region and display it at the upper center. We can see the improvements of using our proposed ChemNNE in both convergence speed and final loss.