anshchkse/MR-GNF-Britain-Centric-Euro-Atlantic

MR-GNF: Britain-Centric Euro-Atlantic Weather Forecasting

MR-GNF is a lightweight regional graph neural weather forecasting model for the UK–Ireland sector and its surrounding Euro-Atlantic context. The model performs one-step atmospheric prediction and supports multi-step forecasting through autoregressive rollout.

This Hugging Face repository contains the released model checkpoints and the minimum code and static artifacts needed to load them. The complete training, data-preparation, and evaluation pipeline is available in the full MR-GNF project repository.

The model accompanies the paper:

Andrii Shchur and Inna Skarga-Bandurova, “MR-GNF: Multi-Resolution Graph Neural Forecasting on Ellipsoidal Meshes for Efficient Regional Weather Prediction.”
arXiv:2603.13563

Available checkpoints

Checkpoint	Description
mr-gnf-general	Main checkpoint for forecasting the complete supported atmospheric state.
mr-gnf-wind	Checkpoint fine-tuned for 10 m wind components: u10@sfc and v10@sfc.
mr-gnf-precipitation	Checkpoint fine-tuned for total precipitation: tp_log@sfc.

The wind and precipitation checkpoints are optional specialist models. During composed inference, their predictions can replace the corresponding channels produced by the general model before the next autoregressive step.

Repository files

File	Description
mr-gnf-general	Main all-variable model checkpoint.
mr-gnf-wind	Wind-specialist checkpoint.
mr-gnf-precipitation	Precipitation-specialist checkpoint.
UK_graph_static_from_stats_vgeo_aligned	Static graph artifact required by the dataloader and model.
UK_geo_pos_static.npz	Geographic coordinates for the regular UK-centred output grid.
stats_025deg_1980_2013.npz	Channel order and normalization statistics derived from the 1980–2013 training period.
unified_graph_weather_dataloader.py	Dataset and Lightning data-module implementation for graph shards.
unified_graph_weather_model.py	General MR-GNF model definition.
unified_graph_weather_model_uv10.py	Wind-specialist model definition.
unified_graph_weather_model_tp.py	Precipitation-specialist model definition.

Time-dependent meteorological input shards are not included in this model repository.

Installation

A CUDA-enabled PyTorch environment is recommended.

pip install torch lightning numpy huggingface_hub
pip install torch-geometric triton

Install versions of PyTorch, CUDA, PyTorch Geometric, and Triton that are mutually compatible with your system.

Download the model repository

from huggingface_hub import snapshot_download

repo_dir = snapshot_download(
    repo_id="YOUR_HF_USERNAME/YOUR_MODEL_REPOSITORY"
)

print(repo_dir)

Replace YOUR_HF_USERNAME/YOUR_MODEL_REPOSITORY with the published Hugging Face repository identifier.

Input requirements

The dataloader expects monthly graph shards such as:

graph_2024-01.pt
graph_2024-02.pt
...

A batch is returned as:

(x_all, geo, pos2d), y_all

with the principal tensors arranged as:

x_all: [batch, input_steps, channels, nodes]
y_all: [batch, forecast_steps, channels, nodes]

Input shards must preserve the channel ordering, node ordering, normalization, and temporal cadence used during training. The channel order stored in the shard metadata is authoritative.

Load the general model

import sys
from pathlib import Path

import torch

REPO_DIR = Path("/path/to/downloaded/model-repository")
SHARDS_ROOT = Path("/path/to/monthly-graph-shards")

sys.path.insert(0, str(REPO_DIR))

from unified_graph_weather_dataloader import GraphWeatherDataModule
from unified_graph_weather_model import UnifiedGraphWeatherGATLightning

STATIC_GRAPH = REPO_DIR / "UK_graph_static_from_stats_vgeo_aligned"
GENERAL_CKPT = REPO_DIR / "mr-gnf-general"

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def load_checkpoint(model, checkpoint_path: Path):
    checkpoint = torch.load(checkpoint_path, map_location="cpu")
    state_dict = checkpoint.get("state_dict", checkpoint)

    # These graph-dependent buffers are rebuilt from the supplied static graph.
    state_dict = {
        key: value
        for key, value in state_dict.items()
        if not (key.endswith("edge_src") or key.endswith("edge_dst"))
    }

    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    print("Missing keys:", missing)
    print("Unexpected keys:", unexpected)


dm = GraphWeatherDataModule(
    shards_root_train=str(SHARDS_ROOT),
    shards_root_val=str(SHARDS_ROOT),
    years_train=range(2024, 2025),
    years_val=range(2024, 2025),
    static_graph_npz=str(STATIC_GRAPH),
    force_tin_tout=None,
    dtype_out="float32",
    shuffle_within_file_train=False,
    shuffle_within_file_val=False,
    batch_size_train=1,
    batch_size_val=1,
    num_workers=0,
    pin_memory=DEVICE.type == "cuda",
    persistent_workers=False,
    device="cuda" if DEVICE.type == "cuda" else "cpu",
)

dm.setup()

model = UnifiedGraphWeatherGATLightning(
    C=dm.C,
    Tin=dm.Tin,
    Tout=1,
    edge_index_base=dm.edge_index_base,
    N_mesh=dm.N_mesh,
    embed_dim=192,
    blocks=4,
    heads_v=4,
    heads_xy=4,
    attn_drop_xy=0.0,
    lr=3e-4,
    weight_decay=1e-4,
    loss="mse",
).to(DEVICE).eval()

load_checkpoint(model, GENERAL_CKPT)

The expected missing keys are the graph-dependent buffers regenerated from the supplied static graph:

net.edge_src
net.edge_dst

Unexpected missing parameters or unexpected checkpoint keys should be investigated before inference.

Run one-step inference

val_loader = dm.val_dataloader()
(x_all, geo, pos2d), _ = next(iter(val_loader))

x_all = x_all.to(DEVICE)
geo = geo.to(DEVICE)
pos2d = pos2d.to(DEVICE)

steps = torch.zeros(x_all.shape[0], device=DEVICE)

with torch.inference_mode():
    with torch.autocast(
        device_type=DEVICE.type,
        dtype=torch.bfloat16,
        enabled=DEVICE.type == "cuda",
    ):
        prediction = model.net(
            (x_all, geo, pos2d),
            diffusion_step=steps,
        )

print(prediction.shape)
# [batch, 1, channels, nodes]

For an autoregressive forecast, append the predicted state to the input sequence, remove the oldest input step, and call the model again.

Load the specialist models

The specialist checkpoints use the same static graph and input tensors as the general model. They also require the original channel order.

from unified_graph_weather_model_uv10 import WindUV10Lightning
from unified_graph_weather_model_tp import PrecipTPLightning

channel_order = dm.channel_order

wind_model = WindUV10Lightning(
    C=dm.C,
    Tin=dm.Tin,
    Tout=1,
    edge_index_base=dm.edge_index_base,
    N_mesh=dm.N_mesh,
    embed_dim=192,
    blocks=4,
    heads_v=4,
    heads_xy=4,
    attn_drop_xy=0.0,
    lr=1e-3,
    weight_decay=1e-4,
    loss_components=1.0,
    loss_magnitude=0.0,
    loss_direction=0.0,
    robust_loss=False,
    chan_u_name="u10@sfc",
    chan_v_name="v10@sfc",
    channel_order=channel_order,
    diffusion_max_step=0,
).to(DEVICE).eval()

precip_model = PrecipTPLightning(
    C=dm.C,
    Tin=dm.Tin,
    Tout=1,
    edge_index_base=dm.edge_index_base,
    N_mesh=dm.N_mesh,
    embed_dim=192,
    blocks=4,
    heads_v=4,
    heads_xy=4,
    attn_drop_xy=0.0,
    lr=3e-4,
    weight_decay=1e-4,
    diffusion_max_step=0,
    channel_order=channel_order,
).to(DEVICE).eval()

load_checkpoint(wind_model, REPO_DIR / "mr-gnf-wind")
load_checkpoint(precip_model, REPO_DIR / "mr-gnf-precipitation")

The wind model outputs the two near-surface wind channels. The precipitation model outputs the precipitation channel. Their outputs can be inserted into the general model prediction before continuing an autoregressive rollout.

Normalization

stats_025deg_1980_2013.npz contains the channel order, mean, and standard deviation used for normalization.

import numpy as np

stats = np.load(REPO_DIR / "stats_025deg_1980_2013.npz", allow_pickle=True)
channel_order = stats["order"].tolist()
mean = stats["mean"]
std = stats["std"]

physical_values = normalized_values * std + mean

Precipitation is represented in log space and can be converted back with:

precipitation = np.maximum(np.expm1(tp_log), 0.0)

Intended use

This release is intended for research and reproducibility work on short-range regional weather forecasting over the UK–Ireland and surrounding Euro-Atlantic sector.

It is not an operational warning system and should not replace forecasts or warnings issued by national meteorological services.

Citation

Please cite the associated paper when using the model or artifacts:

@misc{shchur2026mrgnf,
  title         = {MR-GNF: Multi-Resolution Graph Neural Forecasting on Ellipsoidal Meshes for Efficient Regional Weather Prediction},
  author        = {Andrii Shchur and Inna Skarga-Bandurova},
  year          = {2026},
  eprint        = {2603.13563},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  doi           = {10.48550/arXiv.2603.13563},
  url           = {https://arxiv.org/abs/2603.13563}
}

Paper: https://arxiv.org/abs/2603.13563

Full project repository: https://github.com/AndriiShchur/MR-GNF

License

The model and accompanying code are released under the MIT License. The underlying meteorological data may be subject to separate terms and licenses.