Hoe DeepMind Cyclonen Voorspelt met AI en JAX

Het voorspellen van het weer voor een week vooruit heeft traditioneel gesteund op krachtige supercomputers. Ze besteden uren aan het oplossen van hydrodynamische en thermodynamische vergelijkingen op enorme roosters. Google DeepMind koos een andere aanpak: engineers trainden neurale netwerken om atmosferische processen in enkele seconden te voorspellen. De repository google-deepmind/weathernext bevat broncode en gewichten voor de architecturen van WeatherNext 2, GraphCast en GenCast, inclusief gespecialiseerde modellen voor het volgen van tropische cyclonen.

Laten we verkennen hoe het project is gestructureerd, welke hardwarebronnen je nodig hebt om het uit te voeren, en wat ontwikkelaars of onderzoekers eruit kunnen halen.

Wat er in de Repository Zit

Het project combineert verschillende generaties van DeepMinds meteorologische neurale netwerken. De hoofdfocus ligt op WeatherNext 2 (WN2) — een globaal middellange-termijn weersvoorspellingsmodel. Het werkt met een ruimtelijke resolutie van 0,25 graden (een rooster met een stap van ongeveer 30 km op het aardoppervlak).

Het model neemt de beginsituatie van de atmosfeer en voert autoregressieve runs uit, waarbij het stap-voor-stap voorspellingen genereert van temperatuur, windsnelheid op verschillende hoogtes, geopotential en cycloonbeweging.

De repository bevat drie categorieën modellen:

• WeatherNext 2 (<2025): Het werkpaard. Het is fijngeregeld op operationele ECMWF HRES (High Resolution Forecasts) data, dus het kan direct worden geïnitialiseerd vanuit echte synoptische analyses. Naast standaard metrieken voorspelt het windsnelheid op 100 meter hoogte, wat cruciaal is voor windenergie.
• WeatherNext Cyclones: Een serie modellen (2023, 2024 en 2025 versies) geoptimaliseerd voor het volgen van orkanen en tropische stormen. De 2025 versie genaamd FNV3 is al getest tijdens een echt Atlantisch orkaanseizoen.
• WeatherNext Cyclones Mini: Lichtgewicht modellen met 1-graden resolutie. Ze zijn veel kleiner en geschikt voor lokale tests of draaien op één GPU.

Naast WN2 bevat map docs/ links en materialen over eerdere modellen: het grafgebaseerde GraphCast en het diffusiegebaseerde GenCast.

Hardware- en Systeemvereisten

Het uitvoeren van volledige modellen op 0,25° is resource-intensief. Inferentie van zware checkpoints is geoptimaliseerd voor TPU (specifiek TPU v5p). Als je probeert de volledige versie op een GPU uit te voeren, heb je een kaart nodig op het niveau van NVIDIA H100 vanwege de geheugenvolumes die nodig zijn voor runs met hoge resolutie tensors.

Voor verkenning en experimenten hebben de makers een Mini-versie vrijgegeven. Deze is ontwikkelaarsvriendelijk:

• Draait soepel gratis in Google Colab op TPU v5e-1.
• Werkt op GPU, zelfs met oudere accelerators zoals NVIDIA P100 of T4.

Hoe de Code is Gestructureerd en Waar te Beginnen

De code is geschreven in Python met JAX en library xarray voor het werken met multidimensionale meteorologische data-arrays.

Om de library te installeren, voer je simpelweg uit:

pip install git+https://github.com/google-deepmind/[email protected]

De auteurs raden sterk aan om een specifieke release te pinnen, aangezien dit onderzoekscode is en API-stabiliteit niet gegarandeerd is.

Snelle Start via Colab

Map docs/weathernext2/wn2_demo.ipynb bevat een kant-en-klaar notebook. Het demonstreert de complete workflow met het model:

1. Laad gewichten uit een openbare Google Cloud bucket (dm_graphcast).
2. Prepareer initiële atmosferische condities (ERA5 of HRES data via Zarr-formaat en WeatherBench2 platform).
3. Initialiseer FGN/WN2 architectuur in JAX.
4. Voer autoregressieve rollout uit om een voorspelling van enkele dagen te genereren.
5. Bereken cycloontracks met de ingebouwde tracker.
6. Bereken loss en voer één gradientestap uit (ja, je kunt het model lokaal fijntunen).

Hier is conceptueel hoe de voorspellingsloop er in code uitziet:

# Загрузка весов и конфигурации
model_weights = load_weights("WeatherNextCyclones_Mini_<2024.npz")
initial_state = load_hres_analysis(timestamp="2025-01-01T00:00")

# Авторегрессионный прогон на несколько шагов вперед
forecast_steps = 14  # например, прогноз на 3.5 дня с шагом 6 часов
predictions = run_autoregressive_rollout(
    model=wn2_architecture,
    params=model_weights,
    inputs=initial_state,
    steps=forecast_steps
)

Praktische Toepassingen

Als je niet van plan bent om het model op je eigen synoptische data te fijntunen of wetenschappelijke pipelines te bouwen, is het opzetten van JAX-infrastructuur niet nodig. Google publiceert kant-en-klare WN2-resultaten naar open bronnen:

• OpenMeteo API: Je kunt WN2-data verkrijgen via reguliere REST-verzoeken zonder verbinding te maken met een GPU.
• Google Cloud / BigQuery / Vertex AI: Kant-en-klare datasets voor analyses en grootschalige queries.
• WeatherLab: Een interactieve service voor het visualiseren van orkaanstrajecten.

Voor degenen die werken aan onderzoek naar hernieuwbare energie, klimaatrisicoanalyse of het bouwen van hun eigen ML-diensten, dient de code uit de repository als een kant-en-klare basis. Sectie utils/ bevat handige utilities voor datanormalisatie, het werken met grafische blokken en het berekenen van specifieke meteorologische verliezen.

Conclusie

WeatherNext is een uitstekende demonstratie van hoe deep learning traditionele numerieke weersvoorspelling (NWP) geleidelijk vervangt of aanvult.

Voordelen:

• Open gewichten en demo-notebooks die draaien in Colab.
• Lichtgewicht Mini-modellen voor experimenten op reguliere hardware.
• Kant-en-klare integraties met echte meteorologische data (WeatherBench2, ERA5, HRES).

Nadelen:

• Volledige-resolutie inferentie vereist top-tier H100 of TPU accelerators.
• API garandeert geen backward compatibility tussen versies.

Als je met JAX werkt of wilt begrijpen hoe je autoregressieve passes op enorme tensorroosters uitvoert, probeer dan zeker de WN2 demo-notebook.