Systematic selection of surrogate models for nonequilibrium chemistry

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.

17 pages, 9 figures, submitted to Astronomy & Astrophysics

References in corpus (25)

• DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators
• A KInetic Database for Astrochemistry (KIDA)
• Grackle: a Chemistry and Cooling Library for Astrophysics
• Uncertainties in H2 and HD Chemistry and Cooling and their Role in Early Structure Formation
• On the Relationship Between Molecular Hydrogen and Carbon Monoxide Abundances in Molecular Clouds
• KROME - a package to embed chemistry in astrophysical simulations
• Modeling UV Radiation Feedback from Massive Stars: II. Dispersal of Star-Forming Giant Molecular Clouds by Photoionization and Radiation Pressure
• The circumgalactic medium of massive galaxies at z~3: a test for stellar feedback, galactic outflows, and cold streams
• UCLCHEM: A Gas-Grain Chemical Code
• The impact of chemistry on the structure of high-z galaxies
• Is H3+ cooling ever important in primordial gas?
• CO line emission from galaxies in the Epoch of Reionization
• Photoevaporation of Jeans-unstable molecular clumps
• A chemical model for the interstellar medium in galaxies
• Catastrophic Cooling in Superwinds: Line Emission and Non-equilibrium Ionization
• Reducing the complexity of chemical networks via interpretable autoencoders
• Reaction Studies of Neutral Atomic $\mathrm{C}$ with ${\mathrm{H}}_3^{+}$ using a Merged-Beams Apparatus
• Chemulator: Fast, accurate thermochemistry for dynamical models through emulation
• H2 chemistry in galaxy simulations: an improved supernova feedback model
• Catastrophic Cooling in Superwinds. III. Non-equilibrium Photoionization
• Neural networks: solving the chemistry of the interstellar medium
• The Astrochemical Evolution of Turbulent Giant Molecular Clouds : I - Physical Processes and Method of Solution for Hydrodynamic, Embedded Starless Clouds
• NeuralPDR: Neural Differential Equations as surrogate models for Photodissociation Regions
• Bridging Machine Learning and Cosmological Simulations: Using Neural Operators to emulate Chemical Evolution
• Neural-Network Chemical Emulator for First-Star Formation: Robust Iterative Predictions over a Wide Density Range