Como a DeepMind Prevê Ciclones Usando IA e JAX

Prever o tempo com uma semana de antecedência sempre dependeu de supercomputadores poderosos. Eles passam horas resolvendo equações hidrodinâmicas e termodinâmicas em grades massivas. O Google DeepMind adotou uma abordagem diferente: engenheiros treinaram redes neurais para prever processos atmosféricos em questão de segundos. O repositório google-deepmind/weathernext contém código-fonte e pesos para as arquiteturas WeatherNext 2, GraphCast e GenCast, incluindo modelos especializados para rastrear ciclones tropicais.

Vamos explorar como o projeto é estruturado, quais recursos de hardware você precisará para executá-lo e o que desenvolvedores ou pesquisadores podem extrair dele.

O Que Está Dentro do Repositório

O projeto combina várias gerações de redes neurais meteorológicas da DeepMind. O foco principal é o WeatherNext 2 (WN2) — um modelo global de previsão de médio prazo. Ele opera em uma resolução espacial de 0,25 graus (uma grade com um passo de aproximadamente 30 km na superfície terrestre).

O modelo recebe o estado inicial da atmosfera e executa passos autorregressivos, gerando previsões passo a passo de temperatura, velocidade do vento em várias altitudes, geopotencial e movimento de ciclones.

O repositório contém três categorias de modelos:

• WeatherNext 2 (<2025): O modelo principal. Ele foi ajustado com dados operacionais do ECMWF HRES (Previsões de Alta Resolução), então pode ser inicializado diretamente a partir de análises sinóticas reais. Além das métricas padrão, ele prevê a velocidade do vento a 100 metros de altitude, o que é crítico para a energia eólica.
• WeatherNext Cyclones: Uma série de modelos (versões de 2023, 2024 e 2025) otimizados para rastrear furacões e tempestades tropicais. A versão de 2025 chamada FNV3 já foi testada em combate durante uma temporada real de furacões no Atlântico.
• WeatherNext Cyclones Mini: Modelos leves com resolução de 1 grau. Eles são muito menores e adequados para testes locais ou execução em uma única GPU.

Além do WN2, a pasta docs/ contém links e materiais sobre modelos anteriores: o GraphCast baseado em grafos e o GenCast baseado em difusão.

Requisitos de Hardware e Sistema

A execução de modelos de tamanho completo a 0,25° é intensiva em recursos. A inferência de checkpoints pesados é otimizada para TPU (especificamente TPU v5p). Se você tentar executar a versão completa em uma GPU, precisará de uma placa no nível de NVIDIA H100 devido aos volumes de memória necessários para execuções de tensores de alta resolução.

Para exploração e experimentos, os criadores lançaram uma versão Mini. Esta é amigável para desenvolvedores:

• Roda suavemente no Google Colab gratuito em TPU v5e-1.
• Funciona em GPU mesmo com aceleradores mais antigos como NVIDIA P100 ou T4.

Como o Código é Estruturado e Por Onde Começar

O código é escrito em Python usando JAX e a biblioteca xarray para trabalhar com arrays multidimensionais de dados meteorológicos.

Para instalar a biblioteca, basta executar:

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

Os autores recomendam fortemente fixar uma versão específica, pois este é um código de pesquisa e a estabilidade da API não é garantida.

Início Rápido via Colab

A pasta docs/weathernext2/wn2_demo.ipynb contém um notebook pronto para uso. Ele demonstra o fluxo de trabalho completo com o modelo:

1. Carregar pesos de um bucket público do Google Cloud (dm_graphcast).
2. Preparar condições atmosféricas iniciais (dados ERA5 ou HRES via formato Zarr e plataforma WeatherBench2).
3. Inicializar arquitetura FGN/WN2 em JAX.
4. Executar rollout autorregressivo para gerar uma previsão de vários dias.
5. Calcular trajetórias de ciclones usando o rastreador integrado.
6. Calcular loss e realizar um passo de gradiente (sim, você pode fazer fine-tuning do modelo localmente).

Aqui está conceitualmente como o loop de previsão fica no 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
)

Aplicações Práticas

Se você não planeja fazer fine-tuning do modelo com seus próprios dados sinóticos ou construir pipelines científicos, configurar a infraestrutura JAX não é necessário. O Google publica resultados WN2 prontos em fontes abertas:

• OpenMeteo API: Você pode obter dados WN2 via requisições REST regulares sem conectar a uma GPU.
• Google Cloud / BigQuery / Vertex AI: Datasets prontos para análise e consultas em grande escala.
• WeatherLab: Um serviço interativo para visualizar trajetórias de furacões.

Para quem trabalha com pesquisa de energia renovável, análise de risco climático ou construção de seus próprios serviços de ML, o código do repositório serve como uma base pronta. A seção utils/ contém utilitários convenientes para normalização de dados, trabalho com blocos de grafos e cálculo de losses meteorológicos específicos.

Resumo Final

O WeatherNext é uma excelente demonstração de como o deep learning está gradualmente substituindo ou complementando a previsão numérica do tempo tradicional (NWP).

Prós:

• Pesos abertos e notebooks de demonstração que rodam no Colab.
• Modelos Mini leves para experimentos em hardware comum.
• Integrações prontas com dados meteorológicos reais (WeatherBench2, ERA5, HRES).

Contras:

• A inferência em resolução completa requer aceleradores de última geração como H100 ou TPU.
• A API não garante compatibilidade retroativa entre versões.

Se você trabalha com JAX ou quer entender como executar passos autorregressivos em grades massivas de tensores, definitivamente experimente o notebook de demonstração do WN2.