How DeepMind Predicts Cyclones Using AI and JAX

Predicting weather a week ahead has traditionally relied on powerful supercomputers. They spend hours solving hydrodynamic and thermodynamic equations on massive grids. Google DeepMind took a different approach: engineers trained neural networks to predict atmospheric processes in a matter of seconds. The repository google-deepmind/weathernext contains source code and weights for WeatherNext 2, GraphCast, and GenCast architectures, including specialized models for tracking tropical cyclones.

Let's explore how the project is structured, what hardware resources you'll need to run it, and what developers or researchers can extract from it.

What's Inside the Repository

The project combines several generations of DeepMind's meteorological neural networks. The main focus is on WeatherNext 2 (WN2) — a global medium-range forecasting model. It operates at a spatial resolution of 0.25 degrees (a grid with a step of approximately 30 km on Earth's surface).

The model takes the initial state of the atmosphere and performs autoregressive runs, generating step-by-step forecasts of temperature, wind speed at various altitudes, geopotential, and cyclone movement.

The repository contains three categories of models:

• WeatherNext 2 (<2025): The workhorse. It was fine-tuned on operational ECMWF HRES (High Resolution Forecasts) data, so it can be initialized directly from real synoptic analyses. Beyond standard metrics, it forecasts wind speed at 100 meters altitude, which is critical for wind energy.
• WeatherNext Cyclones: A series of models (2023, 2024, and 2025 versions) optimized for tracking hurricanes and tropical storms. The 2025 version called FNV3 has already been battle-tested during a real Atlantic hurricane season.
• WeatherNext Cyclones Mini: Lightweight models with 1-degree resolution. They are much smaller and suitable for local tests or running on a single GPU.

Beyond WN2, folder docs/ contains links and materials on previous models: the graph-based GraphCast and the diffusion-based GenCast.

Hardware and System Requirements

Running full-sized models at 0.25° is resource-intensive. Inference of heavy checkpoints is optimized for TPU (specifically TPU v5p). If you try to run the full version on a GPU, you'll need a card at the level of NVIDIA H100 due to the memory volumes required for high-resolution tensor runs.

For exploration and experiments, the creators released a Mini version. This one is developer-friendly:

• Runs smoothly in free Google Colab on TPU v5e-1.
• Works on GPU even with older accelerators like NVIDIA P100 or T4.

How the Code Is Structured and Where to Start

The code is written in Python using JAX and library xarray for working with multidimensional meteorological data arrays.

To install the library, simply run:

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

The authors strongly recommend pinning a specific release, as this is research code and API stability is not guaranteed.

Quick Start via Colab

Folder docs/weathernext2/wn2_demo.ipynb contains a ready-to-use notebook. It demonstrates the complete workflow with the model:

1. Load weights from a public Google Cloud bucket (dm_graphcast).
2. Prepare initial atmospheric conditions (ERA5 or HRES data via Zarr format and WeatherBench2 platform).
3. Initialize FGN / WN2 architecture in JAX.
4. Run autoregressive rollout to generate a forecast several days ahead.
5. Calculate cyclone tracks using the built-in tracker.
6. Compute loss and perform one gradient step (yes, you can locally fine-tune the model).

Here's conceptually how the forecasting loop looks in 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
)

Practical Applications

If you're not planning to fine-tune the model on your own synoptic data or build scientific pipelines, setting up JAX infrastructure isn't necessary. Google publishes ready-made WN2 results to open sources:

• OpenMeteo API: You can get WN2 data via regular REST requests without connecting to a GPU.
• Google Cloud / BigQuery / Vertex AI: Ready-made datasets for analytics and large-scale queries.
• WeatherLab: An interactive service for visualizing hurricane trajectories.

For those working on renewable energy research, climate risk analysis, or building their own ML services, the code from the repository serves as a ready foundation. Section utils/ contains convenient utilities for data normalization, working with graph blocks, and calculating specific meteorological losses.

Bottom Line

WeatherNext is an excellent demonstration of how deep learning is gradually replacing or complementing traditional numerical weather prediction (NWP).

Pros:

• Open weights and demo notebooks that run in Colab.
• Lightweight Mini-models for experiments on regular hardware.
• Ready-made integrations with real meteorological data (WeatherBench2, ERA5, HRES).

Cons:

• Full-resolution inference requires top-tier H100 or TPU accelerators.
• API does not guarantee backward compatibility between versions.

If you work with JAX or want to understand how to run autoregressive passes on massive tensor grids, definitely try out the WN2 demo notebook.