Views
No views yet
| 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 |
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.[!NOTE] 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
WeatherNext2FeatureExtractor owns everything physical — the per-variable
normalization statistics, the calendar-derived forcings, and the residual connection back to an atmospheric state.1import numpy as np
2import torch
3from transformers import WeatherNext2ForWeatherForecasting, WeatherNext2FeatureExtractor
4
5model = WeatherNext2ForWeatherForecasting.from_pretrained("kashif/weathernext2", device_map="auto").eval()
6processor = WeatherNext2FeatureExtractor.from_pretrained("kashif/weathernext2")
7
8# `state` maps each input variable to its values, e.g. from an xarray Dataset of HRES analysis.
9# Time-varying variables are [batch, 2, (levels,) lat, lon]; static ones are [lat, lon].
10state = {name: ... for name in processor.input_variables}
11valid_time = np.array([np.datetime64("2024-10-07T06:00:00").astype("datetime64[s]").astype(np.int64)])
12
13inputs = processor(state, seconds_since_epoch=valid_time).to(model.device)
14with torch.no_grad():
15 outputs = model(**inputs, generator=torch.Generator().manual_seed(0))
16
17forecast = processor.postprocess(outputs.prediction, state)
18print(forecast["2m_temperature"].shape) # (1, 721, 1440)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).1step_seconds = processor.time_step_hours * 3600
2for step in range(20): # 5 days
3 inputs = processor(state, seconds_since_epoch=valid_time).to(model.device)
4 with torch.no_grad():
5 outputs = model(**inputs)
6 forecast = processor.postprocess(outputs.prediction, state)
7 # `valid_time` is the time this forecast is valid at, so it stamps the appended frame first.
8 state = processor.advance_state(state, forecast, valid_time)
9 valid_time = valid_time + step_seconds1inputs = processor(state, seconds_since_epoch=valid_time).to(model.device)
2
3members = 8
4predictions = []
5for member in range(members):
6 noise = torch.randn(1, model.config.noise_channels, generator=torch.Generator().manual_seed(member))
7 with torch.no_grad():
8 predictions.append(model(**inputs, noise=noise.to(model.device)).prediction)jax.random.fold_in.1batched = {key: value.repeat(members, *([1] * (value.ndim - 1))) for key, value in inputs.items()}
2with torch.no_grad():
3 outputs = model(**batched, generator=torch.Generator().manual_seed(0))
4# outputs.prediction is (members, channels, lat, lon)sample and batch dimensions.1REPO = "kashif/weathernext2"
2
3def load_member(member: int, revision: str = "main"):
4 return WeatherNext2ForWeatherForecasting.from_pretrained(
5 REPO, subfolder=f"model{member}", revision=revision, device_map="auto"
6 ).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.num_models × num_noise_draws trajectories. Loading one member at a time keeps peak memory at
roughly one model:1import numpy as np
2import torch
3
4processor = WeatherNext2FeatureExtractor.from_pretrained(REPO)
5raw_inputs = processor(state, seconds_since_epoch=valid_time)
6
7forecasts = []
8for member in range(1, 5):
9 model = load_member(member)
10 inputs = raw_inputs.to(model.device)
11 for draw in range(4):
12 noise = torch.randn(
13 1, model.config.noise_channels,
14 generator=torch.Generator().manual_seed(1000 * member + draw),
15 ).to(model.device)
16 with torch.no_grad():
17 prediction = model(**inputs, noise=noise).prediction
18 forecasts.append(processor.postprocess(prediction, state)["2m_temperature"])
19 del model # free before loading the next member
20
21stack = np.concatenate(forecasts, axis=0) # (16, lat, lon)
22ensemble_mean = stack.mean(axis=0)
23ensemble_spread = stack.std(axis=0)state per member and advance them
separately (or keep members on the batch axis, which advance_state handles for you).HF_HOME.1@article{alet2025skillful,
2 title={Skillful joint probabilistic weather forecasting from marginals},
3 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},
4 journal={arXiv preprint arXiv:2506.10772},
5 year={2025}
6}