.. title:: Talk: SpeedyWeather.jl

.. only:: html

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.