Differentiable Atmospheric Modelling with SpeedyWeather.jl

The current generation of hybrid machine learning and physics-informed machine learning is often limited by the missing availability of comprehensive differentiable models: either strongly simplified models have to be used or machine learning (ML) can’t be integrated natively into process-based models and must be trained separately. Here, we present the ongoing development of SpeedyWeather.jl: A general circulation model that’s differentiable, GPU-capable and ready for ML simulations. SpeedyWeather.jl is a spectral atmospheric GCM with a primitive equation core on flexible grid implementations from Gaussian to HEALPix. It contains simple yet interactive representations of ocean, land and sea ice for coupled climate simulations. With a user interface made for modularity and interactivity, it’s ideally suited as a framework for hybrid atmospheric models. For example, new parameterizations can be defined without any lines of code for GPU or differentiability specifics, yet integrate seamlessly into those. We document the process to achieve differentiability of our model using the general purpose automatic differentiation library Enzyme, problems we encountered and solutions we found. We demonstrate the differentiability with a sensitivity analysis of our model, initial developments of data-driven parameterizations, and give an outlook on the development of differentiable Earth system models.