README.md

---
license: cc-by-4.0
tags:
- weather
- forecasting
- climate
- graph-neural-network
- arxiv:2506.10772
library_name: transformers
pipeline_tag: other
---

# WeatherNext 2

WeatherNext 2 is a global medium-range weather forecasting model from Google DeepMind and Google Research. One forward
pass advances the state of the atmosphere by 6 hours on a 0.25° latitude/longitude grid, predicting 13 pressure levels
of temperature, geopotential, wind and humidity together with surface fields, precipitation, 100m winds and a set of
tropical-cyclone diagnostics.

It is a *Functional Generative Network* (FGN): rather than injecting a noise field into the input or running a
diffusion sampler, the model draws a single 32-dimensional noise vector per ensemble member and uses it to modulate
the scale and offset of **every** normalization layer. One draw gives one self-consistent forecast, and an ensemble is
simply several draws — which here is just the batch dimension.

| Variant | Resolution | Mesh nodes | Params | 100m winds |
|---|---|---|---|---|
| [WeatherNext2](https://huggingface.co/kashif/weathernext2) | 0.25° (721×1440) | 40,962 | 183.8M | yes |
| [WeatherNextCyclones](https://huggingface.co/kashif/weathernext-cyclones) | 0.25° (721×1440) | 40,962 | 183.8M | no |

These weights correspond to `WeatherNext2_<2025_model{1..4}`, trained on data through 2024 and fine-tuned for
initialization from operational ECMWF HRES analysis. All four independently trained members (`model1`–`model4`) are included; see [Ensembles](#ensembles).

## Usage

> [!NOTE]
> WeatherNext 2 support is not in a released version of Transformers yet. Until
> [huggingface/transformers#47874](https://github.com/huggingface/transformers/pull/47874) is merged, install from the
> branch:
>
> ```bash
> pip install "git+https://github.com/kashif/transformers.git@add-weathernext2" torch scipy
> ```

The model works in a normalized space; `WeatherNext2FeatureExtractor` owns everything physical — the per-variable
normalization statistics, the calendar-derived forcings, and the residual connection back to an atmospheric state.

```python
import numpy as np
import torch
from transformers import WeatherNext2ForWeatherForecasting, WeatherNext2FeatureExtractor

model = WeatherNext2ForWeatherForecasting.from_pretrained("kashif/weathernext2", device_map="auto").eval()
processor = WeatherNext2FeatureExtractor.from_pretrained("kashif/weathernext2")

# `state` maps each input variable to its values, e.g. from an xarray Dataset of HRES analysis.
# Time-varying variables are [batch, 2, (levels,) lat, lon]; static ones are [lat, lon].
state = {name: ... for name in processor.input_variables}
valid_time = np.array([np.datetime64("2024-10-07T06:00:00").astype("datetime64[s]").astype(np.int64)])

inputs = processor(state, seconds_since_epoch=valid_time).to(model.device)
with torch.no_grad():
    outputs = model(**inputs, generator=torch.Generator().manual_seed(0))

forecast = processor.postprocess(outputs.prediction, state)
print(forecast["2m_temperature"].shape)  # (1, 721, 1440)
```

### Autoregressive rollout

Each 6-hour step draws fresh noise. `advance_state` drops the oldest frame, appends the forecast, recomputes the clock
variables, and discards targets that are not also inputs (precipitation and the cyclone diagnostics).

```python
step_seconds = processor.time_step_hours * 3600
for step in range(20):  # 5 days
    inputs = processor(state, seconds_since_epoch=valid_time).to(model.device)
    with torch.no_grad():
        outputs = model(**inputs)
    forecast = processor.postprocess(outputs.prediction, state)
    # `valid_time` is the time this forecast is valid at, so it stamps the appended frame first.
    state = processor.advance_state(state, forecast, valid_time)
    valid_time = valid_time + step_seconds
```

## Ensembles

There are two independent ensembles here, and the operational product combines both.

### 1. Noise ensemble (within one checkpoint)

This is the FGN mechanism: each member is one draw of the 32-dimensional noise vector through the *same* weights.
Members are fully independent, so the usual way to run them - and what the reference implementation does, one member
per device - is to loop, seeding each draw from its own index:

```python
inputs = processor(state, seconds_since_epoch=valid_time).to(model.device)

members = 8
predictions = []
for member in range(members):
    noise = torch.randn(1, model.config.noise_channels, generator=torch.Generator().manual_seed(member))
    with torch.no_grad():
        predictions.append(model(**inputs, noise=noise.to(model.device)).prediction)
```

Seeding per member (rather than drawing from one stream) means the first N members are reproducible regardless of how
many you end up running — the same property the original implementation gets from `jax.random.fold_in`.

At 0.25° a member needs roughly 50 GB, so looping is normally the only option. Where the whole ensemble does fit,
the members can instead ride on the batch axis:

```python
batched = {key: value.repeat(members, *([1] * (value.ndim - 1))) for key, value in inputs.items()}
with torch.no_grad():
    outputs = model(**batched, generator=torch.Generator().manual_seed(0))
# outputs.prediction is (members, channels, lat, lon)
```

The batch axis serves double duty: it carries ensemble members here, and independent initialization times when several
forecasts are run together. The reference implementation keeps these as separate `sample` and `batch` dimensions.

### 2. Multi-model ensemble (across checkpoints)

The release is four separate training runs rather than one checkpoint, and upstream ships them as four weight files
with no code to combine them. Run 1 is at the repository root; all four are also available as subfolders, so you can
loop uniformly. How you pool them is your choice; what follows is one reasonable recipe.

```python
REPO = "kashif/weathernext2"

def load_member(member: int, revision: str = "main"):
    return WeatherNext2ForWeatherForecasting.from_pretrained(
        REPO, subfolder=f"model{member}", revision=revision, device_map="auto"
    ).eval()
```

`subfolder` composes with `revision`, and works the same way for `WeatherNext2FeatureExtractor` and `AutoConfig` — each
subfolder carries its own `config.json` and `preprocessor_config.json`. The processors are identical across members,
so loading one is enough.

### Putting them together

The full ensemble is `num_models × num_noise_draws` trajectories. Loading one member at a time keeps peak memory at
roughly one model:

```python
import numpy as np
import torch

processor = WeatherNext2FeatureExtractor.from_pretrained(REPO)
raw_inputs = processor(state, seconds_since_epoch=valid_time)

forecasts = []
for member in range(1, 5):
    model = load_member(member)
    inputs = raw_inputs.to(model.device)
    for draw in range(4):
        noise = torch.randn(
            1, model.config.noise_channels,
            generator=torch.Generator().manual_seed(1000 * member + draw),
        ).to(model.device)
        with torch.no_grad():
            prediction = model(**inputs, noise=noise).prediction
        forecasts.append(processor.postprocess(prediction, state)["2m_temperature"])
    del model  # free before loading the next member

stack = np.concatenate(forecasts, axis=0)   # (16, lat, lon)
ensemble_mean = stack.mean(axis=0)
ensemble_spread = stack.std(axis=0)
```

For multi-step forecasts each trajectory carries its own state, so keep one `state` per member and advance them
separately (or keep members on the batch axis, which `advance_state` handles for you).

## Model details

- **Architecture**: encode–process–decode graph network. The lat/lon grid is encoded, projected onto an icosahedral
  mesh by a graph network (ball-query connectivity), processed by a 24-layer transformer whose attention is restricted
  to a 32-hop neighbourhood on the mesh, projected back (in-triangle connectivity), and decoded.
- **Hidden size / layers / heads**: 768 / 24 / 6, feed-forward 3072.
- **Mesh**: icosahedron refined 6 times → 40,962 nodes; ~1.6M grid→mesh and ~3.1M mesh→grid edges.
- **Positional information**: none learned. Position is carried entirely by the mesh geometry and the attention mask,
  both of which are rebuilt deterministically from the config at load time and cached on disk.
- **Inputs**: two frames 6h apart, 13 pressure levels, plus static fields and calendar forcings.
- **Time step**: 6 hours.

## Evaluation

For scorecards see the [technical report](https://huggingface.co/papers/2506.10772) and
[WeatherBench 2](https://sites.research.google/weatherbench/).

Note that this variant takes 100m winds as inputs, which the publicly published sample datasets do not contain, so the
port was verified end-to-end using the sibling
[WeatherNextCyclones](https://huggingface.co/kashif/weathernext-cyclones) checkpoint — identical architecture, same
resolution, no 100m winds. Weight conversion for this checkpoint was verified structurally (all 488 parameter arrays
mapped, 183.8M parameters, no missing or unexpected keys).

## Hardware

A single ensemble member at 0.25° needs roughly 50 GB. The forward pass above took 96 s on CPU; a modern GPU is much
faster. The mesh and graph construction takes a few minutes the first time and is then cached under `HF_HOME`.

## Limitations

This is a research model, not an operational warning system. It does not replace official alerts from national
meteorological agencies. It was trained on ERA5 and fine-tuned on HRES analysis, and is designed to be initialized
from HRES initial conditions rather than reanalysis.

## License

The weights are released by Google DeepMind under [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/). The
original code is Apache-2.0. Original repository: [google-deepmind/weathernext](https://github.com/google-deepmind/weathernext).

## Acknowledgements

Data and products of the European Centre for Medium-range Weather Forecasts (ECMWF), as modified by Google. Modified
Copernicus Climate Change Service information 2023. Neither the European Commission nor ECMWF is responsible for any
use that may be made of the Copernicus information or data it contains. ECMWF HRES datasets copyright statement:
Copyright "© 2023 European Centre for Medium-Range Weather Forecasts (ECMWF)". Source:
[www.ecmwf.int](http://www.ecmwf.int/). License statement: ECMWF open data is published under a Creative Commons
Attribution 4.0 International (CC BY 4.0), [https://creativecommons.org/licenses/by/4.0/](https://creativecommons.org/licenses/by/4.0/).
Disclaimer: ECMWF does not accept any liability whatsoever for any error or omission in the data, their availability,
or for any loss or damage arising from their use.

## Citation

```bibtex
@article{alet2025skillful,
  title={Skillful joint probabilistic weather forecasting from marginals},
  author={Alet, Ferran and Price, Ilan and El-Kadi, Andrew and Masters, Dominic and Markou, Stratis and Andersson, Tom R and Stott, Jacklynn and Lam, Remi and Willson, Matthew and Sanchez-Gonzalez, Alvaro and Battaglia, Peter},
  journal={arXiv preprint arXiv:2506.10772},
  year={2025}
}
```