Comment DeepMind prédit les cyclones grâce à l'IA et JAX

La prévision météorologique à une semaine s'est traditionnellement appuyée sur des supercalculateurs puissants. Ces derniers passent des heures à résoudre des équations hydrodynamiques et thermodynamiques sur d'immenses grilles. Google DeepMind a adopté une approche différente : des ingénieurs ont entraîné des réseaux de neurones à prédire les processus atmosphériques en quelques secondes. Le dépôt google-deepmind/weathernext contient le code source et les poids pour les architectures WeatherNext 2, GraphCast et GenCast, ainsi que des modèles spécialisés pour le suivi des cyclones tropicaux.

Explorons la structure du projet, les ressources matérielles nécessaires pour l'exécuter, et ce que les développeurs ou chercheurs peuvent en extraire.

Ce que contient le dépôt

Le projet combine plusieurs générations de réseaux de neurones météorologiques de DeepMind. L'accent principal est mis sur WeatherNext 2 (WN2) — un modèle de prévision mondiale à moyen terme. Il fonctionne à une résolution spatiale de 0,25 degré (une grille avec un pas d'environ 30 km à la surface de la Terre).

Le modèle prend l'état initial de l'atmosphère et effectue des exécutions autorégressives, générant des prévisions étape par étape de la température, de la vitesse du vent à différentes altitudes, du géopotentiel et du mouvement des cyclones.

Le dépôt contient trois catégories de modèles :

• WeatherNext 2 (<2025) : Le modèle principal. Il a été affiné sur des données opérationnelles ECMWF HRES (High Resolution Forecasts), donc il peut être initialisé directement à partir d'analyses synoptiques réelles. Au-delà des métriques standard, il prévoit la vitesse du vent à 100 mètres d'altitude, ce qui est essentiel pour l'énergie éolienne.
• WeatherNext Cyclones : Une série de modèles (versions 2023, 2024 et 2025) optimisés pour le suivi des ouragans et des tempêtes tropicales. La version 2025 appelée FNV3 a déjà été testée en conditions réelles lors d'une saison des ouragans atlantiques.
• WeatherNext Cyclones Mini : Des modèles légers avec une résolution de 1 degré. Ils sont beaucoup plus petits et adaptés aux tests locaux ou à l'exécution sur un seul GPU.

Au-delà de WN2, le dossier docs/ contient des liens et des matériaux sur les modèles précédents : le GraphCast basé sur les graphes et le GenCast basé sur la diffusion.

Matériel et exigences système

L'exécution de modèles complets à 0,25° est exigeante en ressources. L'inférence des points de contrôle lourds est optimisée pour les TPU (spécifiquement TPU v5p). Si vous essayez d'exécuter la version complète sur un GPU, vous aurez besoin d'une carte de niveau NVIDIA H100 en raison des volumes de mémoire requis pour les exécutions de tenseurs haute résolution.

Pour l'exploration et les expériences, les créateurs ont publié une version Mini. Celle-ci est adaptée aux développeurs :

• S'exécute sans problème dans Google Colab gratuit sur TPU v5e-1.
• Fonctionne sur GPU même avec des accélérateurs plus anciens comme NVIDIA P100 ou T4.

Structure du code et par où commencer

Le code est écrit en Python en utilisant JAX et la bibliothèque xarray pour travailler avec des tableaux de données météorologiques multidimensionnels.

Pour installer la bibliothèque, exécutez simplement :

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

Les auteurs recommandent fortement de figer une version spécifique, car il s'agit de code de recherche et la stabilité de l'API n'est pas garantie.

Démarrage rapide via Colab

Le dossier docs/weathernext2/wn2_demo.ipynb contient un notebook prêt à l'emploi. Il démontre le flux de travail complet avec le modèle :

1. Charger les poids depuis un bucket Google Cloud public (dm_graphcast).
2. Préparer les conditions atmosphériques initiales (données ERA5 ou HRES via le format Zarr et la plateforme WeatherBench2).
3. Initialiser l'architecture FGN / WN2 dans JAX.
4. Exécuter le déroulé autorégressif pour générer une prévision à plusieurs jours.
5. Calculer les trajectoires des cyclones à l'aide du tracker intégré.
6. Calculer la perte et effectuer une étape de gradient (oui, vous pouvez affiner localement le modèle).

Voici conceptuellement à quoi ressemble la boucle de prévision en code :

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

Applications pratiques

Si vous ne prévoyez pas d'affiner le modèle sur vos propres données synoptiques ou de construire des pipelines scientifiques, la mise en place de l'infrastructure JAX n'est pas nécessaire. Google publie des résultats WN2 prêts à l'emploi vers des sources ouvertes :

• API OpenMeteo : Vous pouvez obtenir des données WN2 via des requêtes REST ordinaires sans vous connecter à un GPU.
• Google Cloud / BigQuery / Vertex AI : Des ensembles de données prêts à l'emploi pour l'analyse et les requêtes à grande échelle.
• WeatherLab : Un service interactif pour visualiser les trajectoires des ouragans.

Pour ceux qui travaillent sur la recherche en énergie renouvelable, l'analyse des risques climatiques ou la construction de leurs propres services ML, le code du dépôt sert de base prête à l'emploi. La section utils/ contient des utilitaires pratiques pour la normalisation des données, le travail avec les blocs de graphes et le calcul de pertes météorologiques spécifiques.

En conclusion

WeatherNext est une excellente démonstration de la façon dont l'apprentissage profond remplace progressivement ou complète la prévision numérique du temps (PNT) traditionnelle.

Avantages :

• Poids ouverts et notebooks de démonstration qui s'exécutent dans Colab.
• Modèles Mini légers pour des expériences sur du matériel ordinaire.
• Intégrations prêtes à l'emploi avec des données météorologiques réelles (WeatherBench2, ERA5, HRES).

Inconvénients :

• L'inférence en pleine résolution nécessite des accélérateurs haut de gamme H100 ou TPU.
• L'API ne garantit pas la compatibilité descendante entre les versions.

Si vous travaillez avec JAX ou souhaitez comprendre comment exécuter des passes autorégressives sur d'immenses grilles de tenseurs, essayez absolument le notebook de démonstration WN2.