Jax-esm: a differentiable coupler for jax-based Earth system models

The differentiability of numerical climate models exhibits  many advantages over non-differentiable models. Differentiable climate models would be able to optimize parameters and quickly solve for climate equilibrium. They can also be used to find unstable climate equilibrium states that are impossible to identify in time-forwarding models. Differentiability also enables sensitivity studies, such as the impact of initial conditions on predictions, which is the key concept in the 4-dimensional variational method. Finally, differentiable ability also integrates well with the trending data-driven artificial intelligence model, such as NeuralGCM.  

Currently, physics-based differentiable coupled climate models are still rare. Some existing ones include: ECMWF Integrated Forecasting System (ECMWF-IFS) and Coupled Ocean/Atmosphere Mesoscale Prediction System (COAMPS). The high scientific value of such a tool warrants development of further differentiable modelling systems.

In this work, we present jax-esm, a differentiable coupler for models written in Python with the JAX framework. JAX is a Python library developed by Google that builds on NumPy and adds automatic differentiation and just-in-time (JIT) compilation. It has been used to develop atmospheric models such as NeuralGCM and jax-gcm. In this example, we couple jax-gcm, a JAX-based atmosphere intermediate model, to a slab ocean model. We demonstrate the optimization of ocean mixed-layer depth and solving for climate equilibrium through differentiability.