Milan Klöwer
Milan Klöwer is a NERC Independent Research Fellow at the University of Oxford. He did his postdoc at the Massachusetts Institute of Technology (MIT) working on climate model development in Julia. He started SpeedyWeather.jl, a global atmospheric model designed as a research playground to develop prototype ideas on machine-learned representations of climate processes and computationally efficient climate models. He also works on low precision computing, data compression and information theory, predictability of weather and climate, and software engineering.
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.
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.