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

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.