Systematic selection of surrogate models for nonequilibrium chemistry

TL;DR

This paper introduces CODES, a framework for systematically benchmarking and optimizing neural surrogate models for nonequilibrium chemistry, demonstrating trade-offs between accuracy and efficiency.

Contribution

The authors present CODES, a new systematic benchmarking framework for astrochemical surrogate models, including datasets, metrics, and optimization procedures.

Findings

Fully connected models achieve highest accuracy and reliable uncertainty estimates.

Latent-evolution models show improved robustness under iterative prediction.

Systematic optimization reveals pronounced accuracy-efficiency trade-offs.

Abstract

Nonequilibrium chemistry is central to many astrophysical environments but remains a major computational bottleneck in simulations because solving the associated stiff ODE systems is expensive. Neural surrogates promise large speedups, yet existing studies rarely provide systematic comparisons of architectures or rigorous optimization toward both accuracy and efficiency. We introduce CODES, a principled framework for optimizing and benchmarking astrochemical surrogate models. Using CODES, we compare four neural surrogate architectures across four KROME-generated datasets spanning primordial and molecular-cloud chemistry with up to 287 reactions across 37 species. Dual-objective optimization reveals pronounced accuracy-efficiency trade-offs across architectures. Fully connected models achieve the highest accuracy and most reliable uncertainty estimates, while latent-evolution models show improved robustness under iterative prediction. Our results highlight the importance of systematic optimization and architectural comparison. The datasets, metrics, and benchmarking procedure are publicly released within CODES to enable reproducible surrogate benchmarking.