kashif/weathernext2-mini

WeatherNext 2 Mini

WeatherNext 2 Mini is the lightweight variant of Google DeepMind's WeatherNext 2, at 1° resolution instead of 0.25°. It forecasts the same quantities as the full model — 13 pressure levels of temperature, geopotential, wind and humidity, surface fields, precipitation, and 16 tropical-cyclone diagnostics — and is intended for local testing and for machines that cannot hold the 0.25° model. It is not expected to match the full model's skill.

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.

The architecture matches the 0.25° models at reduced width and depth. See kashif/weathernext2 and kashif/weathernext-cyclones for the full-resolution versions.

Revisions

Each revision is trained on data up to a different cut-off, so you can evaluate on a year the model has not seen.

Revision	Trained through	Corresponds to
main	2023	WeatherNextCyclones_Mini_<2024
<2023	2022	WeatherNextCyclones_Mini_<2023

model = WeatherNext2ForWeatherForecasting.from_pretrained("kashif/weathernext2-mini", revision="<2023")

Unlike the 0.25° releases, the Mini checkpoints ship a single trained network per cut-off rather than a four-member ensemble, so only the noise ensemble applies; see Ensembles.

Usage

WeatherNext 2 support is not in a released version of Transformers yet. Until huggingface/transformers#47874 is merged, install from the branch:

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.

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

model = WeatherNext2ForWeatherForecasting.from_pretrained("kashif/weathernext2-mini").eval()
processor = WeatherNext2FeatureExtractor.from_pretrained("kashif/weathernext2-mini")

# `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)
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, 181, 360)
print(forecast["cyclone_exists_gaussian_unit_mode"].max())  # cyclone probability field

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).

step_seconds = processor.time_step_hours * 3600
for step in range(20):  # 5 days
    inputs = processor(state, seconds_since_epoch=valid_time)
    with torch.no_grad():
        outputs = model(**inputs)
    forecast = processor.postprocess(outputs.prediction, state)
    valid_time = valid_time + step_seconds
    state = processor.advance_state(state, forecast, valid_time)

Turning the gridded cyclone diagnostics into tracks requires the tracker from the original repository; it is not part of this port.

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 ride on the batch axis and stay independent through an autoregressive rollout.

members = 8
inputs = processor(state, seconds_since_epoch=valid_time)
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)

At 1° all members usually fit in one batch. If memory is tight, looping over draws gives identical results with a constant footprint:

single = processor(state, seconds_since_epoch=valid_time)   # batch of 1
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(**single, noise=noise).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.

The Mini release has one trained network per cut-off, so there is no multi-model ensemble here — the noise ensemble above is the whole ensemble. The 0.25° repositories ship four independently trained members each.

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 16-hop neighbourhood on the mesh, projected back (in-triangle connectivity), and decoded.
• Hidden size / layers / heads: 512 / 16 / 4, feed-forward 2048. 56.7M parameters.
• Mesh: icosahedron refined 5 times → 10,242 nodes, with 16-hop attention.
• 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

A single 6-hour step from ECMWF HRES analysis (initialized 2024-10-07 00:00 UTC), scored against the verifying analysis with latitude weighting, using the main revision at 1° resolution:

Field	WeatherNext 2 Mini	Persistence
2m temperature	0.785 K	2.620 K
Temperature @500hPa	0.414 K	1.197 K
Geopotential @500hPa	31.7 m²/s²	224.5 m²/s²

This is a single case, meant as a smoke test of the port rather than a benchmark. For cyclone track and intensity scorecards see the Nature paper; for general forecast skill see WeatherBench 2.

Hardware

At 1° a forward pass fits comfortably on a small GPU (the original release notes a P100 as sufficient) or on CPU. The mesh and graph construction takes a few seconds the first time and is then cached under HF_HOME.

Limitations

This is a research model, not an operational warning system, and it does not replace official alerts, warnings or notices from national meteorological agencies. It was not produced in collaboration with, nor endorsed by, any government meteorological agency. It is designed to be initialized from operational HRES analysis rather than from reanalysis.

License

The weights are released by Google DeepMind under CC-BY-4.0. The original code is Apache-2.0. Original repository: 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. NOAA's International Best Track Archive for Climate Stewardship (IBTrACS) data, first accessed on 1 Dec 2022. 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. 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/. 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

@article{Alet2026,
  title={Operational Tropical Cyclone Forecasting with AI},
  author={Alet, Ferran and Andersson, Tom R. and Price, Ilan and Markou, Stratis and El-Kadi, Andrew and Masters, Dominic and Li, Amy and Merchant, Samier and Williams, Natalie and Thornton, Gregory and MacKay, Ken and Graham, Olivia and Uddin, Akib and Gaiarin, Ben and Shah, Devaja and Kruse, Elinor and Hogsett, Wallace and Zelinsky, David and Cangialosi, John and Martinez, Jonathan and Franklin, James and DeMaria, Mark and Musgrave, Kate and Bain, Caroline L. and Titley, Helen and Stott, Jacklynn and Lam, Remi and Bell, Aaron and Komarek, Paul and Willson, Matthew and Sanchez-Gonzalez, Alvaro and Battaglia, Peter},
  journal={Nature},
  year={2026},
  issn={1476-4687},
  doi={10.1038/s41586-026-10953-2},
  url={https://doi.org/10.1038/s41586-026-10953-2}
}

@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}
}