Wie DeepMind Zyklone mit KI und JAX vorhersagt

Wettervorhersagen für eine Woche im Voraus basierten traditionell auf leistungsstarken Supercomputern. Diese verbringen Stunden damit, hydrodynamische und thermodynamische Gleichungen auf massiven Gittern zu lösen. Google DeepMind verfolgte einen anderen Ansatz: Ingenieure trainierten neuronale Netze, um atmosphärische Prozesse in Sekundenschnelle vorherzusagen. Das Repository google-deepmind/weathernext enthält Quellcode und Gewichte für die Architekturen WeatherNext 2, GraphCast und GenCast, einschließlich spezialisierter Modelle zur Verfolgung tropischer Zyklone.

Lass uns erkunden, wie das Projekt strukturiert ist, welche Hardware-Ressourcen du zum Ausführen benötigst und was Entwickler oder Forscher daraus ableiten können.

Was sich im Repository befindet

Das Projekt kombiniert mehrere Generationen meteorologischer neuronaler Netze von DeepMind. Der Schwerpunkt liegt auf WeatherNext 2 (WN2) – einem globalen Vorhersagemodell für den mittelfristigen Zeitraum. Es arbeitet mit einer räumlichen Auflösung von 0,25 Grad (ein Gitter mit einer Schrittweite von ca. 30 km auf der Erdoberfläche).

Das Modell nimmt den Ausgangszustand der Atmosphäre und führt autoregressive Läufe durch, um schrittweise Prognosen für Temperatur, Windgeschwindigkeit in verschiedenen Höhen, Geopotenzial und Zyklonbewegung zu generieren.

Das Repository enthält drei Modellkategorien:

• WeatherNext 2 (<2025): Das Arbeitspferd. Es wurde mit operativen ECMWF HRES-Daten (High Resolution Forecasts) feinabgestimmt, sodass es direkt aus realen synoptischen Analysen initialisiert werden kann. Über Standardmetriken hinaus prognostiziert es die Windgeschwindigkeit in 100 Metern Höhe, was für die Windenergie entscheidend ist.
• WeatherNext Cyclones: Eine Reihe von Modellen (Versionen 2023, 2024 und 2025), die für die Verfolgung von Hurrikanen und tropischen Stürmen optimiert sind. Die 2025er Version namens FNV3 wurde bereits während einer echten atlantischen Hurrikansaison erprobt.
• WeatherNext Cyclones Mini: Leichtgewichtige Modelle mit 1-Grad-Auflösung. Sie sind viel kleiner und eignen sich für lokale Tests oder das Ausführen auf einer einzelnen GPU.

Über WN2 hinaus enthält Ordner docs/ Links und Materialien zu früheren Modellen: das graphbasierte GraphCast und das diffusionsbasierte GenCast.

Hardware- und Systemanforderungen

Das Ausführen von Modellen in voller Größe bei 0,25° ist ressourcenintensiv. Die Inferenz schwerer Checkpoints ist für TPU (speziell TPU v5p) optimiert. Wenn du versuchst, die Vollversion auf einer GPU auszuführen, benötigst du eine Karte auf dem Niveau einer NVIDIA H100 aufgrund der für hochauflösende Tensorläufe erforderlichen Speichermengen.

Für Erkundung und Experimente haben die Entwickler eine Mini-Version veröffentlicht. Diese ist entwicklerfreundlich:

• Läuft reibungslos in der kostenlosen Google Colab-Umgebung auf TPU v5e-1.
• Funktioniert auf GPUs, sogar mit älteren Beschleunigern wie NVIDIA P100 oder T4.

Wie der Code strukturiert ist und wo du anfangen solltest

Der Code ist in Python unter Verwendung von JAX und Bibliothek xarray für die Arbeit mit mehrdimensionalen meteorologischen Datenarrays geschrieben.

Um die Bibliothek zu installieren, führe einfach aus:

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

Die Autoren empfehlen dringend, eine bestimmte Version zu pinnen, da es sich um Forschungscode handelt und API-Stabilität nicht garantiert wird.

Schnellstart über Colab

Ordner docs/weathernext2/wn2_demo.ipynb enthält ein gebrauchsfertiges Notebook. Es demonstriert den vollständigen Workflow mit dem Modell:

1. Lade Gewichte aus einem öffentlichen Google Cloud Bucket (dm_graphcast).
2. Bereite anfängliche atmosphärische Bedingungen vor (ERA5- oder HRES-Daten über Zarr-Format und die WeatherBench2-Plattform).
3. Initialisiere die FGN/WN2-Architektur in JAX.
4. Führe autoregressive Rollouts durch, um eine Prognose für mehrere Tage im Voraus zu generieren.
5. Berechne Zyklonbahnen mit dem integrierten Tracker.
6. Berechne den Loss und führe einen Gradientenschritt durch (ja, du kannst das Modell lokal feinabstimmen).

So sieht konzeptionell die Vorhersageschleife im Code aus:

# Загрузка весов и конфигурации
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 Anwendungen

Wenn du nicht planst, das Modell auf deinen eigenen synoptischen Daten feinzutunen oder wissenschaftliche Pipelines zu bauen, ist das Einrichten der JAX-Infrastruktur nicht notwendig. Google veröffentlicht vorgefertigte WN2-Ergebnisse an offenen Quellen:

• OpenMeteo API: Du kannst WN2-Daten über reguläre REST-Anfragen ohne GPU-Anbindung erhalten.
• Google Cloud / BigQuery / Vertex AI: Fertige Datensätze für Analysen und umfangreiche Abfragen.
• WeatherLab: Ein interaktiver Dienst zur Visualisierung von Hurrikanbahnen.

Für diejenigen, die an Forschung zu erneuerbaren Energien, Klimarisikoanalysen oder dem Aufbau eigener ML-Dienste arbeiten, dient der Code aus dem Repository als fertige Grundlage. Abschnitt utils/ enthält praktische Hilfsfunktionen für die Datennormalisierung, die Arbeit mit Graph-Blöcken und die Berechnung spezifischer meteorologischer Losses.

Fazit

WeatherNext ist eine ausgezeichnete Demonstration, wie Deep Learning traditionelle numerische Wettervorhersage (NWP) schrittweise ersetzt oder ergänzt.

Vorteile:

• Offene Gewichte und Demo-Notebooks, die in Colab laufen.
• Leichtgewichtige Mini-Modelle für Experimente auf normaler Hardware.
• Fertige Integrationen mit realen meteorologischen Daten (WeatherBench2, ERA5, HRES).

Nachteile:

• Die Inferenz in voller Auflösung erfordert High-End-H100- oder TPU-Beschleuniger.
• Die API garantiert keine Abwärtskompatibilität zwischen Versionen.

Wenn du mit JAX arbeitest oder verstehen möchtest, wie man autoregressive Durchläufe auf massiven Tensorgittern ausführt, solltest du unbedingt das WN2-Demo-Notebook ausprobieren.