SWIFT-AI: Significant Speed-up in Modelling the Stratospheric Ozone Layer

Common representations of the stratospheric ozone layer in climate modeling are widely considered only in a very simplified way. Neglecting the mutual interactions of ozone with atmospheric temperature and dynamics has the effect of making climate projections less accurate. Although, more elaborate and interactive models of the stratospheric ozone layer are available, they require far too much computation time to be coupled with climate models. Our aim with this project was to break new ground and pursue an interdisciplinary strategy that spans the fields of machine learning, atmospheric physics and climate modelling.

In this work, we present an implicit neural representation of the extrapolar stratospheric ozone chemistry (SWIFT-AI). An implicitly defined hyperspace of the stratospheric ozone chemistry offers a continuous and even differentiable representation that can be parameterized by artificial neural networks. We analysed different parameter-efficient variants of multilayer perceptrons. This was followed by an intensive, as far as possible energy-efficient search for hyperparameters involving Bayesian optimisation and early stopping techniques.

Our data source is the Lagrangian chemistry and transport model ATLAS. Using its full model of stratospheric ozone chemistry, we focused on simulating a wide range of stratospheric variability that will occur in future climate (e.g. temperature and meridional circulation changes). We conducted a simulation for several years and created a data-set with over 200E+6 input and output pairs. Each output is the 24h ozone tendency of a trajectory. We performed a dimensionality reduction of the input parameters by using the concept of chemical families and by performing a sensitivity analysis to choose a set of robust input parameters.

We coupled the resulting machine learning models with the Lagrangian chemistry and transport model ATLAS, substituting the full stratospheric chemistry model. We validated a two-year simulation run by comparing to the differences in accuracy and computation time from both the full stratospheric chemistry model and the previous polynomial approach of extrapolar SWIFT. We found that SWIFT-AI consistently outperforms the previous polynomial approach of SWIFT, both in terms of test data and simulation results. We discovered that the computation time of SWIFT-AI is more than twice as fast as the previous polynomial approach SWIFT and 700 times faster than the full stratospheric chemistry scheme of ATLAS, resulting in minutes instead of weeks of computation time per model year – a speed-up of several orders of magnitude.

To ensure reproducibility and transparency, we developed a machine learning pipeline, published a benchmark dataset and made our repository open to the public.

In summary, we could show that the application of state-of-the-art machine learning methods to the field of atmospheric physics holds great potential. The achieved speed-up of an interactive and very precise ozone layer enables a novel way of representing the ozone layer in climate models. This in turn will increase the quality of climate projections, which are crucial for policy makers and of great importance for our planet.