Differentiable Programming for Atmospheric Models: Experiences and Perspectives

Differentiable programming enables automatic differentiation (AD) tools to compute gradients through code without manually defining derivatives. AD tools can differentiate through entire software stacks, composing many functions and algorithms via the chain rule. With models that incorporate differentiable programming long-standing challenges like systematic calibration, comprehensive sensitivity analyses, and uncertainty quantification can be tackled, and machine learning (ML) methods can be integrated directly into the process-based core of earth-system models (ESMs) to incorporate additional information from observations. Through the advent of ML, several AD tools are gaining traction. A new generation of powerful tools like JAX, Zygote and Enzyme enable differentiable programming for models of varying complexity including highly complex coupled ESMs. Here we present an overview about the perspectives of differentiable programming for ESMs, using experience from two of our applications in atmospheric modelling. First off, we set up PseudoSpectralNet, a differentiable quasi-geostrophic model in Julia with Zygote. This is a hybrid model combining neural networks with a dynamical core, showcasing how the stability and accuracy of ML models is improved by integrating a process-based dynamical core into our model. Additionally, ongoing work uses Enzyme to achieve a differentiable version of the significantly more complex SpeedyWeather.jl atmospheric model. We will discuss advantages of both approaches and give an outlook into future possibilities with differentiable models.