Maximilian Gelbrecht
Researching differentiable programming and machine learning for Earth system models and dynamical systems
Sessions
Traditionally, climate models are difficult to run for end users, and even harder to customize or interface with machine learning. We want to change that. Here, we present the ongoing development of SpeedyWeather.jl: A general circulation model that’s differentiable, GPU-capable and ready for machine learning integration. SpeedyWeather.jl is a spectral atmospheric general circulation model with an everything-flexible attitude. In this talk, we will give an overview of SpeedyWeather.jl’s development of the last year, in which we worked towards differentiability with Enzyme, GPU-capability with KernelAbstractions and Reactant and rewrote our parametrizations for better performance and more customisability.
Global land surface and hydrological models are crucial components of Earth System Models (ESMs). In addition to providing realistic boundary conditions for the atmosphere and ocean components, they also play a key role in understanding Earth’s changing energy imbalance and the response of the terrestrial carbon and water cycles to anthropogenic climate change. Unlike atmosphere and ocean models, however, land models lack a fluid dynamical core and rely heavily on empirical parameterizations to represent many key processes. As such, there is a continued need for a new generation of land models which can facilitate the incorporation of data-driven components. Here we present Terrarium.jl, a Julia-based land modeling framework for GPU-accelerated and automatically differentiable simulations of soil, snow, and vegetation dynamics, along with their corresponding land-atmosphere exchange fluxes. We highlight how Julia’s key features enable unprecedented modularity in the model design and seamless GPU parallelization through KernelAbstractions.jl. We further demonstrate the value of GPU acceleration and differentiability through a series of performance benchmarks and sensitivity analyses. We also detail our initial experiments in achieving stable coupling to a reduced-complexity atmosphere model, SpeedyWeather.jl.
Hybrid climate modelling combines numerical models with machine-learned components. We present the development of multiple machine-learned surface climate processes and their integration into the climate model SpeedyWeather.jl using PyTorch and Lux.jl. Despite the offline training, the hybrid model is designed to generalise in space and to different climates. We address speed vs. accuracy tradeoffs using SymbolicRegression.jl and discuss online learning with Enzyme.jl.
Fortran climate models are being adapted to GPUs by automatically translating loop-by-loop into a kernel. In Julia, we have more flexibility to develop the climate model SpeedyWeather.jl for the GPU. Many parts are easy to accelerate, leverage multiple dispatch on the GPU and a high level of kernel fusion for modularity and performance, while being optionally hardware-specific. The spherical harmonic transforms remain a complex bottleneck but we employ a multi-algorithm approach with custom linear algebra kernels using Reactant, Fourier and Legendre transforms.
Climate models rely on parameterizations that are traditionally tuned manually. We present a differentiable calibration framework using Enzyme.jl to compute exact reverse mode gradients of energy-balance diagnostics in SpeedyWeather.jl. By batching single- timestep gradients across chaotic dynamics, we enable systematic, reproducible optimization of shortwave radiation parameters, establishing an extensible workflow for objective calibration in Earth system models.