GAIA-Chem: A Framework for Global AI-Accelerated Atmospheric Chemistry Modelling

The inclusion of atmospheric chemistry in global climate projections is currently limited by the high computational expense of modelling the many reactions of chemical species. Recent rapid advancements in artificial intelligence (AI) provide us with new tools for reducing the cost of numerical simulations. The application of these tools to atmospheric chemistry is still somewhat nascent and multiple challenges remain due to the reaction complexities and the high number of chemical species. In this work, we present GAIA-Chem, a global AI-accelerated atmospheric chemistry framework for large-scale, multi-fidelity, data-driven chemical simulations. GAIA-Chem provides an environment for testing different approaches to data-driven species simulation. GAIA-Chem includes curated training and validation datasets, support for offline and online training schemes, and comprehensive metrics for model intercomparison. We use GAIA-Chem to evaluate two Deep Neural Network (DNN) models: a standard autoencoder scheme based on convolutional LSTM nodes, and a transformer-based model. We show computational speedups of 1,280x over numerical methods for the chemical solver and a 2.8x reduction in RMSE when compared to previous works.