Cómo DeepMind Predice Ciclones Usando IA y JAX

Predecir el clima con una semana de anticipación tradicionalmente ha dependido de potentes supercomputadoras. Estas pasan horas resolviendo ecuaciones hidrodinámicas y termodinámicas en enormes mallas. Google DeepMind adoptó un enfoque diferente: los ingenieros entrenaron redes neuronales para predecir procesos atmosféricos en cuestión de segundos. El repositorio google-deepmind/weathernext contiene código fuente y pesos para las arquitecturas WeatherNext 2, GraphCast y GenCast, incluyendo modelos especializados para rastrear ciclones tropicales.

Exploremos cómo está estructurado el proyecto, qué recursos de hardware necesitarás para ejecutarlo y qué pueden extraer los desarrolladores o investigadores de él.

Qué hay dentro del repositorio

El proyecto combina varias generaciones de redes neuronales meteorológicas de DeepMind. El enfoque principal está en WeatherNext 2 (WN2) — un modelo global de pronóstico de rango medio. Opera a una resolución espacial de 0.25 grados (una malla con un paso de aproximadamente 30 km en la superficie terrestre).

El modelo toma el estado inicial de la atmósfera y realiza ejecuciones autorregresivas, generando pronósticos paso a paso de temperatura, velocidad del viento a varias altitudes, geopotencial y movimiento de ciclones.

El repositorio contiene tres categorías de modelos:

• WeatherNext 2 (<2025): El caballo de batalla. Se ajustó con datos operativos del ECMWF HRES (High Resolution Forecasts), por lo que puede inicializarse directamente desde análisis sinópticos reales. Más allá de las métricas estándar, pronostica la velocidad del viento a 100 metros de altitud, lo cual es crítico para la energía eólica.
• WeatherNext Cyclones: Una serie de modelos (versiones 2023, 2024 y 2025) optimizados para rastrear huracanes y tormentas tropicales. La versión 2025 llamada FNV3 ya fue probada en combate durante una temporada real de huracanes del Atlántico.
• WeatherNext Cyclones Mini: Modelos ligeros con resolución de 1 grado. Son mucho más pequeños y adecuados para pruebas locales o para ejecutar en una sola GPU.

Más allá de WN2, la carpeta docs/ contiene enlaces y materiales sobre modelos anteriores: el basado en grafos GraphCast y el basado en difusión GenCast.

Hardware y requisitos del sistema

Ejecutar modelos de tamaño completo a 0.25° requiere muchos recursos. La inferencia de puntos de control pesados está optimizada para TPU (específicamente TPU v5p). Si intentas ejecutar la versión completa en una GPU, necesitarás una tarjeta del nivel de NVIDIA H100 debido a los volúmenes de memoria requeridos para ejecuciones de tensores de alta resolución.

Para exploración y experimentos, los creadores publicaron una versión Mini. Esta es amigable para desarrolladores:

• Se ejecuta sin problemas en Google Colab gratuito en TPU v5e-1.
• Funciona en GPU incluso con aceleradores más antiguos como NVIDIA P100 o T4.

Cómo está estructurado el código y por dónde empezar

El código está escrito en Python usando JAX y la librería xarray para trabajar con arrays de datos meteorológicos multidimensionales.

Para instalar la librería, simplemente ejecuta:

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

Los autores recomiendan encarecidamente fijar una versión específica, ya que este es código de investigación y no se garantiza la estabilidad de la API.

Inicio rápido a través de Colab

La carpeta docs/weathernext2/wn2_demo.ipynb contiene un notebook listo para usar. Demuestra el flujo de trabajo completo con el modelo:

1. Cargar pesos desde un bucket público de Google Cloud (dm_graphcast).
2. Preparar las condiciones atmosféricas iniciales (datos ERA5 o HRES a través del formato Zarr y la plataforma WeatherBench2).
3. Inicializar la arquitectura FGN/WN2 en JAX.
4. Ejecutar despliegue autorregresivo para generar un pronóstico de varios días.
5. Calcular las trayectorias de ciclones usando el tracker integrado.
6. Calcular la pérdida y realizar un paso de gradiente (sí, puedes ajustar finamente el modelo localmente).

Aquí está conceptualmente cómo se ve el bucle de pronóstico en código:

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

Aplicaciones prácticas

Si no planeas ajustar finamente el modelo con tus propios datos sinópticos o construir pipelines científicos, no es necesario configurar la infraestructura de JAX. Google publica resultados de WN2 listos para usar en fuentes abiertas:

• API de OpenMeteo: Puedes obtener datos de WN2 a través de solicitudes REST regulares sin conectar a una GPU.
• Google Cloud / BigQuery / Vertex AI: Datasets listos para usar para análisis y consultas a gran escala.
• WeatherLab: Un servicio interactivo para visualizar trayectorias de huracanes.

Para quienes trabajan en investigación de energía renovable, análisis de riesgo climático o construcción de sus propios servicios de ML, el código del repositorio sirve como base lista. La sección utils/ contiene utilidades convenientes para normalización de datos, trabajo con bloques de grafos y cálculo de pérdidas meteorológicas específicas.

En resumen

WeatherNext es una excelente demostración de cómo el aprendizaje profundo está reemplazando gradualmente o complementando la predicción meteorológica numérica tradicional (NWP).

Ventajas:

• Pesos abiertos y notebooks de demostración que se ejecutan en Colab.
• Modelos Mini ligeros para experimentos en hardware común.
• Integraciones listas para usar con datos meteorológicos reales (WeatherBench2, ERA5, HRES).

Desventajas:

• La inferencia a resolución completa requiere aceleradores de primera línea H100 o TPU.
• La API no garantiza compatibilidad hacia atrás entre versiones.

Si trabajas con JAX o quieres entender cómo ejecutar pasadas autorregresivas en enormes mallas de tensores, definitivamente prueba el notebook de demostración de WN2.