NuGNN: a Graph Neural Network for Nuclear Reaction Network Equations
Nuclear reaction networks are a major computational bottleneck in astrophysical simulations when large isotope sets are required, because of the stiffness of the network equations and the repeated calls to Jacobian-based solvers required by implicit methods. In this work, we develop a deep learning surrogate solver for a large 690-isotope nuclear reaction network under general Type I X-ray burst conditions using a graph neural network, NuGNN. Unlike conventional fully connected or convolutional neural networks, NuGNN directly reflects the structure of the reaction network through heterogeneous isotope and reaction nodes and message-passing along reaction connections. The model is trained on data spanning many orders of magnitude in stellar temperature and density and in simulation time step size. We compare NuGNN with a Res-U-Net and fully connected neural network and find that NuGNN consistently achieves significantly better accuracy with errors of only a few percent. More importantly, when implemented in the network evolution code in place of the original solver, NuGNN successfully reproduces the final abundance patterns, whereas the other architectures fail to do so. We also show that the trained model can substantially improve computational speed, demonstrating its practical potential for large-scale simulations. These results show that graph neural networks provide a robust and promising framework for accurate surrogate modeling of large nuclear reaction networks.