Bridging Machine Learning and Cosmological Simulations: Using Neural Operators to emulate Chemical Evolution

Abstract

The computational expense of solving non-equilibrium chemistry equations in astrophysical simulations poses a significant challenge, particularly in high-resolution, large-scale cosmological models. In this work, we explore the potential of machine learning, specifically Neural Operators, to emulate the Grackle chemistry solver, which is widely used in cosmological hydrodynamical simulations. Neural Operators offer a mesh-free, data-driven approach to approximate solutions to coupled ordinary differential equations governing chemical evolution, gas cooling, and heating. We construct and train multiple Neural Operator architectures (DeepONet variants) using a dataset derived from cosmological simulations to optimize accuracy and efficiency. Our results demonstrate that the trained models accurately reproduce Grackle's outputs with an average error of less than 0.6 dex in most cases, though deviations increase in highly dynamic chemical environments. Compared to Grackle, the machine learning models provide computational speedups of up to a factor of six in large-scale simulations, highlighting their potential for reducing computational bottlenecks in astrophysical modeling. However, challenges remain, particularly in iterative applications where accumulated errors can lead to numerical instability. Additionally, the performance of these machine learning models is constrained by their need for well-represented training datasets and the limited extrapolation capabilities of deep learning methods. While promising, further development is required for Neural Operator-based emulators to be fully integrated into astrophysical simulations. Future work should focus on improving stability over iterative timesteps and optimizing implementations for hardware acceleration. This study provides an initial step toward the broader adoption of machine learning approaches in astrophysical chemistry solvers.