SpeedyWeather.jl: Differentiable, GPU-accelerated and machine-learned climate modelling

Speaker: Milan Klöwer, University of Oxford

Authors: Milan Klöwer, Gregory Munday, Niklas Viebig, Maximilian Gelbrecht, Brian Groenke

Abstract

Climate models are traditionally hard to run, harder to customise, and harder still to interface with machine learning. We want to change that. In this talk we present SpeedyWeather.jl, a spectral atmospheric general circulation model written in Julia, made differentiable with Enzyme.jl and GPU-accelerated with KernelAbstractions and Reactant, with a redesigned parameterization system that makes physical processes easy to swap, calibrate, and replace with machine- learned components. We illustrate this with two applications: gradient-based calibration of shortwave radiation parameters against the observed global energy budget, and machine-learned surface roughness schemes over land, ocean and sea ice, trained offline but designed to generalise in space and to different climates. We first give an overview of the past year of development: differentiability, GPU support, and new parameterizations of radiation alongside simple sea ice, land and snow models that enable climate rather than only weather-timescale simulations. Regarding automatic differentiation, we now compute exact reverse-mode gradients of physical diagnostics — top-of-atmosphere radiation, surface fluxes — with respect to chosen parameters, with Enzyme propagating the loss back through the full atmospheric physics including radiative transfer and surface processes. We differentiate individual timesteps but batch across chaotic weather timescales to stabilise gradients, and train on continuously simulated data so slow processes can adapt despite single-timestep gradients, bypassing checkpointing. Since SpeedyWeather’s parameters are distributed across large nested structs, we implemented a parameter handling scheme for convenient updates and model reconstruction. Focusing on shortwave radiation, including cloud and surface albedos, we obtain improved agreement with the observed energy balance and a reproducible, extensible calibration workflow.