---
license: mit
language:
- en
library_name: pytorch
pipeline_tag: time-series-forecasting
tags:
- weather-forecasting
- graph-neural-network
- graph-attention-network
- autoregressive-forecasting
- regional-weather
- euro-atlantic
- united-kingdom
- ireland
- wind
- precipitation
metrics:
- mae
- rmse
---

# 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](https://github.com/AndriiShchur/MR-GNF).

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](https://arxiv.org/abs/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.

```bash
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

```python
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:

```text
graph_2024-01.pt
graph_2024-02.pt
...
```

A batch is returned as:

```python
(x_all, geo, pos2d), y_all
```

with the principal tensors arranged as:

```text
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

```python
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:

```text
net.edge_src
net.edge_dst
```

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

## Run one-step inference

```python
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.

```python
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.

```python
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:

```python
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:

```bibtex
@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](https://arxiv.org/abs/2603.13563)

Full project repository: [https://github.com/AndriiShchur/MR-GNF](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.