# How to Reproduce WeatherNext Results: A Complete Step-by-Step Guide

> Reproduce WeatherNext results with this step-by-step guide. Download weights, install the package, prepare data, and run the autoregressive rollout function for accurate weather forecasting.

- Repository: [Google DeepMind/weathernext](https://github.com/google-deepmind/weathernext)
- Tags: how-to-guide
- Published: 2026-08-16

---

**To reproduce WeatherNext results, download pretrained weights from Google Cloud, install the package with pip, prepare ERA5 or HRES input data, instantiate the ForwardPass architecture, and run the autoregressive rollout function with your desired forecast horizon.**

The [google-deepmind/weathernext](https://github.com/google-deepmind/weathernext) repository contains a family of state-of-the-art deep learning models for global weather and tropical cyclone forecasting. This guide walks through the exact pipeline used by Google DeepMind researchers to generate the results reported in the WeatherNext 2 and WeatherNext Cyclones papers, with specific references to source files and runnable code examples.

## Understanding the WeatherNext Model Family

Before diving into reproduction, it's important to identify which model variant matches your use case:

| Model | Resolution | Primary Paper | Best For |
|-------|-----------|---------------|----------|
| **WeatherNext 2 (WN2)** | 0.25° (~30 km) | *WeatherNext 2: Our most advanced weather forecasting model* | Global medium-range atmospheric forecasts, including 100 m wind |
| **WeatherNext Cyclones** | 0.25° | *Operational tropical cyclone forecasting with AI* | Cyclone tracking and intensity forecasts |
| **WeatherNext Cyclones Mini** | 1° | Same paper (mini version) | Lightweight testing on a single GPU/TPU |

All three variants share the same core architecture but differ in spatial resolution and computational requirements.

## The Three-Stage Inference Pipeline

According to the source code in [`weathernext/weathernext2/architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/architecture.py) and [`weathernext/utils/rollout.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/rollout.py), WeatherNext inference follows a strict three-stage pipeline:

### Stage 1: Data Preparation

Load and preprocess initial condition datasets—typically ERA5 reanalysis or HRES operational data—using xarray and Zarr formats. The utilities in [`weathernext/utils/data_utils.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_utils.py) handle **normalization**, **masking**, and **modality conversion**.

### Stage 2: Model Construction

Instantiate the `ForwardPass` Haiku module defined at lines 49-89 of [`weathernext/weathernext2/architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/architecture.py). This module wires together four key components:

- `encode_mesh_spatial_features`: Builds learnable representations on an icosahedral mesh
- **Points-to-mesh transformer**: Converts grid-based observations to mesh features
- `mesh_model_ctor`: A **mesh-GNN** for message passing between spatial locations
- **Mesh-to-grid decoder**: Outputs an `xarray.Dataset` matching your target template's spatial layout

### Stage 3: Auto-Regressive Rollout

Execute the forecast loop using `autoregressive_rollout` from [`weathernext/utils/rollout.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/rollout.py) (lines 168-210). This function iterates over your forecast horizon, feeding predictions back as inputs, while handling **sharding**, **batched sampling**, and **ensembling**.

## Step-by-Step Reproduction Instructions

### 1. Install the Package

Pin to a specific release tag to ensure reproducibility:

```bash
pip install git+https://github.com/google-deepmind/weathernext.git@v0.3.0

```

The README at lines 34-38 recommends version pinning for stable results.

### 2. Download Pretrained Weights

Google DeepMind hosts checkpoints in the `dm_graphcast` Cloud bucket:

```bash
gsutil -m cp -r gs://dm_graphcast/WeatherNext2_<2025>* .

```

Checkpoint filenames follow the pattern `WeatherNext2_<2025_model{1,2,3,4}.npz` as documented at lines 57-62 of [`README.md`](https://github.com/google-deepmind/weathernext/blob/main/README.md).

### 3. Obtain Input Data

**ERA5 (recommended for research):**

```python
import xarray as xr
import weatherbench2 as wb2

ds = wb2.load_dataset("era5", split="train")  # or "hres" for operational data

```

**HRES (operational):** Also available through WeatherBench2's HRES driver per lines 55-61 of [`README.md`](https://github.com/google-deepmind/weathernext/blob/main/README.md).

### 4. Build and Initialize the Model

```python
import weathernext.weathernext2.architecture as arch
import weathernext.utils.model_utils as mu
import jax

# Load pretrained checkpoint

ckpt = mu.load_checkpoint("WeatherNext2_<2025_model1.npz")

# Instantiate ForwardPass with architecture hyperparameters

forward = arch.ForwardPass(
    latent_dense_kwargs=dict(width=256, num_layers=4),
    output_dense_kwargs=dict(width=256, num_layers=2, output_size=ckpt["output_dim"]),
    spatial_features_kwargs=dict(),
    mesh_num_splits=5,
    points_to_mesh_model_ctor=mu.default_points_to_mesh_ctor(),
    mesh_model_ctor=mu.default_mesh_gnn_ctor(),
    mesh_to_grid_model_ctor=mu.default_mesh_to_grid_ctor(),
)

# Initialize parameters with dummy inputs matching checkpoint shapes

params = forward.init(
    jax.random.PRNGKey(0),
    inputs=ckpt["sample_inputs"],
    targets_template=ckpt["sample_targets"],
    forcings=ckpt["sample_forcings"],
    is_training=False
)

```

Constructor arguments map directly to the class header documentation at lines 54-73 of [`architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/architecture.py).

### 5. Run the Auto-Regressive Rollout

```python
from weathernext.utils import rollout

predictions = rollout.autoregressive_rollout(
    params=params,
    forward=forward,
    inputs=initial_state,      # Your ERA5/HRES initial condition

    forcings=static_forcings,   # Static fields like orography

    horizon=14,                # 14-day forecast

    sample_kwargs=dict(num_samples=4),  # 4-member ensemble

)

```

The rollout utility manages chunking and the `MultipleRolloutsLooped` ensembling logic internally.

### 6. Visualize Results

```python
import matplotlib.pyplot as plt

# Extract 2-meter temperature at forecast day 0

preds = predictions["t2m"]
plt.contourf(preds.lat, preds.lon, preds.isel(time=0), cmap="coolwarm")
plt.title("Day 0 2 m temperature")
plt.show()

```

### 7. Optional: Run Cyclone Tracking

For tropical cyclone applications, feed forecast fields to the direct tracker:

```python
from weathernext.cyclones import direct_tracker

tracks = direct_tracker.run_tracker(predictions)

```

The tracker implementation resides in [`weathernext/cyclones/direct_tracker.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/cyclones/direct_tracker.py) (lines 1-30).

## Complete End-to-End Example

This single script reproduces a full WeatherNext 2 run on one TPU device:

```python
import jax
import xarray as xr
from weathernext.weathernext2 import architecture as arch
from weathernext.utils import rollout, model_utils as mu

# 1️⃣ Load initial condition and static forcings

init_ds = xr.open_zarr("gs://weatherbench2/era5/train/0")
forcings = xr.open_zarr("gs://weatherbench2/hres/forcings")

# 2️⃣ Build model architecture

forward = arch.ForwardPass(
    latent_dense_kwargs=dict(width=256, num_layers=4),
    output_dense_kwargs=dict(width=256, num_layers=2, output_size=71),
    spatial_features_kwargs=dict(),
    mesh_num_splits=5,
    points_to_mesh_model_ctor=mu.default_points_to_mesh_ctor(),
    mesh_model_ctor=mu.default_mesh_gnn_ctor(),
    mesh_to_grid_model_ctor=mu.default_mesh_to_grid_ctor(),
)

# 3️⃣ Initialize from checkpoint

ckpt = mu.load_checkpoint("WeatherNext2_<2025_model1.npz")
params = forward.init(
    jax.random.PRNGKey(42),
    inputs=ckpt["sample_inputs"],
    targets_template=ckpt["sample_targets"],
    forcings=ckpt["sample_forcings"],
    is_training=False
)

# 4️⃣ Execute 14-day rollout with 8-member ensemble

forecast = rollout.autoregressive_rollout(
    params=params,
    forward=forward,
    inputs=init_ds,
    forcings=forcings,
    horizon=14,
    sample_kwargs=dict(num_samples=8),
)

# 5️⃣ Visualize Day-7 temperature

import matplotlib.pyplot as plt
temp_day7 = forecast["t2m"].isel(time=7)
plt.figure(figsize=(8, 4))
temp_day7.plot(cmap="coolwarm")
plt.title("Day 7 2-m temperature")
plt.show()

```

## Quick Start with the Official Demo

For immediate results without local setup, use the **Colab demo notebook** at `docs/weathernext2/wn2_demo.ipynb`. It automatically:

- Downloads pretrained weights from the public Google Cloud bucket
- Loads sample ERA5 initial conditions
- Runs a 14-day rollout (TPU `v5e-1` for Mini, `v5p` for full resolution)
- Visualizes temperature, wind speed, geopotential height, and cyclone tracks

## Key Source Files Reference

| File | Purpose |
|------|---------|
| [`weathernext/weathernext2/architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/architecture.py) | `ForwardPass` Haiku module definition (lines 49-89) |
| [`weathernext/utils/rollout.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/rollout.py) | `autoregressive_rollout` and chunking logic (lines 168-210) |
| [`weathernext/utils/model_utils.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/model_utils.py) | Checkpoint loading and default model constructors |
| [`weathernext/utils/data_utils.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_utils.py) | Normalization, derived variables, masking utilities |
| [`weathernext/cyclones/direct_tracker.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/cyclones/direct_tracker.py) | Best-track cyclone forecasting (lines 1-30) |
| `docs/weathernext2/wn2_demo.ipynb` | End-to-end Colab demonstration |

## Summary

- **Install** with `pip install git+https://github.com/google-deepmind/weathernext.git@v0.3.0`
- **Download weights** via `gsutil` from `gs://dm_graphcast/`
- **Prepare data** using WeatherBench2's ERA5 or HRES loaders
- **Build** the `ForwardPass` architecture from [`weathernext/weathernext2/architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/architecture.py)
- **Run** `autoregressive_rollout` from [`weathernext/utils/rollout.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/rollout.py) with your target horizon
- **Track cyclones** optionally with `direct_tracker.run_tracker`

Following this pipeline exactly reproduces the evaluation metrics reported in the WeatherNext 2 and WeatherNext Cyclones papers.

## Frequently Asked Questions

### What hardware do I need to reproduce WeatherNext results?

The full-resolution WeatherNext 2 model requires TPU v5p for operational forecasts. WeatherNext Cyclones Mini runs on a single TPU v5e-1 or GPU, making it suitable for experimentation. The Colab demo notebook automatically selects appropriate runtime configurations.

### Can I use my own meteorological data instead of ERA5?

Yes. The [`weathernext/utils/data_utils.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_utils.py) utilities accept any xarray Dataset with compatible coordinate names (latitude, longitude, time, pressure levels). Ensure your variables match the expected input modality—temperature, geopotential, wind components, humidity, and surface pressure.

### How do I generate ensemble forecasts with WeatherNext?

Pass `sample_kwargs=dict(num_samples=N)` to `autoregressive_rollout` where N is your desired ensemble size. The function handles batched sampling and returns a forecast with an additional `sample` dimension. Set `N=1` for deterministic single-run forecasts.

### Where are the training scripts for WeatherNext?

The published repository focuses on **inference and reproduction** of pretrained results. Training implementations remain internal to Google DeepMind. For research purposes, the architecture definitions in [`weathernext/weathernext2/architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/architecture.py) provide sufficient detail to reconstruct the training objective described in the papers.