Upload folder using huggingface_hub

README.md

@@ -0,0 +1,159 @@
+---
+datasets:
+- OneScience/ERA5
+frameworks:
+- PyTorch
+language:
+- en
+- zh
+license: mit
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Global Forecast
+- ERA5
+- OneForecast
+tasks: []
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">OneForecast</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+OneForecast was developed by the team of Prof. Xiaomeng Huang in the Department of Earth System Science at Tsinghua University, in collaboration with multiple institutions. The paper has been accepted by ICML 2025.
+
+Paper: OneForecast: A Universal Framework for Global and Regional Weather Forecasting
+
+https://arxiv.org/abs/2502.00338
+
+# Model Description
+
+OneForecast is a universal framework for nested global-regional weather forecasting based on graph neural networks (GNNs). Its core goal is to address the challenges of existing AI weather models in balancing low-resolution global forecasts with high-resolution regional forecasts, as well as issues such as over-smoothing in extreme event forecasting.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global Weather Forecast Training | Train the single-step OneForecast model using ERA5 HDF5 data. |
+| Local Quick Validation | Use synthetic data to verify the data protocol, model construction, training, inference, and result visualization. |
+| Multi-GPU Training | Data-parallel training on multiple GPUs/DCUs via PyTorch DDP and `torchrun`. |
+| ModelScope / OneCode Execution | Download as a standalone model package, install the OneScience dependencies, and run. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/OneForecast --local-dir ./OneForecast
+cd OneForecast
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data Introduction
+
+The OneScience community provides ERA5 data for training (due to file size limits, the current repository contains a slice of the full dataset). Users can download it with the command below and confirm that the data path in `conf/config.yaml` is set correctly:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+### Generate Synthetic Data
+
+Synthetic data is only used to verify the data protocol and program flow; it does not represent scientific forecast quality:
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Training starts from random initialization by default and saves the model to `data/checkpoint/model_bak.tar`.
+
+Multi-GPU:
+
+```bash
+torchrun --nproc-per-node=4 scripts/train.py
+```
+
+### Fine-tuning
+
+Fine-tuning starts from the training checkpoint `data/checkpoint/model_bak.tar` by default and saves the result to `data/checkpoint/model_finetuned.tar`:
+
+```bash
+python scripts/finetune.py
+```
+
+### Training Weights
+
+This repository provides weights trained on ERA5 reanalysis data in the `weights/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference loads the training checkpoint `data/checkpoint/model_bak.tar` by default, uses the test-year data, and writes predictions to `outputs/predictions/`:
+
+```bash
+python scripts/inference.py --config conf/config.yaml
+```
+
+### Result Visualization
+
+```bash
+python scripts/result.py
+```
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation & License
+
+- This repository is a reproduction of the original OneForecast paper.

conf/config.yaml

@@ -0,0 +1,141 @@
+model:
+  name: OneForecast
+  input_channels: 69
+  output_channels: 69
+  # Native ERA5 721x1440 is sampled every sixth point to 121x240, then cropped to 120x240.
+  grid_height: 120
+  grid_width: 240
+  dt_hours: 6
+  mesh_level: 5
+  processor_layers: 16
+  hidden_layers: 1
+  hidden_dim: 512
+  num_heads_edge: 4
+  num_heads_node: 4
+  weight_init: scratch
+  official_checkpoint_path: /root/private_data/workspaces/yangzt01/OneForecast/best_ckpt.tar
+
+runtime:
+  seed: 42
+  device: dcu
+  output_dir: ./outputs
+  distributed_backend: nccl
+
+training:
+  start_epoch: 0
+  max_epoch: 5
+  learning_rate: 0.00025
+  weight_decay: 0.0
+  checkpoint_dir: ./data/checkpoint
+  model_name: model_bak
+  save_every_epoch: 1
+  max_batches: 1
+
+finetuning:
+  steps: 2
+  max_epoch: 2
+  learning_rate: 0.00025
+  max_batches: 1
+  model_source: trained
+  trained_model_path: ./data/checkpoint/model_bak.tar
+  official_checkpoint_path: /root/private_data/workspaces/yangzt01/OneForecast/best_ckpt.tar
+  output_path: ./data/checkpoint/model_finetuned.tar
+
+datapipe:
+  name: ERA5
+  task: global_weather_forecasting
+  dataset_dir: ./data
+  train_years: [2000]
+  valid_years: [2001]
+  test_years: [2002]
+  input_steps: 1
+  output_steps: 1
+  normalize: true
+  batch_size: 1
+  num_workers: 0
+  # Runtime distribution is enabled automatically when launched with torchrun.
+  distributed: false
+
+# This exact order matches the official OneForecast global dataset.
+  variables:
+    - Z50
+    - Z100
+    - Z150
+    - Z200
+    - Z250
+    - Z300
+    - Z400
+    - Z500
+    - Z600
+    - Z700
+    - Z850
+    - Z925
+    - Z1000
+    - Q50
+    - Q100
+    - Q150
+    - Q200
+    - Q250
+    - Q300
+    - Q400
+    - Q500
+    - Q600
+    - Q700
+    - Q850
+    - Q925
+    - Q1000
+    - T50
+    - T100
+    - T150
+    - T200
+    - T250
+    - T300
+    - T400
+    - T500
+    - T600
+    - T700
+    - T850
+    - T925
+    - T1000
+    - U50
+    - U100
+    - U150
+    - U200
+    - U250
+    - U300
+    - U400
+    - U500
+    - U600
+    - U700
+    - U850
+    - U925
+    - U1000
+    - V50
+    - V100
+    - V150
+    - V200
+    - V250
+    - V300
+    - V400
+    - V500
+    - V600
+    - V700
+    - V850
+    - V925
+    - V1000
+    - U10M
+    - V10M
+    - T2M
+    - MSLP
+
+inference:
+  model_source: trained
+  trained_model_path: ./data/checkpoint/model_bak.tar
+  official_checkpoint_path: /root/private_data/workspaces/yangzt01/OneForecast/best_ckpt.tar
+  output_dir: ./outputs/predictions
+  max_batches: 1
+
+visualization:
+  input_dir: ./outputs/predictions
+  output_dir: ./outputs/visualizations
+  channels: [0, 18, 36, 54, 67, 68]

config.json

@@ -0,0 +1,97 @@
+{
+  "model_name": "OneForecast",
+  "model_type": "oneforecast",
+  "architectures": [
+    "OneForecast"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "global-regional-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/oneforecast.py",
+    "scope": "portable encoder-processor-decoder mesh graph neural network with PyTorch index-tensor graph kernels usable on CPU, CUDA, and DCU; official checkpoint compatible"
+  },
+  "architecture": {
+    "family": "encoder-processor-decoder graph neural network on a refined icosahedral mesh",
+    "input_grid_shape": [
+      120,
+      240
+    ],
+    "grid_sampling": "ERA5 0.25-degree 721x1440 sampled every sixth point, then cropped to 120x240",
+    "input_channels": 69,
+    "output_channels": 69,
+    "mesh_level": 5,
+    "mesh_local_refinement_regions": [
+      {
+        "description": "South and East Asia",
+        "lat_range": [
+          0.0,
+          30.0
+        ],
+        "lon_range": [
+          105.0,
+          160.0
+        ]
+      },
+      {
+        "description": "Central North America",
+        "lat_range": [
+          10.0,
+          30.0
+        ],
+        "lon_range": [
+          -95.0,
+          -35.0
+        ]
+      }
+    ],
+    "hidden_dim": 512,
+    "processor_layers": 16,
+    "hidden_layers": 1,
+    "num_heads_edge": 4,
+    "num_heads_node": 4,
+    "aggregation": "sum",
+    "activation": "SiLU",
+    "normalization": "LayerNorm"
+  },
+  "data": {
+    "dataset": "ERA5",
+    "grid_spatial_resolution_degrees": 1.5,
+    "time_step_hours": 6,
+    "input_steps": 1,
+    "output_steps": 1,
+    "surface_variables": [
+      "10m_u_component_of_wind",
+      "10m_v_component_of_wind",
+      "2m_temperature",
+      "mean_sea_level_pressure"
+    ],
+    "atmospheric_variables": [
+      "geopotential",
+      "specific_humidity",
+      "temperature",
+      "u_component_of_wind",
+      "v_component_of_wind"
+    ],
+    "pressure_levels_hpa": [
+      50,
+      100,
+      150,
+      200,
+      250,
+      300,
+      400,
+      500,
+      600,
+      700,
+      850,
+      925,
+      1000
+    ]
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/oneforecast.py",
+    "model/era5_adapter.py"
+  ]
+}

configuration.json

@@ -0,0 +1 @@
+{"framework":"PyTorch","task":"global_weather_forecasting","hardware":"DCU"}

model/era5_adapter.py

@@ -0,0 +1,254 @@
+"""OneScience ERA5 adapter for the official OneForecast 69-channel contract."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import tempfile
+from typing import Any, Iterable
+
+import numpy as np
+
+SOURCE_GRID = (721, 1440)
+ONEFORECAST_FILE_GRID = (121, 240)
+SPATIAL_STRIDE = 6
+
+OFFICIAL_VARIABLES = tuple(
+    [f"Z{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"Q{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"T{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"U{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"V{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + ["U10M", "V10M", "T2M", "MSLP"]
+)
+
+VARIABLE_ALIASES = {
+    **{f"Z{x}": f"geopotential_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"Q{x}": f"specific_humidity_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"T{x}": f"temperature_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"U{x}": f"u_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"V{x}": f"v_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    "U10M": "10m_u_component_of_wind",
+    "V10M": "10m_v_component_of_wind",
+    "T2M": "2m_temperature",
+    "MSLP": "mean_sea_level_pressure",
+}
+
+
+def _decode_variables(values: Iterable[Any]) -> list[str]:
+    return [value.decode() if isinstance(value, bytes) else str(value) for value in values]
+
+
+class OneForecastERA5Adapter:
+    """Validate files and construct OneScience's ERA5 DataLoader."""
+
+    def __init__(self, dataset_dir: str | Path, years: Iterable[int], batch_size: int = 1,
+                 input_steps: int = 1, output_steps: int = 1, normalize: bool = True,
+                 num_workers: int = 0, distributed: bool = False) -> None:
+        self.dataset_dir = Path(dataset_dir).expanduser().resolve()
+        self.years = [int(year) for year in years]
+        self.batch_size = batch_size
+        self.input_steps = input_steps
+        self.output_steps = output_steps
+        self.normalize = normalize
+        self.num_workers = num_workers
+        self.distributed = distributed
+        self.source_variables: list[str] = []
+        self.channel_indices: list[int] = []
+        self.global_means: np.ndarray | None = None
+        self.global_stds: np.ndarray | None = None
+        self.time_step_hours: int | None = None
+        self.source_grid: tuple[int, int] | None = None
+        self._external_stats: tuple[Path, Path] | None = None
+        self._layout_dir: tempfile.TemporaryDirectory[str] | None = None
+        self._validate_files()
+
+    def _year_path(self, year: int) -> Path:
+        for path in (self.dataset_dir / "data" / f"{year}.h5", self.dataset_dir / f"{year}.h5"):
+            if path.is_file():
+                return path
+        raise FileNotFoundError(f"ERA5 file for year {year} was not found below {self.dataset_dir}")
+
+    def _validate_files(self) -> None:
+        try:
+            import h5py
+        except ImportError as exc:
+            raise RuntimeError("h5py is required to validate ERA5 HDF5 files") from exc
+        if not self.years:
+            raise ValueError("At least one ERA5 year is required")
+        reference_variables: list[str] | None = None
+        reference_indices: list[int] | None = None
+        for year in self.years:
+            path = self._year_path(year)
+            with h5py.File(path, "r") as handle:
+                if "fields" not in handle:
+                    raise ValueError(f"{path} does not contain a fields dataset")
+                fields = handle["fields"]
+                if len(fields.shape) != 4:
+                    raise ValueError(f"{path}: fields must have shape [T, C, H, W], got {fields.shape}")
+                variables = _decode_variables(fields.attrs.get("variables", []))
+                source_variables = [
+                    name if name in variables else VARIABLE_ALIASES[name]
+                    for name in OFFICIAL_VARIABLES
+                    if name in variables or VARIABLE_ALIASES[name] in variables
+                ]
+                missing = [
+                    name for name in OFFICIAL_VARIABLES
+                    if name not in variables and VARIABLE_ALIASES[name] not in variables
+                ]
+                if missing:
+                    raise ValueError(f"{path}: missing official variables: {missing}")
+                indices = [variables.index(name) for name in source_variables]
+                if reference_variables is not None and variables != reference_variables:
+                    raise ValueError(f"{path}: complete variable metadata differs between yearly files")
+                if reference_indices is not None and indices != reference_indices:
+                    raise ValueError(f"{path}: official channel indices differ between yearly files")
+                reference_variables, reference_indices = variables, indices
+                self.source_variables = source_variables
+                self.channel_indices = indices
+                if fields.shape[1] != len(variables):
+                    raise ValueError(f"{path}: variables metadata does not match channel dimension")
+                if fields.shape[1] != 69 or tuple(fields.shape[2:]) not in (SOURCE_GRID, ONEFORECAST_FILE_GRID):
+                    raise ValueError(
+                        f"{path}: expected fields [T, 69, 721, 1440] or [T, 69, 121, 240], got {fields.shape}"
+                    )
+                grid = tuple(fields.shape[2:])
+                if self.source_grid is not None and grid != self.source_grid:
+                    raise ValueError(f"{path}: spatial grid differs between yearly files")
+                self.source_grid = grid
+                if fields.shape[0] < self.input_steps + self.output_steps:
+                    raise ValueError(f"{path}: not enough time steps for configured window")
+                if "time_step" not in fields.attrs:
+                    raise ValueError(f"{path}: fields.attrs['time_step'] is required by ERA5Datapipe")
+                time_step = int(fields.attrs["time_step"])
+                if time_step != 6 or (self.time_step_hours is not None and time_step != self.time_step_hours):
+                    raise ValueError(f"{path}: expected a consistent 6-hour time_step, got {time_step}")
+                self.time_step_hours = time_step
+                if "global_means" in handle and "global_stds" in handle:
+                    means = np.asarray(handle["global_means"])
+                    stds = np.asarray(handle["global_stds"])
+                else:
+                    candidates = (
+                        (self.dataset_dir / "stats" / "global_means.npy",
+                         self.dataset_dir / "stats" / "global_stds.npy"),
+                        (self.dataset_dir / "mean.npy", self.dataset_dir / "std.npy"),
+                        (self.dataset_dir.parent / "mean.npy", self.dataset_dir.parent / "std.npy"),
+                    )
+                    stats_paths = next(((mean, std) for mean, std in candidates
+                                        if mean.is_file() and std.is_file()), None)
+                    if stats_paths is None:
+                        raise ValueError(f"{path}: embedded or external ERA5 statistics are required")
+                    self._external_stats = stats_paths
+                    means, stds = (np.load(item) for item in stats_paths)
+                expected_shape = (1, len(variables), 1, 1)
+                if means.shape != expected_shape or stds.shape != expected_shape:
+                    raise ValueError(f"{path}: statistics must have shape {expected_shape}")
+                if not np.isfinite(means).all() or not np.isfinite(stds).all() or not (stds > 0).all():
+                    raise ValueError(f"{path}: statistics must be finite and standard deviations positive")
+                if self.global_means is not None and not np.array_equal(means, self.global_means):
+                    raise ValueError(f"{path}: global_means differ between yearly files")
+                if self.global_stds is not None and not np.array_equal(stds, self.global_stds):
+                    raise ValueError(f"{path}: global_stds differ between yearly files")
+                self.global_means, self.global_stds = means, stds
+
+    def _onescience_dataset_dir(self) -> Path:
+        if self._layout_dir is not None:
+            return Path(self._layout_dir.name)
+        self._layout_dir = tempfile.TemporaryDirectory(prefix="oneforecast_era5_")
+        root = Path(self._layout_dir.name)
+        data_dir = root / "data"
+        data_dir.mkdir()
+        for year in self.years:
+            source_path = self._year_path(year)
+            target_path = data_dir / f"{year}.h5"
+            if self.source_grid == SOURCE_GRID:
+                import h5py
+
+                with h5py.File(source_path, "r") as source_handle:
+                    source_fields = source_handle["fields"]
+                    layout = h5py.VirtualLayout(
+                        shape=(source_fields.shape[0], source_fields.shape[1], *ONEFORECAST_FILE_GRID),
+                        dtype=source_fields.dtype,
+                    )
+                    virtual_source = h5py.VirtualSource(str(source_path), "fields", shape=source_fields.shape)
+                    layout[:] = virtual_source[:, :, ::SPATIAL_STRIDE, ::SPATIAL_STRIDE]
+                    with h5py.File(target_path, "w", libver="latest") as target_handle:
+                        fields = target_handle.create_virtual_dataset("fields", layout)
+                        for name, value in source_fields.attrs.items():
+                            fields.attrs[name] = value
+            else:
+                target_path.symlink_to(source_path)
+        if self._external_stats is not None:
+            stats_dir = root / "stats"
+            stats_dir.mkdir()
+            (stats_dir / "global_means.npy").symlink_to(self._external_stats[0])
+            (stats_dir / "global_stds.npy").symlink_to(self._external_stats[1])
+
+        return root
+
+    def get_dataloader(self, mode: str):
+        """Delegate loading to OneScience, then align native ERA5 to OneForecast's grid."""
+        try:
+            from onescience.datapipes.climate.era5 import ERA5Datapipe
+        except ImportError as exc:
+            raise RuntimeError("OneScience ERA5Datapipe is required for data loading") from exc
+        datapipe = ERA5Datapipe(
+            dataset_dir=str(self._onescience_dataset_dir()), used_years=self.years,
+            used_variables=self.source_variables, distributed=self.distributed,
+            input_steps=self.input_steps, output_steps=self.output_steps,
+            normalize=self.normalize, batch_size=self.batch_size, num_workers=self.num_workers,
+        )
+        loader, sampler = datapipe.get_dataloader(mode=mode)
+        return _SpatiallyAdaptedLoader(loader, self.source_grid), sampler
+
+    def inspect(self) -> dict[str, Any]:
+        try:
+            import h5py
+        except ImportError as exc:
+            raise RuntimeError("h5py is required to inspect ERA5 HDF5 files") from exc
+        path = self._year_path(self.years[0])
+        with h5py.File(path, "r") as handle:
+            fields = handle["fields"]
+            variables = _decode_variables(fields.attrs["variables"])
+            indices = [variables.index(name) for name in self.source_variables]
+            return {"path": str(path), "fields_shape": list(fields.shape),
+                     "source_grid": list(fields.shape[2:]),
+                     "oneforecast_file_grid": list(ONEFORECAST_FILE_GRID),
+                     "oneforecast_model_grid": [120, 240],
+                     "spatial_transform": "identity" if tuple(fields.shape[2:]) == ONEFORECAST_FILE_GRID else "stride_6",
+                     "time_step_hours": int(fields.attrs["time_step"]),
+                     "variable_count": len(variables), "official_channel_indices": indices,
+                     "source_variables": self.source_variables,
+                     "statistics_shape": list(self.global_means.shape),
+                     "statistics_shared_across_years": True,
+                     "official_variables_match": len(indices) == len(OFFICIAL_VARIABLES)}
+
+    def selected_statistics(self) -> tuple[np.ndarray, np.ndarray]:
+        """Return normalization statistics in the model's 69-channel order."""
+        if self.global_means is None or self.global_stds is None:
+            raise RuntimeError("ERA5 statistics have not been validated")
+        return self.global_means[:, self.channel_indices], self.global_stds[:, self.channel_indices]
+
+
+def _adapt_spatial(value: Any, source_grid: tuple[int, int] | None) -> Any:
+    if not hasattr(value, "shape") or len(value.shape) < 2:
+        return value
+    if tuple(value.shape[-2:]) == ONEFORECAST_FILE_GRID:
+        return value
+    if tuple(value.shape[-2:]) != SOURCE_GRID or source_grid != SOURCE_GRID:
+        return value
+    return value[..., ::SPATIAL_STRIDE, ::SPATIAL_STRIDE]
+
+
+class _SpatiallyAdaptedLoader:
+    """Preserve the DataLoader interface while adapting fields after ERA5Datapipe."""
+
+    def __init__(self, loader: Any, source_grid: tuple[int, int] | None) -> None:
+        self.loader = loader
+        self.source_grid = source_grid
+
+    def __len__(self) -> int:
+        return len(self.loader)
+
+    def __iter__(self):
+        for batch in self.loader:
+            yield tuple(_adapt_spatial(value, self.source_grid) for value in batch)

model/oneforecast.py

@@ -0,0 +1,553 @@
+"""Portable OneForecast model and official checkpoint compatibility helpers.
+
+The parameter hierarchy mirrors the official model. Graph operations use
+PyTorch index tensors instead of CUDA-only CuGraph kernels, making the model
+usable on CPU, CUDA, and DCU PyTorch builds.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+import sys
+import types
+from typing import Any, NamedTuple
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+class GraphData:
+    """Minimal homogeneous or bipartite graph used by the portable kernels."""
+
+    def __init__(self, src: Tensor, dst: Tensor, num_src: int, num_dst: int) -> None:
+        self.src = src.to(torch.long)
+        self.dst = dst.to(torch.long)
+        self.num_src = num_src
+        self.num_dst = num_dst
+
+    def to(self, device: torch.device | str) -> "GraphData":
+        self.src = self.src.to(device)
+        self.dst = self.dst.to(device)
+        return self
+
+
+def _aggregate(values: Tensor, dst: Tensor, num_dst: int, reduction: str) -> Tensor:
+    output = values.new_zeros((num_dst,) + values.shape[1:])
+    index = dst.view((-1,) + (1,) * (values.ndim - 1)).expand_as(values)
+    output.scatter_add_(0, index, values)
+    if reduction == "mean":
+        counts = values.new_zeros(num_dst)
+        counts.scatter_add_(0, dst, torch.ones_like(dst, dtype=values.dtype))
+        output = output / counts.clamp_min(1).view((-1,) + (1,) * (values.ndim - 1))
+    elif reduction != "sum":
+        raise ValueError(f"Unsupported aggregation: {reduction}")
+    return output
+
+
+def _edge_softmax(logits: Tensor, dst: Tensor, num_dst: int) -> Tensor:
+    index = dst[:, None].expand_as(logits)
+    maxima = logits.new_full((num_dst, logits.shape[1]), -torch.inf)
+    maxima.scatter_reduce_(0, index, logits, reduce="amax", include_self=True)
+    exp = torch.exp(logits - maxima[dst])
+    denominator = logits.new_zeros((num_dst, logits.shape[1]))
+    denominator.scatter_add_(0, index, exp)
+    return exp / denominator[dst].clamp_min(torch.finfo(exp.dtype).tiny)
+
+
+class MeshGraphMLP(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int = 512, hidden_dim: int = 512,
+                 hidden_layers: int | None = 1, activation_fn: nn.Module | None = None,
+                 norm_type: str | None = "LayerNorm", recompute_activation: bool = False) -> None:
+        super().__init__()
+        del recompute_activation
+        activation_fn = activation_fn or nn.SiLU()
+        if hidden_layers is None:
+            self.model = nn.Identity()
+            return
+        layers: list[nn.Module] = [nn.Linear(input_dim, hidden_dim), activation_fn]
+        for _ in range(hidden_layers - 1):
+            layers.extend([nn.Linear(hidden_dim, hidden_dim), nn.SiLU()])
+        layers.append(nn.Linear(hidden_dim, output_dim))
+        if norm_type is not None:
+            if norm_type != "LayerNorm":
+                raise ValueError("The portable model supports LayerNorm only")
+            layers.append(nn.LayerNorm(output_dim))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.model(x)
+
+
+class MeshGraphEdgeMLPSum(nn.Module):
+    """Concat-trick edge MLP with the official parameter names and initialization."""
+
+    def __init__(self, efeat_dim: int, src_dim: int, dst_dim: int,
+                 output_dim: int = 512, hidden_dim: int = 512,
+                 hidden_layers: int = 1, activation_fn: nn.Module | None = None,
+                 norm_type: str | None = "LayerNorm", recompute_activation: bool = False) -> None:
+        super().__init__()
+        del recompute_activation
+        activation_fn = activation_fn or nn.SiLU()
+        initial = nn.Linear(efeat_dim + src_dim + dst_dim, hidden_dim)
+        weights = torch.split(initial.weight, [efeat_dim, src_dim, dst_dim], dim=1)
+        self.lin_efeat = nn.Parameter(weights[0])
+        self.lin_src = nn.Parameter(weights[1])
+        self.lin_dst = nn.Parameter(weights[2])
+        self.bias = initial.bias
+        layers: list[nn.Module] = [activation_fn]
+        for _ in range(hidden_layers - 1):
+            layers.extend([nn.Linear(hidden_dim, hidden_dim), nn.SiLU()])
+        layers.append(nn.Linear(hidden_dim, output_dim))
+        if norm_type is not None:
+            if norm_type != "LayerNorm":
+                raise ValueError("The portable model supports LayerNorm only")
+            layers.append(nn.LayerNorm(output_dim))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, efeat: Tensor, nfeat: Tensor | tuple[Tensor, Tensor],
+                graph: GraphData) -> Tensor:
+        src_feat, dst_feat = (nfeat, nfeat) if isinstance(nfeat, Tensor) else nfeat
+        hidden = F.linear(efeat, self.lin_efeat)
+        hidden = hidden + F.linear(src_feat[graph.src], self.lin_src)
+        hidden = hidden + F.linear(dst_feat[graph.dst], self.lin_dst, self.bias)
+        return self.model(hidden)
+
+
+class OneForecastEncoderEmbedder(nn.Module):
+    def __init__(self, input_dim_grid_nodes: int = 69, input_dim_mesh_nodes: int = 3,
+                 input_dim_edges: int = 4, output_dim: int = 512,
+                 hidden_dim: int = 512, hidden_layers: int = 1) -> None:
+        super().__init__()
+        kwargs = dict(output_dim=output_dim, hidden_dim=hidden_dim, hidden_layers=hidden_layers)
+        self.grid_node_mlp = MeshGraphMLP(input_dim_grid_nodes, **kwargs)
+        self.mesh_node_mlp = MeshGraphMLP(input_dim_mesh_nodes, **kwargs)
+        self.mesh_edge_mlp = MeshGraphMLP(input_dim_edges, **kwargs)
+        self.grid2mesh_edge_mlp = MeshGraphMLP(input_dim_edges, **kwargs)
+
+    def forward(self, grid: Tensor, mesh: Tensor, g2m: Tensor,
+                mesh_edges: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+        return (self.grid_node_mlp(grid), self.mesh_node_mlp(mesh),
+                self.grid2mesh_edge_mlp(g2m), self.mesh_edge_mlp(mesh_edges))
+
+
+class OneForecastDecoderEmbedder(nn.Module):
+    def __init__(self, input_dim_edges: int = 4, output_dim: int = 512,
+                 hidden_dim: int = 512, hidden_layers: int = 1) -> None:
+        super().__init__()
+        self.mesh2grid_edge_mlp = MeshGraphMLP(
+            input_dim_edges, output_dim, hidden_dim, hidden_layers)
+
+    def forward(self, edges: Tensor) -> Tensor:
+        return self.mesh2grid_edge_mlp(edges)
+
+
+class MeshGraphEncoder(nn.Module):
+    def __init__(self, hidden_dim: int = 512, hidden_layers: int = 1,
+                 aggregation: str = "sum") -> None:
+        super().__init__()
+        self.aggregation = aggregation
+        self.edge_mlp = MeshGraphEdgeMLPSum(hidden_dim, hidden_dim, hidden_dim,
+                                            hidden_dim, hidden_dim, hidden_layers)
+        self.src_node_mlp = MeshGraphMLP(hidden_dim, hidden_dim, hidden_dim, hidden_layers)
+        self.dst_node_mlp = MeshGraphMLP(hidden_dim * 2, hidden_dim, hidden_dim, hidden_layers)
+
+    def forward(self, edges: Tensor, grid: Tensor, mesh: Tensor,
+                graph: GraphData) -> tuple[Tensor, Tensor]:
+        edges = self.edge_mlp(edges, (grid, mesh), graph)
+        aggregated = _aggregate(edges, graph.dst, graph.num_dst, self.aggregation)
+        return grid + self.src_node_mlp(grid), mesh + self.dst_node_mlp(torch.cat((aggregated, mesh), -1))
+
+
+class MeshGraphDecoder(nn.Module):
+    def __init__(self, hidden_dim: int = 512, hidden_layers: int = 1,
+                 aggregation: str = "sum") -> None:
+        super().__init__()
+        self.aggregation = aggregation
+        self.edge_mlp = MeshGraphEdgeMLPSum(hidden_dim, hidden_dim, hidden_dim,
+                                            hidden_dim, hidden_dim, hidden_layers)
+        self.node_mlp = MeshGraphMLP(hidden_dim * 2, hidden_dim, hidden_dim, hidden_layers)
+
+    def forward(self, edges: Tensor, grid: Tensor, mesh: Tensor, graph: GraphData) -> Tensor:
+        edges = self.edge_mlp(edges, (mesh, grid), graph)
+        aggregated = _aggregate(edges, graph.dst, graph.num_dst, self.aggregation)
+        return grid + self.node_mlp(torch.cat((aggregated, grid), -1))
+
+
+class MeshEdgeBlockMultiHeadGated(nn.Module):
+    def __init__(self, hidden_dim: int = 512, hidden_layers: int = 1,
+                 num_heads: int = 4) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.edge_mlp = MeshGraphEdgeMLPSum(hidden_dim, hidden_dim, hidden_dim,
+                                            hidden_dim, hidden_dim, hidden_layers)
+        gating_hidden = max(16, hidden_dim // 8)
+        self.gate_net = nn.Sequential(nn.Linear(hidden_dim * 3, gating_hidden), nn.SiLU(),
+                                      nn.Linear(gating_hidden, 3 * num_heads), nn.Sigmoid())
+
+    def forward(self, edges: Tensor, nodes: Tensor, graph: GraphData) -> tuple[Tensor, Tensor]:
+        raw = torch.cat((edges, nodes[graph.src], nodes[graph.dst]), -1)
+        gates = self.gate_net(raw).view(-1, self.num_heads, 3).mean(1)
+        updated = self.edge_mlp(edges, nodes, graph)
+        return edges + updated * gates.mean(-1, keepdim=True), nodes
+
+
+class MeshNodeBlockMultiHeadAttn(nn.Module):
+    def __init__(self, hidden_dim: int = 512, hidden_layers: int = 1,
+                 aggregation: str = "sum", num_heads: int = 4) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.aggregation = aggregation
+        self.node_mlp = MeshGraphMLP(hidden_dim * (num_heads + 1), hidden_dim,
+                                     hidden_dim, hidden_layers)
+        attention_hidden = max(16, hidden_dim // 8)
+        self.attn_net = nn.Sequential(nn.Linear(hidden_dim, attention_hidden), nn.SiLU(),
+                                      nn.Linear(attention_hidden, num_heads))
+
+    def forward(self, edges: Tensor, nodes: Tensor, graph: GraphData) -> tuple[Tensor, Tensor]:
+        scores = _edge_softmax(self.attn_net(edges), graph.dst, graph.num_dst)
+        messages = edges[:, None, :].expand(-1, self.num_heads, -1) * scores[:, :, None]
+        aggregated = _aggregate(messages, graph.dst, graph.num_dst, self.aggregation).flatten(1)
+        return edges, nodes + self.node_mlp(torch.cat((aggregated, nodes), -1))
+
+
+class OneForecastProcessor(nn.Module):
+    def __init__(self, processor_layers: int, hidden_dim: int = 512,
+                 hidden_layers: int = 1, aggregation: str = "sum",
+                 num_heads_edge: int = 4, num_heads_node: int = 4) -> None:
+        super().__init__()
+        layers: list[nn.Module] = []
+        for _ in range(processor_layers):
+            layers.append(MeshEdgeBlockMultiHeadGated(hidden_dim, hidden_layers, num_heads_edge))
+            layers.append(MeshNodeBlockMultiHeadAttn(hidden_dim, hidden_layers, aggregation, num_heads_node))
+        self.processor_layers = nn.ModuleList(layers)
+
+    def forward(self, edges: Tensor, nodes: Tensor, graph: GraphData) -> tuple[Tensor, Tensor]:
+        for layer in self.processor_layers:
+            edges, nodes = layer(edges, nodes, graph)
+        return edges, nodes
+
+
+class TriangularMesh(NamedTuple):
+    vertices: np.ndarray
+    faces: np.ndarray
+
+
+def _icosahedron() -> TriangularMesh:
+    from scipy.spatial.transform import Rotation
+
+    phi = (1 + np.sqrt(5)) / 2
+    vertices = []
+    for first in (1.0, -1.0):
+        for second in (phi, -phi):
+            vertices.extend(((first, second, 0.0), (0.0, first, second), (second, 0.0, first)))
+    vertices = np.asarray(vertices, dtype=np.float32) / np.linalg.norm([1.0, phi])
+    faces = np.asarray(((0,1,2),(0,6,1),(8,0,2),(8,4,0),(3,8,2),(3,2,7),(7,2,1),
+                        (0,4,6),(4,11,6),(6,11,5),(1,5,7),(4,10,11),(4,8,10),(10,8,3),
+                        (10,3,9),(11,10,9),(11,9,5),(5,9,7),(9,3,7),(1,6,5)), dtype=np.int32)
+    angle = (np.pi - 2 * np.arcsin(phi / np.sqrt(3))) / 2
+    vertices = vertices @ Rotation.from_euler("y", angle).as_matrix()
+    return TriangularMesh(vertices.astype(np.float32), faces)
+
+
+def _split_mesh(mesh: TriangularMesh) -> TriangularMesh:
+    vertices = list(mesh.vertices)
+    children: dict[tuple[int, int], int] = {}
+    faces = []
+    for a, b, c in mesh.faces:
+        mids = []
+        for pair in ((a, b), (b, c), (c, a)):
+            key = tuple(sorted(map(int, pair)))
+            if key not in children:
+                position = mesh.vertices[list(pair)].mean(0)
+                position /= np.linalg.norm(position)
+                children[key] = len(vertices)
+                vertices.append(position)
+            mids.append(children[key])
+        ab, bc, ca = mids
+        faces.extend(((a, ab, ca), (ab, b, bc), (ca, bc, c), (ab, bc, ca)))
+    return TriangularMesh(np.asarray(vertices, dtype=np.float32), np.asarray(faces, dtype=np.int32))
+
+
+def _mesh_hierarchy(level: int) -> list[TriangularMesh]:
+    meshes = [_icosahedron()]
+    for _ in range(level):
+        meshes.append(_split_mesh(meshes[-1]))
+    return meshes
+
+
+def _faces_to_edges(faces: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    return (np.concatenate((faces[:, 0], faces[:, 1], faces[:, 2])),
+            np.concatenate((faces[:, 1], faces[:, 2], faces[:, 0])))
+
+
+def _latlon_to_xyz(latlon: Tensor) -> Tensor:
+    values = torch.deg2rad(latlon)
+    lat, lon = values[:, 0], values[:, 1]
+    return torch.stack((torch.cos(lat) * torch.cos(lon), torch.cos(lat) * torch.sin(lon), torch.sin(lat)), 1)
+
+
+def _node_features(xyz: Tensor) -> Tensor:
+    # The official implementation applies trigonometric functions to the
+    # degree-valued xyz2latlon output; retain that behavior for parity.
+    lat = torch.rad2deg(torch.asin(xyz[:, 2]))
+    lon = torch.rad2deg(torch.atan2(xyz[:, 1], xyz[:, 0]))
+    return torch.stack((torch.cos(lat), torch.sin(lon), torch.cos(lon)), -1)
+
+
+def _edge_features(src_pos: Tensor, dst_pos: Tensor, src: Tensor, dst: Tensor) -> Tensor:
+    source, target = src_pos[src], dst_pos[dst]
+    lat = torch.asin(target[:, 2])
+    lon = torch.atan2(target[:, 1], target[:, 0])
+    cos_lon, sin_lon = torch.cos(-lon), torch.sin(-lon)
+    source = torch.stack((cos_lon * source[:, 0] - sin_lon * source[:, 1],
+                          sin_lon * source[:, 0] + cos_lon * source[:, 1], source[:, 2]), -1)
+    target = torch.stack((cos_lon * target[:, 0] - sin_lon * target[:, 1],
+                          sin_lon * target[:, 0] + cos_lon * target[:, 1], target[:, 2]), -1)
+    cos_lat, sin_lat = torch.cos(lat), torch.sin(lat)
+    source = torch.stack((cos_lat * source[:, 0] + sin_lat * source[:, 2], source[:, 1],
+                          -sin_lat * source[:, 0] + cos_lat * source[:, 2]), -1)
+    target = torch.stack((cos_lat * target[:, 0] + sin_lat * target[:, 2], target[:, 1],
+                          -sin_lat * target[:, 0] + cos_lat * target[:, 2]), -1)
+    displacement = source - target
+    norm = torch.linalg.norm(displacement, dim=-1, keepdim=True)
+    maximum = norm.max()
+    return torch.cat((displacement / maximum, norm / maximum), -1)
+
+
+def _local_refine(mesh: TriangularMesh, lat_min: float, lat_max: float,
+                  lon_min: float, lon_max: float) -> TriangularMesh:
+    centroids = mesh.vertices[mesh.faces].mean(axis=1)
+    # Match the official xyz2latlon call, which assumes radius=1 for centroids.
+    latitudes = np.rad2deg(np.arcsin(centroids[:, 2]))
+    longitudes = np.rad2deg(np.arctan2(centroids[:, 1], centroids[:, 0]))
+    selected = ((latitudes >= lat_min) & (latitudes <= lat_max)
+                & (longitudes >= lon_min) & (longitudes <= lon_max))
+    refined = _split_mesh(TriangularMesh(mesh.vertices, mesh.faces[selected]))
+    combined_vertices = np.concatenate((refined.vertices, mesh.vertices), axis=0)
+    combined_faces = np.concatenate((refined.faces, mesh.faces[~selected] + len(refined.vertices)), axis=0)
+    rounded = np.round(combined_vertices, decimals=6)
+    unique: dict[tuple[float, float, float], int] = {}
+    remap = np.empty(len(rounded), dtype=np.int64)
+    vertices = []
+    for index, coordinates in enumerate(rounded):
+        key = tuple(coordinates.tolist())
+        if key not in unique:
+            unique[key] = len(vertices)
+            vertices.append(combined_vertices[index])
+        remap[index] = unique[key]
+    return TriangularMesh(np.asarray(vertices, dtype=np.float32), remap[combined_faces].astype(np.int32))
+
+
+def _build_graphs(height: int, width: int, mesh_level: int) -> tuple[GraphData, GraphData, GraphData, Tensor, Tensor, Tensor, Tensor]:
+    from scipy.spatial import cKDTree
+
+    latitudes = torch.linspace(-90, 90, height + 1)[:-1]
+    longitudes = torch.linspace(-180, 180, width + 1)[1:]
+    latlon = torch.stack(torch.meshgrid(latitudes, longitudes, indexing="ij"), -1).reshape(-1, 2)
+    grid_xyz = _latlon_to_xyz(latlon)
+    hierarchy = _mesh_hierarchy(mesh_level)
+    finest = hierarchy[-1]
+    refined = _local_refine(finest, 0.0, 30.0, 105.0, 160.0)
+    refined = _local_refine(refined, 10.0, 30.0, -95.0, -35.0)
+    mesh_vertices = refined.vertices
+    mesh_faces = np.concatenate([mesh.faces for mesh in hierarchy] + [refined.faces], axis=0)
+    mesh_src, mesh_dst = _faces_to_edges(mesh_faces)
+    mesh_src = np.concatenate((mesh_src, mesh_dst))
+    mesh_dst = np.concatenate((mesh_dst, mesh_src[:len(mesh_dst)]))
+    pairs = np.unique(np.stack((mesh_src, mesh_dst), 1), axis=0)
+    mesh_src_t = torch.from_numpy(pairs[:, 0])
+    mesh_dst_t = torch.from_numpy(pairs[:, 1])
+    mesh_xyz = torch.from_numpy(mesh_vertices)
+    mesh_graph = GraphData(mesh_src_t, mesh_dst_t, len(mesh_vertices), len(mesh_vertices))
+
+    finest_src, finest_dst = _faces_to_edges(finest.faces)
+    max_edge = np.linalg.norm(finest.vertices[finest_src] - finest.vertices[finest_dst], axis=1).max()
+    distances, neighbors = cKDTree(mesh_vertices).query(grid_xyz.numpy(), k=4)
+    valid = distances <= 0.6 * max_edge
+    g2m_src, neighbor_slot = np.nonzero(valid)
+    g2m_dst = neighbors[g2m_src, neighbor_slot]
+    g2m_graph = GraphData(torch.from_numpy(g2m_src), torch.from_numpy(g2m_dst), len(grid_xyz), len(mesh_vertices))
+
+    centroids = mesh_vertices[mesh_faces].mean(axis=1)
+    face_indices = cKDTree(centroids).query(grid_xyz.numpy(), k=1)[1]
+    m2g_src = mesh_faces[face_indices].reshape(-1)
+    m2g_dst = np.repeat(np.arange(len(grid_xyz)), 3)
+    m2g_graph = GraphData(torch.from_numpy(m2g_src), torch.from_numpy(m2g_dst), len(mesh_vertices), len(grid_xyz))
+    mesh_nodes = _node_features(mesh_xyz)
+    mesh_edges = _edge_features(mesh_xyz, mesh_xyz, mesh_graph.src, mesh_graph.dst)
+    g2m_edges = _edge_features(grid_xyz, mesh_xyz, g2m_graph.src, g2m_graph.dst)
+    m2g_edges = _edge_features(mesh_xyz, grid_xyz, m2g_graph.src, m2g_graph.dst)
+    return mesh_graph, g2m_graph, m2g_graph, mesh_nodes, mesh_edges, g2m_edges, m2g_edges
+
+
+class OneForecast(nn.Module):
+    """Official OneForecast message-passing architecture with portable graph kernels."""
+
+    def __init__(self, input_res: tuple[int, int] = (120, 240), input_dim_grid_nodes: int = 69,
+                 output_dim_grid_nodes: int = 69, mesh_level: int = 5,
+                 processor_layers: int = 16, hidden_layers: int = 1,
+                 hidden_dim: int = 512, aggregation: str = "sum",
+                 num_heads_edge: int = 4, num_heads_node: int = 4,
+                 build_graph: bool = True) -> None:
+        super().__init__()
+        if processor_layers <= 2:
+            raise ValueError("Expected at least 3 processor layers")
+        self.register_buffer("device_buffer", torch.empty(0))
+        self.input_res = tuple(input_res)
+        self.input_dim_grid_nodes = input_dim_grid_nodes
+        self.output_dim_grid_nodes = output_dim_grid_nodes
+        self.mesh_level = mesh_level
+        self.encoder_embedder = OneForecastEncoderEmbedder(
+            input_dim_grid_nodes, 3, 4, hidden_dim, hidden_dim, hidden_layers)
+        self.decoder_embedder = OneForecastDecoderEmbedder(4, hidden_dim, hidden_dim, hidden_layers)
+        self.encoder = MeshGraphEncoder(hidden_dim, hidden_layers, aggregation)
+        self.processor_encoder = OneForecastProcessor(
+            1, hidden_dim, hidden_layers, aggregation, num_heads_edge, num_heads_node)
+        self.processor = OneForecastProcessor(
+            processor_layers - 2, hidden_dim, hidden_layers, aggregation, num_heads_edge, num_heads_node)
+        self.processor_decoder = OneForecastProcessor(
+            1, hidden_dim, hidden_layers, aggregation, num_heads_edge, num_heads_node)
+        self.decoder = MeshGraphDecoder(hidden_dim, hidden_layers, aggregation)
+        self.finale = MeshGraphMLP(hidden_dim, output_dim_grid_nodes, hidden_dim, hidden_layers, norm_type=None)
+        self._graph_ready = False
+        if build_graph:
+            self.build_graph()
+
+    def build_graph(self) -> None:
+        values = _build_graphs(*self.input_res, self.mesh_level)
+        self.mesh_graph, self.g2m_graph, self.m2g_graph = values[:3]
+        for name, value in zip(("mesh_ndata", "mesh_edata", "g2m_edata", "m2g_edata"), values[3:]):
+            self.register_buffer(name, value, persistent=False)
+        self._graph_ready = True
+
+    def forward(self, grid_nfeat: Tensor) -> Tensor:
+        if not self._graph_ready:
+            raise RuntimeError("Graph construction was disabled for this model instance")
+        if grid_nfeat.shape != (1, self.input_dim_grid_nodes, *self.input_res):
+            raise ValueError(f"Expected input shape (1, {self.input_dim_grid_nodes}, {self.input_res[0]}, {self.input_res[1]}), got {tuple(grid_nfeat.shape)}")
+        grid = grid_nfeat[0].reshape(self.input_dim_grid_nodes, -1).T
+        grid, mesh, g2m, mesh_edges = self.encoder_embedder(
+            grid, self.mesh_ndata, self.g2m_edata, self.mesh_edata)
+        grid, mesh = self.encoder(g2m, grid, mesh, self.g2m_graph)
+        mesh_edges, mesh = self.processor_encoder(mesh_edges, mesh, self.mesh_graph)
+        mesh_edges, mesh = self.processor(mesh_edges, mesh, self.mesh_graph)
+        _, mesh = self.processor_decoder(mesh_edges, mesh, self.mesh_graph)
+        grid = self.decoder(self.decoder_embedder(self.m2g_edata), grid, mesh, self.m2g_graph)
+        output = self.finale(grid).T.reshape(self.output_dim_grid_nodes, *self.input_res)
+        return output.unsqueeze(0)
+
+    def to(self, *args: Any, **kwargs: Any) -> "OneForecast":
+        super().to(*args, **kwargs)
+        if self._graph_ready:
+            device = self.device_buffer.device
+            self.mesh_graph.to(device)
+            self.g2m_graph.to(device)
+            self.m2g_graph.to(device)
+        return self
+
+
+@dataclass(frozen=True)
+class CheckpointReport:
+    checkpoint_path: str
+    checkpoint_keys: int
+    model_keys: int
+    missing_keys: tuple[str, ...]
+    unexpected_keys: tuple[str, ...]
+    shape_mismatches: tuple[str, ...]
+
+    @property
+    def compatible(self) -> bool:
+        return not (self.missing_keys or self.unexpected_keys or self.shape_mismatches)
+
+
+def _install_scalarfloat_safe_global() -> type[float]:
+    """Allow the known ruamel ScalarFloat metadata type without importing ruamel."""
+    module_name = "ruamel.yaml.scalarfloat"
+    module = sys.modules.get(module_name)
+    if module is not None and hasattr(module, "ScalarFloat"):
+        scalar_float = module.ScalarFloat
+    else:
+        ruamel = sys.modules.setdefault("ruamel", types.ModuleType("ruamel"))
+        yaml_module = sys.modules.setdefault("ruamel.yaml", types.ModuleType("ruamel.yaml"))
+        module = types.ModuleType(module_name)
+        scalar_float = type("ScalarFloat", (float,), {})
+        scalar_float.__module__ = module_name
+        module.ScalarFloat = scalar_float
+        yaml_module.scalarfloat = module
+        ruamel.yaml = yaml_module
+        sys.modules[module_name] = module
+    torch.serialization.add_safe_globals([scalar_float])
+    return scalar_float
+
+
+def read_official_checkpoint(path: str | Path) -> tuple[dict[str, Tensor], dict[str, Any]]:
+    path = Path(path).expanduser().resolve()
+    _install_scalarfloat_safe_global()
+    checkpoint = torch.load(path, map_location="cpu", weights_only=True, mmap=True)
+    if not isinstance(checkpoint, dict) or "model_state" not in checkpoint:
+        raise ValueError(f"{path} does not contain an official model_state")
+    state = checkpoint["model_state"]
+    if not isinstance(state, dict):
+        raise TypeError("checkpoint model_state must be a mapping")
+    cleaned = {key.removeprefix("module."): value for key, value in state.items()}
+    metadata = {key: value for key, value in checkpoint.items() if key not in {"model_state", "optimizer_state_dict"}}
+    return cleaned, metadata
+
+
+def _compare_checkpoint_state(model: nn.Module, checkpoint_state: dict[str, Tensor],
+                              path: str | Path) -> CheckpointReport:
+    model_state = model.state_dict()
+    missing = tuple(sorted(set(model_state) - set(checkpoint_state)))
+    unexpected = tuple(sorted(set(checkpoint_state) - set(model_state)))
+    mismatches = tuple(sorted(
+        f"{key}: checkpoint={tuple(checkpoint_state[key].shape)} model={tuple(model_state[key].shape)}"
+        for key in set(model_state) & set(checkpoint_state)
+        if model_state[key].shape != checkpoint_state[key].shape
+    ))
+    return CheckpointReport(str(Path(path).expanduser().resolve()), len(checkpoint_state),
+                            len(model_state), missing, unexpected, mismatches)
+
+
+def check_checkpoint_compatibility(model: nn.Module, path: str | Path) -> CheckpointReport:
+    checkpoint_state, _ = read_official_checkpoint(path)
+    return _compare_checkpoint_state(model, checkpoint_state, path)
+
+
+def load_official_checkpoint(model: nn.Module, path: str | Path, strict: bool = True) -> CheckpointReport:
+    state, _ = read_official_checkpoint(path)
+    report = _compare_checkpoint_state(model, state, path)
+    if strict and not report.compatible:
+        raise RuntimeError(f"Official checkpoint is incompatible: {report}")
+    compatible = {key: value for key, value in state.items()
+                  if key in model.state_dict() and value.shape == model.state_dict()[key].shape}
+    model.load_state_dict(compatible, strict=strict)
+    return report
+
+
+def build_model(config: dict[str, Any], build_graph: bool = True) -> OneForecast:
+    settings = config["model"]
+    model = OneForecast(
+        input_res=(settings["grid_height"], settings["grid_width"]),
+        input_dim_grid_nodes=settings["input_channels"],
+        output_dim_grid_nodes=settings["output_channels"],
+        mesh_level=settings.get("mesh_level", 5),
+        processor_layers=settings.get("processor_layers", 16),
+        hidden_layers=settings.get("hidden_layers", 1),
+        hidden_dim=settings.get("hidden_dim", 512),
+        num_heads_edge=settings.get("num_heads_edge", 4),
+        num_heads_node=settings.get("num_heads_node", 4),
+        build_graph=build_graph,
+    )
+    initialization = settings.get("weight_init", "scratch")
+    if initialization == "official":
+        load_official_checkpoint(model, settings["checkpoint_path"])
+    elif initialization != "scratch":
+        raise ValueError("model.weight_init must be 'scratch' or 'official'")
+    return model
+
+
+__all__ = ["CheckpointReport", "OneForecast", "build_model", "check_checkpoint_compatibility",
+           "load_official_checkpoint", "read_official_checkpoint"]

scripts/fake_data.py

@@ -0,0 +1,102 @@
+"""Generate metadata-compatible ERA5 HDF5 fixtures at the native 0.25 degree grid."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import h5py
+import numpy as np
+
+
+OFFICIAL_VARIABLES = (
+    [f"Z{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"Q{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"T{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"U{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + [f"V{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)]
+    + ["U10M", "V10M", "T2M", "MSLP"]
+)
+
+VARIABLE_ALIASES = {
+    **{f"Z{x}": f"geopotential_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"Q{x}": f"specific_humidity_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"T{x}": f"temperature_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"U{x}": f"u_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    **{f"V{x}": f"v_component_of_wind_{x}" for x in (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)},
+    "U10M": "10m_u_component_of_wind",
+    "V10M": "10m_v_component_of_wind",
+    "T2M": "2m_temperature",
+    "MSLP": "mean_sea_level_pressure",
+}
+REAL_VARIABLES = tuple(VARIABLE_ALIASES[name] for name in OFFICIAL_VARIABLES)
+
+
+def _synthetic_field(time_index: int, channel: int, height: int, width: int) -> np.ndarray:
+    lat = np.linspace(1.0, -1.0, height, dtype=np.float32)[:, None]
+    lon = np.linspace(0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32)[None, :]
+    phase = np.float32(channel * 0.17)
+    field = np.sin(lon + phase) + 0.4 * np.cos(np.float32(time_index / 3.0) + phase) + 0.2 * lat
+    return np.asarray(field, dtype=np.float32)
+
+
+def generate_fake_h5(output_dir: Path, years: list[int], stats_years: set[int], time_steps: int,
+                     height: int, width: int, seed: int) -> None:
+    data_dir = output_dir / "data"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    sums = np.zeros(len(OFFICIAL_VARIABLES), dtype=np.float64)
+    squared_sums = np.zeros_like(sums)
+    value_count = 0
+
+    for offset, year in enumerate(years):
+        path = data_dir / f"{year}.h5"
+        with h5py.File(path, "w") as handle:
+            dataset = handle.create_dataset(
+                "fields", shape=(time_steps, len(OFFICIAL_VARIABLES), height, width),
+                dtype="float32", chunks=(1, 1, height, width),
+            )
+            dataset.attrs["variables"] = list(REAL_VARIABLES)
+            dataset.attrs["time_step"] = 6
+            for time_index in range(time_steps):
+                for channel in range(len(OFFICIAL_VARIABLES)):
+                    field = _synthetic_field(time_index + offset, channel, height, width)
+                    dataset[time_index, channel] = field
+                    if year in stats_years:
+                        sums[channel] += field.sum(dtype=np.float64)
+                        squared_sums[channel] += np.square(field, dtype=np.float64).sum()
+            if year in stats_years:
+                value_count += time_steps * height * width
+        print(f"{path}: fields={(time_steps, len(OFFICIAL_VARIABLES), height, width)}, variables={len(OFFICIAL_VARIABLES)}")
+
+    means = (sums / value_count).reshape(1, -1, 1, 1)
+    variances = squared_sums / value_count - np.square(means.reshape(-1))
+    stds = np.sqrt(np.maximum(variances, 1e-12)).reshape(1, -1, 1, 1)
+    stats_dir = output_dir / "stats"
+    stats_dir.mkdir(parents=True, exist_ok=True)
+    np.save(stats_dir / "global_means.npy", means)
+    np.save(stats_dir / "global_stds.npy", stds)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output-dir", type=Path, default=Path("./data"))
+    parser.add_argument("--years", nargs="+", type=int, default=[2000, 2001, 2002])
+    parser.add_argument("--stats-years", nargs="+", type=int, default=None)
+    parser.add_argument("--time-steps", type=int, default=3)
+    parser.add_argument("--height", type=int, default=721)
+    parser.add_argument("--width", type=int, default=1440)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+    if len(args.years) < 1 or min(args.time_steps, args.height, args.width) < 1:
+        parser.error("years, time-steps, height, and width must be positive")
+    if (args.height, args.width) != (721, 1440):
+        parser.error("ERA5 fixtures must use the native 0.25 degree grid 721x1440")
+    stats_years = set(args.stats_years or args.years[:1])
+    if not stats_years.issubset(args.years):
+        parser.error("stats-years must be included in years")
+    generate_fake_h5(args.output_dir, args.years, stats_years, args.time_steps,
+                     args.height, args.width, args.seed)
+
+
+if __name__ == "__main__":
+    main()

scripts/finetune.py

@@ -0,0 +1,131 @@
+"""Multi-step autoregressive fine-tuning for OneForecast."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.checkpoint import checkpoint
+import yaml
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from model.era5_adapter import OFFICIAL_VARIABLES, OneForecastERA5Adapter
+from model.oneforecast import build_model, read_official_checkpoint
+from scripts.train import _LossScaleFunction, _reduce_metrics, _relative_channel_l2, _set_seed, _setup_distributed
+
+
+def _resolve_path(value: str | Path, config_path: Path) -> Path:
+    path = Path(value).expanduser()
+    return path if path.is_absolute() else (config_path.parent.parent / path).resolve()
+
+
+def _prepare_batch(batch: tuple, steps: int) -> tuple[torch.Tensor, torch.Tensor]:
+    inputs, targets = batch[0], batch[1]
+    if inputs.ndim != 4 or targets.ndim != 5:
+        raise ValueError(f"Expected [B,C,H,W] inputs and [B,S,C,H,W] targets, got {inputs.shape} and {targets.shape}")
+    if targets.shape[1] != steps:
+        raise ValueError(f"Expected {steps} target steps, got {targets.shape[1]}")
+    if inputs.shape[-2] == 121:
+        inputs = inputs[..., :120, :]
+    if targets.shape[-2] == 121:
+        targets = targets[..., :120, :]
+    if inputs.shape[-2:] != (120, 240) or targets.shape[-2:] != (120, 240):
+        raise ValueError(f"Expected official model grid 120x240, got {inputs.shape} and {targets.shape}")
+    return torch.nan_to_num(inputs.float()), torch.nan_to_num(targets.float())
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
+    parser.add_argument("--model-source", choices=("trained", "official"), default=None)
+    parser.add_argument("--max-epochs", type=int, default=None)
+    parser.add_argument("--max-batches", type=int, default=None)
+    args = parser.parse_args()
+    config_path = args.config.resolve()
+    with config_path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    if tuple(config["datapipe"]["variables"]) != OFFICIAL_VARIABLES:
+        raise ValueError("datapipe.variables must exactly match the official 69-channel order")
+
+    settings = config["datapipe"]
+    finetune = config["finetuning"]
+    if args.max_epochs is not None:
+        finetune["max_epoch"] = args.max_epochs
+    if args.max_batches is not None:
+        finetune["max_batches"] = args.max_batches
+    steps = int(finetune["steps"])
+    if steps < 2:
+        raise ValueError("finetuning.steps must be at least 2")
+    source = args.model_source or finetune.get("model_source", "trained")
+    checkpoint_path = _resolve_path(
+        finetune["trained_model_path"] if source == "trained" else finetune["official_checkpoint_path"],
+        config_path,
+    )
+    output_path = _resolve_path(finetune["output_path"], config_path)
+    dataset_dir = _resolve_path(settings["dataset_dir"], config_path)
+
+    device, rank, world_size, distributed = _setup_distributed(
+        config["runtime"].get("device", "cpu"), config["runtime"].get("distributed_backend", "nccl")
+    )
+    _set_seed(int(config["runtime"].get("seed", 42)))
+    config["model"]["weight_init"] = "scratch"
+    model = build_model(config).to(device)
+    state, _ = read_official_checkpoint(checkpoint_path)
+    model.load_state_dict(state)
+    if distributed:
+        ddp_devices = {"device_ids": [device.index], "output_device": device.index} if device.type == "cuda" else {}
+        model = DistributedDataParallel(model, broadcast_buffers=False, **ddp_devices)
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(finetune["learning_rate"]))
+    adapter = OneForecastERA5Adapter(
+        dataset_dir, settings["train_years"], batch_size=settings["batch_size"],
+        input_steps=1, output_steps=steps, normalize=settings["normalize"],
+        num_workers=settings["num_workers"], distributed=distributed,
+    )
+    loader, sampler = adapter.get_dataloader("train")
+    max_batches = int(finetune.get("max_batches", -1))
+
+    for epoch in range(int(finetune["max_epoch"])):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        epoch_loss = 0.0
+        batches = 0
+        for batch in loader:
+            inputs, targets = _prepare_batch(batch, steps)
+            current = inputs.to(device)
+            targets = targets.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            loss = torch.zeros((), device=device)
+            for step in range(steps):
+                current = checkpoint(model, current, use_reentrant=False)
+                scaled = _LossScaleFunction.apply(current, 1e-5)
+                step_loss, _ = _relative_channel_l2(scaled, targets[:, step])
+                loss = loss + step_loss
+            loss.backward()
+            optimizer.step()
+            epoch_loss += float(loss.detach())
+            batches += 1
+            if max_batches >= 0 and batches >= max_batches:
+                break
+        mean_loss = _reduce_metrics(epoch_loss, batches, device, distributed)
+        if rank == 0:
+            print({"epoch": epoch + 1, "steps": steps, "loss": mean_loss,
+                   "batches_per_rank": batches, "world_size": world_size})
+
+    if rank == 0:
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        state = model.module.state_dict() if distributed else model.state_dict()
+        torch.save({"model_state": state, "epoch": int(finetune["max_epoch"]),
+                    "finetune_steps": steps, "world_size": world_size}, output_path)
+        print({"checkpoint": str(output_path)})
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,148 @@
+"""OneForecast inference entry point with the shared ERA5 adapter."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import yaml
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from model.era5_adapter import OFFICIAL_VARIABLES, OneForecastERA5Adapter
+from model.oneforecast import build_model, check_checkpoint_compatibility, read_official_checkpoint
+
+
+def _resolve_path(value: str | Path, config_path: Path) -> Path:
+    path = Path(value).expanduser()
+    return path if path.is_absolute() else (config_path.parent.parent / path).resolve()
+
+
+def _load_config(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    config["datapipe"]["dataset_dir"] = str(_resolve_path(config["datapipe"]["dataset_dir"], path))
+    config["model"]["official_checkpoint_path"] = str(
+        _resolve_path(config["model"]["official_checkpoint_path"], path)
+    )
+    config["model"]["checkpoint_path"] = config["model"]["official_checkpoint_path"]
+    config["inference"]["trained_model_path"] = str(
+        _resolve_path(config["inference"]["trained_model_path"], path)
+    )
+    config["inference"]["official_checkpoint_path"] = str(
+        _resolve_path(config["inference"]["official_checkpoint_path"], path)
+    )
+    config["inference"]["output_dir"] = str(_resolve_path(config["inference"]["output_dir"], path))
+    return config
+
+
+def _resolve_device(name: str) -> torch.device:
+    """Map the logical DCU name to the backend exposed by this PyTorch build."""
+    requested = str(name).lower()
+    if requested == "dcu":
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        privateuse = torch._C._get_privateuse1_backend_name()
+        if privateuse != "privateuseone":
+            return torch.device(privateuse)
+        raise RuntimeError("runtime.device=dcu, but this PyTorch build exposes no usable accelerator")
+    if requested == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(requested)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("runtime.device=cuda, but torch.cuda.is_available() is False")
+    return device
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
+    parser.add_argument("--check-data", action="store_true")
+    parser.add_argument("--check-model", action="store_true")
+    parser.add_argument("--check-checkpoint", action="store_true")
+    parser.add_argument("--model-source", choices=("trained", "official"), default=None)
+    args = parser.parse_args()
+    config = _load_config(args.config.resolve())
+    if tuple(config["datapipe"]["variables"]) != OFFICIAL_VARIABLES:
+        raise ValueError("datapipe.variables must exactly match the official 69-channel order")
+    if args.model_source is not None:
+        config["inference"]["model_source"] = args.model_source
+    if args.check_data:
+        settings = config["datapipe"]
+        adapter = OneForecastERA5Adapter(
+            settings["dataset_dir"], settings["test_years"], batch_size=1,
+            input_steps=settings["input_steps"], output_steps=settings["output_steps"],
+            normalize=settings["normalize"], num_workers=settings["num_workers"],
+        )
+        print(adapter.inspect())
+        return
+    if args.check_model:
+        configured_init = config["model"].get("weight_init", "scratch")
+        config["model"]["weight_init"] = "scratch"
+        with __import__("torch").device("meta"):
+            model = build_model(config, build_graph=False)
+        print({"model": type(model).__name__, "parameters": sum(p.numel() for p in model.parameters()),
+               "configured_weight_init": configured_init})
+        return
+    if args.check_checkpoint:
+        with __import__("torch").device("meta"):
+            model = build_model(config, build_graph=False)
+        report = check_checkpoint_compatibility(
+            model, config["model"]["official_checkpoint_path"]
+        )
+        print(report)
+        if not report.compatible:
+            raise SystemExit(1)
+        return
+    settings = config["datapipe"]
+    if settings["input_steps"] != 1 or settings["output_steps"] != 1:
+        raise SystemExit("OneForecast inference currently requires input_steps=1 and output_steps=1")
+    device = _resolve_device(config["runtime"].get("device", "cpu"))
+    config["model"]["weight_init"] = "scratch"
+    model = build_model(config).to(device)
+    source = config["inference"].get("model_source", "trained")
+    checkpoint_path = config["inference"][
+        "trained_model_path" if source == "trained" else "official_checkpoint_path"
+    ]
+    state, _ = read_official_checkpoint(checkpoint_path)
+    model.load_state_dict(state)
+    model.eval()
+    adapter = OneForecastERA5Adapter(
+        _resolve_path(settings["dataset_dir"], args.config), settings["test_years"],
+        batch_size=1, input_steps=1, output_steps=1,
+        normalize=settings["normalize"], num_workers=settings["num_workers"],
+    )
+    loader, _ = adapter.get_dataloader("test")
+    output_dir = Path(config["inference"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    max_batches = int(config["inference"].get("max_batches", -1))
+    processed = 0
+    with torch.no_grad():
+        for index, batch in enumerate(loader):
+            inputs, targets = batch[0].float().to(device), batch[1].float().to(device)
+            if inputs.ndim == 5 or targets.ndim == 5:
+                raise ValueError("OneForecast currently supports input_steps=1 and output_steps=1 only")
+            if inputs.ndim != 4:
+                raise ValueError(f"Expected batched input with four dimensions, got {inputs.shape}")
+            if inputs.shape[-2] == 121:
+                inputs = inputs[..., :120, :]
+            if targets.shape[-2] == 121:
+                targets = targets[..., :120, :]
+            if inputs.shape[-2:] != (120, 240) or targets.shape[-2:] != (120, 240):
+                raise ValueError(f"Expected official model grid 120x240, got {inputs.shape} and {targets.shape}")
+            prediction = model(torch.nan_to_num(inputs))
+            if settings["normalize"]:
+                means, stds = adapter.selected_statistics()
+                prediction = prediction.cpu() * torch.from_numpy(stds).float() + torch.from_numpy(means).float()
+            np.save(output_dir / f"prediction_{index:05d}.npy", prediction.cpu().numpy())
+            processed += 1
+            if max_batches >= 0 and index + 1 >= max_batches:
+                break
+    print({"output_dir": str(output_dir), "batches": processed})
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,52 @@
+"""Create quick field images from OneForecast prediction files."""
+
+from __future__ import annotations
+
+from pathlib import Path
+import argparse
+import numpy as np
+import yaml
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
+    args = parser.parse_args()
+    with args.config.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    root = args.config.resolve().parent.parent
+    input_dir = Path(config["visualization"]["input_dir"])
+    output_dir = Path(config["visualization"]["output_dir"])
+    if not input_dir.is_absolute():
+        input_dir = root / input_dir
+    if not output_dir.is_absolute():
+        output_dir = root / output_dir
+    output_dir.mkdir(parents=True, exist_ok=True)
+    files = sorted(input_dir.glob("prediction_*.npy"))
+    if not files:
+        raise SystemExit(f"No prediction files found in {input_dir}")
+    import matplotlib.pyplot as plt
+
+    channels = config["visualization"].get("channels", [0])
+    for source in files:
+        prediction = np.load(source)
+        if prediction.shape != (1, 69, 120, 240):
+            raise ValueError(f"Expected official prediction shape [1, 69, 120, 240], got {prediction.shape}")
+        field = prediction[0]
+        for channel in channels:
+            if channel < 0 or channel >= field.shape[0]:
+                raise ValueError(f"Channel {channel} is outside prediction shape {field.shape}")
+            figure, axis = plt.subplots(figsize=(8, 3.5))
+            image = axis.imshow(field[channel], cmap="coolwarm", aspect="auto")
+            axis.set_title(f"{source.stem}, channel {channel}")
+            axis.set_xlabel("longitude index")
+            axis.set_ylabel("latitude index")
+            figure.colorbar(image, ax=axis, shrink=0.8)
+            figure.tight_layout()
+            figure.savefig(output_dir / f"{source.stem}_ch{channel}.png", dpi=160)
+            plt.close(figure)
+    print({"input_dir": str(input_dir), "output_dir": str(output_dir), "files": len(files)})
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,284 @@
+"""OneForecast training entry point with integrated data checking."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+import sys
+import random
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.nn import functional as F
+from torch.autograd import Function
+from torch.nn.parallel import DistributedDataParallel
+
+import yaml
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from model.era5_adapter import OFFICIAL_VARIABLES, OneForecastERA5Adapter
+from model.oneforecast import build_model, check_checkpoint_compatibility
+
+
+def _resolve_path(value: str | Path, config_path: Path) -> Path:
+    path = Path(value).expanduser()
+    return path if path.is_absolute() else (config_path.parent.parent / path).resolve()
+
+
+def _load_config(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    config["datapipe"]["dataset_dir"] = str(_resolve_path(config["datapipe"]["dataset_dir"], path))
+    config["model"]["official_checkpoint_path"] = str(
+        _resolve_path(config["model"]["official_checkpoint_path"], path)
+    )
+    config["model"]["checkpoint_path"] = config["model"]["official_checkpoint_path"]
+    config["training"]["checkpoint_dir"] = str(_resolve_path(config["training"]["checkpoint_dir"], path))
+    return config
+
+
+def _set_seed(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def _resolve_device(name: str) -> torch.device:
+    """Map the logical DCU name to the backend exposed by this PyTorch build."""
+    requested = str(name).lower()
+    if requested == "dcu":
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        privateuse = torch._C._get_privateuse1_backend_name()
+        if privateuse != "privateuseone":
+            return torch.device(privateuse)
+        raise RuntimeError("runtime.device=dcu, but this PyTorch build exposes no usable accelerator")
+    if requested == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(requested)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("runtime.device=cuda, but torch.cuda.is_available() is False")
+    return device
+
+
+def _setup_distributed(device_name: str, backend: str = "nccl") -> tuple[torch.device, int, int, bool]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    distributed = world_size > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if distributed:
+        device = _resolve_device(device_name)
+        if device.type == "cuda":
+            torch.cuda.set_device(local_rank)
+            device = torch.device("cuda", local_rank)
+        dist.init_process_group(backend=backend, init_method="env://")
+        return device, dist.get_rank(), world_size, True
+    return _resolve_device(device_name), 0, 1, False
+
+
+def _reduce_metrics(total: float, count: int, device: torch.device, distributed: bool) -> float:
+    metrics = torch.tensor([total, count], dtype=torch.float64, device=device)
+    if distributed:
+        dist.all_reduce(metrics, op=dist.ReduceOp.SUM)
+    return float(metrics[0] / metrics[1].clamp_min(1))
+
+
+def _loader_batch(batch: tuple) -> tuple[torch.Tensor, torch.Tensor]:
+    inputs, targets = batch[0], batch[1]
+    if inputs.ndim == 5 or targets.ndim == 5:
+        raise ValueError("OneForecast currently supports input_steps=1 and output_steps=1 only")
+    if inputs.ndim != 4 or targets.ndim != 4:
+        raise ValueError(f"Expected batched fields with four dimensions, got {inputs.shape} and {targets.shape}")
+    if inputs.shape[-2] == 121:
+        inputs = inputs[..., :120, :]
+    if targets.shape[-2] == 121:
+        targets = targets[..., :120, :]
+    if inputs.shape[-2:] != (120, 240) or targets.shape[-2:] != (120, 240):
+        raise ValueError(f"Expected official model grid 120x240, got {inputs.shape} and {targets.shape}")
+    return torch.nan_to_num(inputs.float()), torch.nan_to_num(targets.float())
+
+
+class _LossScaleFunction(Function):
+    @staticmethod
+    def forward(ctx, values: torch.Tensor, eps: float) -> torch.Tensor:
+        ctx.eps = eps
+        return values
+
+    @staticmethod
+    def backward(ctx, gradients: torch.Tensor) -> tuple[torch.Tensor, None]:
+        channels = gradients.shape[1]
+        weights = 1.0 / gradients.norm(p=2, dim=(-1, -2), keepdim=True).clamp_min(ctx.eps)
+        weights = weights / weights.sum(dim=1, keepdim=True).clamp_min(ctx.eps)
+        return channels * weights * gradients, None
+
+
+def _relative_channel_l2(prediction: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    difference = (prediction - target).flatten(2).norm(p=2, dim=2)
+    target_norm = target.flatten(2).norm(p=2, dim=2).clamp_min(1e-10)
+    channel_loss = (difference / target_norm).mean(dim=0)
+    return channel_loss.mean(), channel_loss
+
+
+def check_data(config: dict) -> dict:
+    settings = config["datapipe"]
+    adapter = OneForecastERA5Adapter(
+        settings["dataset_dir"], settings["train_years"],
+        batch_size=settings["batch_size"], input_steps=settings["input_steps"],
+        output_steps=settings["output_steps"], normalize=settings["normalize"],
+        num_workers=settings["num_workers"],
+    )
+    report = adapter.inspect()
+    print(report)
+    return report
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, default=Path("conf/config.yaml"))
+    parser.add_argument("--check-data", action="store_true")
+    parser.add_argument("--check-model", action="store_true")
+    parser.add_argument("--check-checkpoint", action="store_true")
+    parser.add_argument("--check-distributed", action="store_true")
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--distributed-backend", default=None)
+    parser.add_argument("--max-epochs", type=int, default=None)
+    parser.add_argument("--max-batches", type=int, default=None)
+    parser.add_argument("--weight-init", choices=("scratch", "official"), default=None)
+    args = parser.parse_args()
+    config = _load_config(args.config.resolve())
+    if args.device is not None:
+        config["runtime"]["device"] = args.device
+    if args.distributed_backend is not None:
+        config["runtime"]["distributed_backend"] = args.distributed_backend
+    if args.max_epochs is not None:
+        config["training"]["max_epoch"] = args.max_epochs
+    if args.max_batches is not None:
+        config["training"]["max_batches"] = args.max_batches
+    if tuple(config["datapipe"]["variables"]) != OFFICIAL_VARIABLES:
+        raise ValueError("datapipe.variables must exactly match the official 69-channel order")
+    if args.weight_init is not None:
+        config["model"]["weight_init"] = args.weight_init
+    if config["model"].get("weight_init") == "official":
+        config["model"]["checkpoint_path"] = config["model"]["official_checkpoint_path"]
+    if args.check_data:
+        check_data(config)
+        return
+    if args.check_model:
+        configured_init = config["model"].get("weight_init", "scratch")
+        config["model"]["weight_init"] = "scratch"
+        with __import__("torch").device("meta"):
+            model = build_model(config, build_graph=False)
+        print({"model": type(model).__name__, "parameters": sum(p.numel() for p in model.parameters()),
+               "configured_weight_init": configured_init})
+        return
+    if args.check_checkpoint:
+        with __import__("torch").device("meta"):
+            model = build_model(config, build_graph=False)
+        report = check_checkpoint_compatibility(
+            model, config["model"]["official_checkpoint_path"]
+        )
+        print(report)
+        if not report.compatible:
+            raise SystemExit(1)
+        return
+    if args.check_distributed:
+        device, rank, world_size, distributed = _setup_distributed(
+            config["runtime"].get("device", "cpu"), config["runtime"].get("distributed_backend", "nccl")
+        )
+        settings = config["datapipe"]
+        adapter = OneForecastERA5Adapter(
+            settings["dataset_dir"], settings["train_years"], batch_size=settings["batch_size"],
+            input_steps=settings["input_steps"], output_steps=settings["output_steps"],
+            normalize=settings["normalize"], num_workers=0, distributed=distributed,
+        )
+        loader, sampler = adapter.get_dataloader("train")
+        sample_indices = list(iter(sampler)) if sampler is not None else list(range(len(loader.dataset)))
+        print({"rank": rank, "world_size": world_size, "distributed": distributed,
+               "backend": dist.get_backend() if distributed else None, "device": str(device),
+               "sampler": type(sampler).__name__ if sampler is not None else None,
+               "sample_indices": sample_indices})
+        if distributed:
+            dist.barrier()
+            dist.destroy_process_group()
+        return
+    settings = config["datapipe"]
+    if settings["input_steps"] != 1 or settings["output_steps"] != 1:
+        raise SystemExit("OneForecast training currently requires input_steps=1 and output_steps=1")
+    device, rank, world_size, distributed = _setup_distributed(
+        config["runtime"].get("device", "cpu"), config["runtime"].get("distributed_backend", "nccl")
+    )
+    _set_seed(int(config["runtime"].get("seed", 42)))
+    model = build_model(config).to(device)
+    if distributed:
+        ddp_devices = {"device_ids": [device.index], "output_device": device.index} if device.type == "cuda" else {}
+        model = DistributedDataParallel(model, broadcast_buffers=False, **ddp_devices)
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=float(config["training"]["learning_rate"]),
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=max(1, int(config["training"]["max_epoch"])),
+    )
+    train_adapter = OneForecastERA5Adapter(
+        _resolve_path(settings["dataset_dir"], args.config), settings["train_years"],
+        batch_size=settings["batch_size"], input_steps=1, output_steps=1,
+        normalize=settings["normalize"], num_workers=settings["num_workers"], distributed=distributed,
+    )
+    valid_adapter = OneForecastERA5Adapter(
+        _resolve_path(settings["dataset_dir"], args.config), settings["valid_years"],
+        batch_size=settings["batch_size"], input_steps=1, output_steps=1,
+        normalize=settings["normalize"], num_workers=settings["num_workers"], distributed=distributed,
+    )
+    train_loader, train_sampler = train_adapter.get_dataloader("train")
+    valid_loader, valid_sampler = valid_adapter.get_dataloader("val")
+    checkpoint_dir = Path(config["training"]["checkpoint_dir"])
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    max_batches = int(config["training"].get("max_batches", -1))
+    for epoch in range(int(config["training"]["start_epoch"]), int(config["training"]["max_epoch"])):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        if valid_sampler is not None:
+            valid_sampler.set_epoch(epoch)
+        model.train()
+        train_loss = 0.0
+        train_batches = 0
+        for batch in train_loader:
+            inputs, targets = _loader_batch(batch)
+            optimizer.zero_grad(set_to_none=True)
+            prediction = _LossScaleFunction.apply(model(inputs.to(device)), 1e-5)
+            loss, _ = _relative_channel_l2(prediction, targets.to(device))
+            loss.backward()
+            optimizer.step()
+            train_loss += float(loss.detach())
+            train_batches += 1
+            if max_batches >= 0 and train_batches >= max_batches:
+                break
+        model.eval()
+        valid_loss = 0.0
+        with torch.no_grad():
+            valid_batches = 0
+            for batch in valid_loader:
+                inputs, targets = _loader_batch(batch)
+                prediction = model(inputs.to(device))
+                valid_loss += float(F.mse_loss(prediction, targets.to(device)))
+                valid_batches += 1
+                if max_batches >= 0 and valid_batches >= max_batches:
+                    break
+        train_mean = _reduce_metrics(train_loss, train_batches, device, distributed)
+        valid_mean = _reduce_metrics(valid_loss, valid_batches, device, distributed)
+        if rank == 0:
+            print({"epoch": epoch + 1, "train_loss": train_mean, "valid_loss": valid_mean,
+                   "world_size": world_size})
+        if rank == 0 and (epoch + 1) % int(config["training"].get("save_every_epoch", 1)) == 0:
+            model_name = config["training"].get("model_name", "model_bak")
+            state = model.module.state_dict() if distributed else model.state_dict()
+            torch.save({"model_state": state, "epoch": epoch + 1, "world_size": world_size},
+                       checkpoint_dir / f"{model_name}.tar")
+        scheduler.step()
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()