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

Inclusion of atmospheric chemistry in global climate projections is currently limited by the high computational expense of modelling the chemical transport. For example, in the latest IPCC Assessment Report (AR6), two-thirds of the CMIP6 models used either table-driven species values, or no chemistry at all.

Recent rapid advancements in artificial intelligence offer us new tools for reducing the cost of numerical prediction. The application of these tools to the modelling of atmospheric chemistry is still somewhat nascent and there remain multiple challenges due to the reaction complexities and the high dimensionality of chemical species. For climate timescales, model stability and regional biases are also limiting factors.

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 a standard training and testing dataset based on 8 years of chemistry simulations from GEOS-Chem, plus meteorological data from MERRA2. The framework allows for both offline and coupled online training methods.

We also introduce a new data-driven model called GAIANet. GAIANet is a transformer-based architecture that uses Fourier-space neural operators for resolution independent token mixing. We use GAIA-Chem to compare two DNN models; a standard autoencoder based on Convolutional LSTM nodes, and our GAIANet model. We demonstrate the benefit of enhanced model fidelity in reducing predicted error values and discuss the merits of transformer-based approaches for AI modelling of atmospheric chemistry.