kashif/weathernext2

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	0.25° (721×1440)	40,962	183.8M	yes
WeatherNextCyclones	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.

Usage

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

pip install "git+https://github.com/kashif/transformers.git@add-weathernext2" torch accelerate

The model predicts normalized fields. WeatherNext2FeatureExtractor handles normalization, clock forcings and conversion back to physical units. Initialize from two chronological HRES analysis frames, six hours apart, with ascending latitude, the checkpoint's pressure levels and all configured input fields, including static fields. See the official demo and the model documentation for data preparation. The example below assumes those physical fields have been loaded into state.

This checkpoint also requires 100m wind inputs, which are absent from the published sample datasets.

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

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

# Replace the ellipses with physical HRES fields:
# time-varying fields: [batch, 2, (levels,) lat, lon]; static fields: [lat, lon].
state = {
    name: ...
    for name in processor.input_variables
    if name not in processor.forcing_variables
}
valid_time = np.array([np.datetime64("2024-10-07T06:00:00").astype("datetime64[s]").astype(np.int64)])
conditioning_times = valid_time[:, None] - (
    processor.time_step_hours * 3600 * np.arange(processor.num_input_timesteps, 0, -1)
)
state.update(processor.compute_forcings(conditioning_times))

inputs = processor(
    state, seconds_since_epoch=valid_time, generator=torch.Generator().manual_seed(0)
).to(device=model.device, dtype=model.dtype)
with torch.no_grad():
    outputs = model(**inputs)
forecast = processor.postprocess(outputs.prediction, state)
print(forecast["2m_temperature"].shape)  # torch.Size([1, 721, 1440])

Noise is sampled by the processor, not by the model forward. The processor includes both noise and the shared geometry attention_mask in inputs; do not pass a second noise argument or repeat the mask over ensemble members.

Autoregressive rollout

generate() prepares inputs, advances each member using its own predicted state, and samples fresh noise at each six-hour step. seconds_since_epoch is the valid time of the first forecast, not the last conditioning frame.

rollout = model.generate(
    state=state,
    feature_extractor=processor,
    seconds_since_epoch=valid_time,
    num_steps=20,  # 5 days
    generator=torch.Generator().manual_seed(0),
)
frames = rollout.forecasts
next_state = rollout.state

Ensembles

Noise ensemble

Members share weights but receive independent noise. generate(num_members=8) repeats each initial condition and runs the resulting members along the batch axis. batch_size=1 limits the rows evaluated together; returned forecasts and conditioning states still contain the whole ensemble.

ensemble = model.generate(
    state=state,
    feature_extractor=processor,
    seconds_since_epoch=valid_time,
    num_steps=1,
    num_members=8,
    batch_size=1,
    generator=torch.Generator().manual_seed(0),
)
print(ensemble.forecasts[0]["2m_temperature"].shape)  # torch.Size([8, 721, 1440])

For multiple initial conditions, member m of condition i is row i * num_members + m. A direct forward can use processor(..., num_members=8) instead. Step/member defaults can also be set in model.generation_config.

Multi-model ensemble

The root checkpoint is training run 1; model1–model4 subfolders contain all four runs. Each subfolder includes its own config and processor assets. Combine independent noise draws from each run to pool both ensemble types.

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

# Release the single-run model used above before loading the ensemble runs.
del model
forecasts = []
for member in range(1, 5):
    member_model = load_member(member)
    member_output = member_model.generate(
        state=state,
        feature_extractor=processor,
        seconds_since_epoch=valid_time,
        num_steps=1,
        num_members=4,
        batch_size=1,
        generator=torch.Generator().manual_seed(member),
    )
    temperature = member_output.forecasts[0]["2m_temperature"].cpu()
    forecasts.append(temperature.unflatten(0, (len(valid_time), 4)))
    del member_model, member_output

stack = torch.cat(forecasts, dim=1)  # [initial_batch, 16, lat, lon]
ensemble_mean = stack.mean(dim=1)
ensemble_spread = stack.std(dim=1, correction=0)

Move forecasts to the CPU before pooling them to avoid retaining every run's outputs on the accelerator.

Fine-tuning

Pass physical next-step targets with shape [batch, (levels,) lat, lon] and at least two ensemble members. The processor produces normalized labels; the model returns the original unbiased ensemble CRPS.

model = WeatherNext2ForWeatherForecasting.from_pretrained(repo, device_map="auto")
train_inputs = processor(
    state,
    seconds_since_epoch=valid_time,
    targets=targets,
    num_members=2,
    generator=torch.Generator().manual_seed(0),
).to(device=model.device, dtype=model.dtype)
model.train()
loss = model(**train_inputs).loss
loss.backward()

The loss uses latitude and pressure-level weights plus the checkpoint's loss_weights, vars_with_nan_targets and loss_for_nan_targets. Missing cyclone targets are masked in both CRPS terms. Training memory requirements differ from inference; the optional Hub kernel optimizations below are for inference.

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. Exact geometry is stored as persistent buffers in model.geometry; the processor loads the separate attention_mask.safetensors asset and passes the shared mask to the model.
• 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 and WeatherBench 2.

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 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 and optional kernels

Memory and runtime depend on resolution, dtype, attention implementation, ensemble size and device. Geometry and the attention mask are included in the checkpoint assets; loading does not rebuild the mesh and does not require SciPy or trimesh.

For supported devices, optional Hub kernels tile the grid encoder, grid-to-mesh and mesh-to-grid operations, and forecast head. Install kernels, then enable them when loading:

model = WeatherNext2ForWeatherForecasting.from_pretrained(
    repo, device_map="auto", use_kernels=True
).eval()

chunk_size controls the inference tile size and can be supplied when loading. Smaller tiles reduce intermediate allocations but may increase overhead. To limit the number of ensemble members evaluated together, use generate(..., batch_size=1); returned forecasts and conditioning states still retain every member.

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. 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. 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{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}
}