Jak DeepMind przewiduje cyklony za pomocą AI i JAX

Prognozowanie pogody na tydzień do przodu tradycyjnie opierało się na potężnych superkomputerach. Spędzają one godziny na rozwiązywaniu równań hydrodynamicznych i termodynamicznych na ogromnych siatkach. Google DeepMind poszło inną drogą: inżynierowie wytrenowali sieci neuronowe do przewidywania procesów atmosferycznych w ciągu sekund. Repozytorium google-deepmind/weathernext zawiera kod źródłowy i wagi dla architektur WeatherNext 2, GraphCast i GenCast, w tym wyspecjalizowane modele do śledzenia cyklonów tropikalnych.

Poznajmy strukturę projektu, zasoby sprzętowe potrzebne do jego uruchomienia oraz co programiści i badacze mogą z niego wyciągnąć.

Co zawiera repozytorium

Projekt łączy kilka generacji meteorologicznych sieci neuronowych DeepMind. Główny nacisk kładziony jest na WeatherNext 2 (WN2) — globalny model prognozowania średniego zasięgu. Działa on w rozdzielczości przestrzennej 0,25 stopnia (siatka o kroku około 30 km na powierzchni Ziemi).

Model przyjmuje początkowy stan atmosfery i wykonuje przebiegi autoregresyjne, generując prognozy krok po kroku dotyczące temperatury, prędkości wiatru na różnych wysokościach, geopotencjału oraz ruchu cyklonów.

Repozytorium zawiera trzy kategorie modeli:

• WeatherNext 2 (<2025): Główny model. Został dostrojony na podstawie operacyjnych danych ECMWF HRES (High Resolution Forecasts), więc może być inicjowany bezpośrednio z rzeczywistych analiz synoptycznych. Poza standardowymi metrykami prognozuje prędkość wiatru na wysokości 100 metrów, co jest kluczowe dla energetyki wiatrowej.
• WeatherNext Cyclones: Seria modeli (wersje z 2023, 2024 i 2025) zoptymalizowanych do śledzenia huraganów i burz tropikalnych. Wersja z 2025 roku o nazwie FNV3 została już przetestowana w warunkach rzeczywistych podczas sezonu huraganów na Atlantyku.
• WeatherNext Cyclones Mini: Lekkie modele o rozdzielczości 1 stopnia. Są znacznie mniejsze i nadają się do testów lokalnych lub uruchamiania na pojedynczym GPU.

Oprócz WN2, folder docs/ zawiera linki i materiały dotyczące wcześniejszych modeli: opartego na grafach GraphCast oraz opartego na dyfuzji GenCast.

Wymagania sprzętowe i systemowe

Uruchamianie modeli w pełnym rozmiarze przy 0,25° wymaga sporych zasobów. Wnioskowanie z ciężkich checkpointów jest zoptymalizowane pod kątem TPU (konkretnie TPU v5p). Jeśli spróbujesz uruchomić pełną wersję na GPU, będziesz potrzebować karty na poziomie NVIDIA H100 ze względu na ilość pamięci wymaganą do przebiegów tensorów o wysokiej rozdzielczości.

Do eksploracji i eksperymentów twórcy udostępnili wersję Mini. Ta jest przyjazna dla developerów:

• Działa bezproblemowo w darmowym Google Colab na TPU v5e-1.
• Działa na GPU nawet ze starszymi akceleratorami jak NVIDIA P100 czy T4.

Struktura kodu i od czego zacząć

Kod jest napisany w Pythonie z użyciem JAX i biblioteki xarray do pracy z wielowymiarowymi tablicami danych meteorologicznych.

Aby zainstalować bibliotekę, wystarczy uruchomić:

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

Autorzy zdecydowanie zalecają przypięcie konkretnej wersji, ponieważ jest to kod badawczy i stabilność API nie jest gwarantowana.

Szybki start przez Colab

Folder docs/weathernext2/wn2_demo.ipynb zawiera gotowy do użycia notebook. Demonstruje on kompletny workflow z modelem:

1. Załaduj wagi z publicznego bucketa Google Cloud (dm_graphcast).
2. Przygotuj początkowe warunki atmosferyczne (dane ERA5 lub HRES przez format Zarr i platformę WeatherBench2).
4. Uruchom przebieg autoregresyjny, aby wygenerować prognozę na kilka dni do przodu.
5. Oblicz trasy cyklonów za pomocą wbudowanego trackera.
6. Oblicz stratę i wykonaj jeden krok gradientowy (tak, możesz lokalnie dostroić model).

Oto jak koncepcyjnie wygląda pętla prognozowania w kodzie:

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

Praktyczne zastosowania

Jeśli nie planujesz dostrajania modelu na własnych danych synoptycznych ani budowania naukowych pipeline'ów, konfiguracja infrastruktury JAX nie jest konieczna. Google publikuje gotowe wyniki WN2 w źródłach otwartych:

• OpenMeteo API: Możesz uzyskać dane WN2 przez regularne żądania REST bez łączenia się z GPU.
• Google Cloud / BigQuery / Vertex AI: Gotowe zestawy danych do analizy i zapytań na dużą skalę.
• WeatherLab: Interaktywna usługa do wizualizacji trajektorii huraganów.

Dla osób pracujących nad badaniami nad energią odnawialną, analizą ryzyka klimatycznego lub budowaniem własnych usług ML, kod z repozytorium służy jako gotowa podstawa. Sekcja utils/ zawiera wygodne narzędzia do normalizacji danych, pracy z blokami grafowymi i obliczania specyficznych strat meteorologicznych.

Podsumowanie

WeatherNext to doskonała demonstracja tego, jak głębokie uczenie stopniowo zastępuje lub uzupełnia tradycyjne numeryczne prognozowanie pogody (NWP).

Zalety:

• Otwarta wagi i demo notebooki działające w Colab.
• Lekkie modele Mini do eksperymentów na zwykłym sprzęcie.
• Gotowe integracje z rzeczywistymi danymi meteorologicznymi (WeatherBench2, ERA5, HRES).

Wady:

• Wnioskowanie w pełnej rozdzielczości wymaga akceleratorów najwyższej klasy H100 lub TPU.
• API nie gwarantuje kompatybilności wstecznej między wersjami.

Jeśli pracujesz z JAX lub chcesz zrozumieć, jak uruchamiać autoregresyjne przebiegi na ogromnych siatkach tensorów, koniecznie wypróbuj demo notebook WN2.