Thread by @RyanKeisler

📢 Time to share a project I’ve been working on:

Forecasting Global Weather with Graph Neural Networks

🧵 (1/N)

Context: numerical weather prediction (NWP) has had a huge, positive impact on society. Decades of improvements (more data, better models, more compute) have resulted in increasingly accurate weather forecasts and growing adoption of NWP in real-world applications. (2/N)

And while statistical techniques have been used within NWP for decades, the core dynamical engines of these models are based on the physical principles that govern the atmosphere and ocean. (3/N)

More recently, spurred on by progress in ML, there has been a surge of interest in data-driven techniques for weather forecasting. The idea is to improve upon the already extremely successful NWP program through some combination of better/faster/more forecasts. (4/N)

ML+weather is a huge research area (nowcasting, numerics/PDEs, high-impact events…) but in this thread I’ll focus on one area: using ML to emulate the big, global models like GFS or ECMWF that give us 10-day forecasts & are the foundation for many other models/applications (5/N)

The work I'm sharing today is a data-driven, ML system for forecasting global weather. It's inspired by similar work from @raspstephan, @wx_jon, and others, but (i) simulates a significantly denser physical system and (ii) uses graph neural networks (GNNs) instead of CNNs. (6/N)

The motivation to try GNNs came from seeing the success that @spectralhippo, @PeterWBattaglia, and @DeepMind collaborators had in simulating physical systems on complex mesh geometries using message-passing GNNs:

https://twitter.com/spectralhippo/status/1405935950413414405

(7/N)

The flexibility of GNNs provides benefits well suited to NWP like: (i) straightforward handling of the spherical geometry of earth, (ii) the potential for learning multi-resolution models, and (iii) the potential for adaptive mesh refinement (not explored in this work). (8/N)

In this work, the ML model steps forward the current 3D atmospheric state — stored as 6 prognostic variables (temp, wind, etc) on 13 pressure levels — by 6 hours, and steps are chained to make longer forecasts. The model is trained on ERA5 reanalysis data or GFS forecasts. (9/N)

Here is an example of a typical 6-hour step, with ERA5 on the left and the ML model’s predictions on the right (on a 1 deg lat/lon grid). The model has learned to accurately predict the 6-hour changes in these variables using only the initial state. (10/N)

Multiple 6-hour steps are chained together to produce multi-day forecasts. Here is an example 3-day rollout showing 850 hPa humidity. The ML model generally tracks the large-scale flows seen in ERA5, although the ML predictions do become smoother over time. (11/N)

Looks pretty good, but how does it compare to previous approaches? For starters, this model improves upon the forecast performance of previous data-driven models, as shown below. (12/N)

And surprisingly, it seems the forecast performance is comparable to operational, high-res, physical models from GFS & ECMWF, at least when evaluated on 1-deg scales & using reanalysis init. Please see Sec 4.5 of paper for more info on this nuanced & qualified comparison. (13/N)

Finally, I tried to bring this thing to life a bit by connecting it to live GFS forecasts and producing a hybrid physics+ML system. The punchline is that this system can “look ahead” and see where the GFS forecast is going before it gets there. (14/N)

I spent some time in the paper trying to provide a balanced discussion of why this model works well. I think it’s probably a combination of many things: using GNNs, modeling a dense physical system à la traditional NWP, a specific loss normalization, & GPU memory/tricks. (15/N)

If you’d like to learn more, please check out the paper, linked below.

pdf: arxiv.org/pdf/2202.07575…
site: rkeisler.github.io/graph_weather

(16/N)

Inspired by the work of many people (& their collaborators) some of whom I will shamelessly tag below: @raspstephan @thuereyGroup @wx_jon @spectralhippo @PeterWBattaglia @BethanyLusch @shoyer @PDueben @oliwm @jeremy_mcgibbon @NoahBrenowitz @cosmo_shirley @NalKalchbrenner (17/N)

Anyway, it has been a lot of fun, and I'm excited to share it. (fin)