Upload folder using huggingface_hub

README.md

@@ -0,0 +1,170 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Subseasonal Forecast
+- Coupled Atmosphere-Ocean-Land
+- ERA5
+tasks: []
+datasets:
+- OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">FengWu-W2S</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+FengWu-W2S (FengWu Weather-to-Subseasonal) is a seamless global weather-to-subseasonal forecasting model extending FengWu.
+
+Paper: FengWu-W2S: A deep learning model for seamless weather-to-subseasonal forecast of global atmosphere
+
+https://arxiv.org/abs/2411.10191
+
+# Model Description
+
+The model uses a six-hour time step for autoregressive forecasts of up to 42 days. Coupled atmospheric, ocean, and land branches, together with diversity perturbations, are used to improve extended-range forecast skill. This repository contains a compact implementation of the coupled interfaces for reproducible workflow checks.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Weather-to-subseasonal forecast research | Train a 78-channel coupled model with six-hourly global atmospheric, ocean, and land fields. |
+| Local quick validation | Use small-grid synthetic HDF5 data to check training, fine-tuning, inference, and forecast visualization. |
+| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Launch PyTorch DistributedDataParallel with `torchrun`. |
+
+# 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/FengWu-W2S --local-dir ./FengWu-W2S
+cd FengWu-W2S
+```
+
+### 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 an ERA5 data slice for training. Download it and confirm that the data path in `conf/config.yaml` is correct:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+Real files must contain all 78 channels listed in `conf/config.yaml`, including atmospheric, ocean, and land variables, with six-hour time spacing. The configured group indices define the coupled branches.
+
+### Generate Synthetic Data
+
+The default configuration describes a full `721x1440` grid. For a practical local smoke test, generate a small fixture explicitly:
+
+```bash
+python scripts/fake_data.py --height 32 --width 64 --timesteps 12
+```
+
+Synthetic HDF5 files are for workflow validation only and do not represent the ERA5 reanalysis or the paper's forecast quality.
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=2 scripts/train.py
+```
+
+For a short smoke run, reduce the workload while keeping the generated data and model grid consistent:
+
+```bash
+python scripts/train.py --max-epoch 1 --batch-size 1 --num-workers 0 --rollout-steps 1
+```
+
+The default checkpoint is saved to `data/checkpoints/model_bak.pth`.
+
+### Fine-tuning
+
+Resume from the base checkpoint with the lower fine-tuning learning rate:
+
+```bash
+python scripts/train.py --finetune
+```
+
+### Training Weights
+
+This repository provides a `weight/` directory for FengWu-W2S checkpoints. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth` by default and writes forecasts grouped by initialization year to `result/output/<year>/`, together with `result/output/index.json`:
+
+```bash
+python scripts/inference.py
+```
+
+Use `--steps` and `--limit` to bound a local smoke test; `--stochastic` enables perturbation sampling.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The script computes per-channel RMSE, normalized RMSE, and anomaly ACC, and generates forecast-comparison, lead-time skill, channel-ranking, and training-loss figures.
+
+# Official OneScience Resources
+
+| 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 and License
+
+- Paper: https://arxiv.org/abs/2411.10191
+- This directory is an independent reproduction of the paper method and does not represent official code, weights, or training results released by the authors.

conf/config.yaml

@@ -0,0 +1,133 @@
+model:
+  name: FengWu-W2S
+  default_checkpoint: ./data/checkpoints/model_bak.pth
+  input_steps: 2
+  rollout_steps: 2
+  inference_steps: 4
+  patch_size: 4
+  hidden_channels: 32
+  latent_channels: 64
+  num_blocks: 2
+  perturbation_scale: 0.01
+  learning_rate: 0.0005
+  finetune_learning_rate: 0.00005
+  adam_beta1: 0.9
+  adam_beta2: 0.9
+  probability_loss: true
+  kl_weight: 0.0001
+  initial_perturbation_scale: 0.01
+  weight_decay: 0.00001
+  max_epoch: 10
+  finetune_epoch: 40
+  patience: 10
+  checkpoint_dir: ./data/checkpoints
+
+data:
+  data_dir: ./data
+  train_years: [1999,2000]
+  val_years: [2001]
+  test_years: [2002]
+  time_step_hours: 6
+  channels:
+    - geopotential_50
+    - geopotential_100
+    - geopotential_150
+    - geopotential_200
+    - geopotential_250
+    - geopotential_300
+    - geopotential_400
+    - geopotential_500
+    - geopotential_600
+    - geopotential_700
+    - geopotential_850
+    - geopotential_925
+    - geopotential_1000
+    - temperature_50
+    - temperature_100
+    - temperature_150
+    - temperature_200
+    - temperature_250
+    - temperature_300
+    - temperature_400
+    - temperature_500
+    - temperature_600
+    - temperature_700
+    - temperature_850
+    - temperature_925
+    - temperature_1000
+    - u_component_of_wind_50
+    - u_component_of_wind_100
+    - u_component_of_wind_150
+    - u_component_of_wind_200
+    - u_component_of_wind_250
+    - u_component_of_wind_300
+    - u_component_of_wind_400
+    - u_component_of_wind_500
+    - u_component_of_wind_600
+    - u_component_of_wind_700
+    - u_component_of_wind_850
+    - u_component_of_wind_925
+    - u_component_of_wind_1000
+    - v_component_of_wind_50
+    - v_component_of_wind_100
+    - v_component_of_wind_150
+    - v_component_of_wind_200
+    - v_component_of_wind_250
+    - v_component_of_wind_300
+    - v_component_of_wind_400
+    - v_component_of_wind_500
+    - v_component_of_wind_600
+    - v_component_of_wind_700
+    - v_component_of_wind_850
+    - v_component_of_wind_925
+    - v_component_of_wind_1000
+    - specific_humidity_50
+    - specific_humidity_100
+    - specific_humidity_150
+    - specific_humidity_200
+    - specific_humidity_250
+    - specific_humidity_300
+    - specific_humidity_400
+    - specific_humidity_500
+    - specific_humidity_600
+    - specific_humidity_700
+    - specific_humidity_850
+    - specific_humidity_925
+    - specific_humidity_1000
+    - 2m_temperature
+    - total_precipitation
+    - mean_sea_level_pressure
+    - 10m_u_component_of_wind
+    - 10m_v_component_of_wind
+    - 100m_u_component_of_wind
+    - 100m_v_component_of_wind
+    - outgoing_longwave_radiation
+    - sea_surface_temperature
+    - significant_wave_height
+    - mean_wave_direction
+    - mean_wave_period
+    - soil_moisture_0_7cm
+  groups:
+    atmosphere: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64]
+    ocean: [73, 74, 75, 76]
+    land: [65, 66, 67, 68, 69, 70, 71, 72, 77]
+  normalization_epsilon: 1.0e-6
+  dataloader:
+    batch_size: 1
+    num_workers: 0
+    pin_memory: false
+    drop_last: false
+
+fake_data:
+  # Small defaults keep the workflow runnable on a login node. Real ERA5 can use 721x1440.
+  # height: 32
+  height: 721
+  # width: 64
+  width: 1440
+  timesteps: 12
+  seed: 42
+
+runtime:
+  # Fail fast instead of silently training on CPU when the DCU runtime is absent.
+  device: cuda
+  amp: false

config.json

@@ -0,0 +1,41 @@
+{
+  "model_name": "FengWu-W2S",
+  "model_type": "fengwu_w2s",
+  "architectures": [
+    "FengWuW2S"
+  ],
+  "framework": "PyTorch",
+  "domain": "coupled-atmosphere-ocean-land",
+  "task": "seamless-weather-to-subseasonal-forecasting",
+  "implementation": {
+    "entry_point": "model/fengwu_w2s.py",
+    "scope": "compact coupled weather, ocean, and land model with stochastic perturbations for autoregressive rollout"
+  },
+  "architecture": {
+    "family": "coupled branch encoders with latent fusion and perturbation module",
+    "input_steps": 2,
+    "rollout_steps": 2,
+    "inference_steps": 4,
+    "input_channels": 78,
+    "atmosphere_channels": 65,
+    "ocean_channels": 4,
+    "land_channels": 9,
+    "hidden_channels": 32,
+    "latent_channels": 64,
+    "patch_size": 4,
+    "blocks": 2,
+    "perturbation_scale": 0.01
+  },
+  "data": {
+    "dataset": "ERA5 with coupled ocean and land fields",
+    "spatial_resolution_degrees": 0.25,
+    "time_step_hours": 6,
+    "protocol": "78-channel weather-to-subseasonal contract"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/fengwu_w2s.py",
+    "scripts/data_loader.py",
+    "scripts/train.py"
+  ]
+}

configuration.json

@@ -0,0 +1,14 @@
+{
+  "framework": "PyTorch",
+  "task": "subseasonal_forecasting",
+  "model": "FengWu-W2S",
+  "input_format": "BTCHW",
+  "protocol": "independent_reproduction",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py",
+  "paper": "https://arxiv.org/abs/2411.10191",
+  "license": "Apache License 2.0"
+}

model/fengwu_w2s.py

@@ -0,0 +1,252 @@
+"""Self-contained FengWu-W2S style coupled weather model.
+
+The original FengWu-W2S extends FengWu with coupled ocean-atmosphere-land
+states and perturbations for seamless subseasonal rollout.  This module keeps
+those interfaces explicit while using compact convolutional blocks so the
+workflow can be smoke-tested on small synthetic grids.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Iterable, Mapping, Optional, Sequence
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def _group_count(channels: int, maximum: int = 8) -> int:
+    for groups in range(min(maximum, channels), 0, -1):
+        if channels % groups == 0:
+            return groups
+    return 1
+
+
+class ResidualConvBlock(nn.Module):
+    """Local periodic convolution block used by each FengWu branch."""
+
+    def __init__(self, channels: int, dilation: int = 1) -> None:
+        super().__init__()
+        padding = dilation
+        self.norm1 = nn.GroupNorm(_group_count(channels), channels)
+        self.conv1 = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=3,
+            padding=padding,
+            dilation=dilation,
+            padding_mode="circular",
+        )
+        self.norm2 = nn.GroupNorm(_group_count(channels), channels)
+        self.conv2 = nn.Conv2d(
+            channels, channels, kernel_size=3, padding=1, padding_mode="circular"
+        )
+        self.activation = nn.GELU()
+
+    def forward(self, x: Tensor) -> Tensor:
+        residual = x
+        x = self.conv1(self.activation(self.norm1(x)))
+        x = self.conv2(self.activation(self.norm2(x)))
+        return residual + x
+
+
+class BranchEncoder(nn.Module):
+    """Patch embedding followed by compact Swin/FengWu-like local blocks."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        latent_channels: int,
+        patch_size: int,
+        num_blocks: int,
+    ) -> None:
+        super().__init__()
+        self.patch_size = max(1, int(patch_size))
+        self.stem = nn.Conv2d(
+            in_channels,
+            hidden_channels,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+        blocks = []
+        for block_index in range(max(1, int(num_blocks))):
+            blocks.append(ResidualConvBlock(hidden_channels, dilation=1 + block_index % 2))
+        self.blocks = nn.Sequential(*blocks)
+        self.projection = nn.Conv2d(hidden_channels, latent_channels, kernel_size=1)
+
+    def forward(self, x: Tensor, output_size: Sequence[int]) -> Tensor:
+        # Padding here is implicit in interpolation, which also handles odd
+        # real-world grids such as 721x1440 with a compact patch embedding.
+        x = self.stem(x)
+        x = self.blocks(x)
+        x = self.projection(x)
+        return F.interpolate(x, size=tuple(output_size), mode="bilinear", align_corners=False)
+
+
+class CoupledFusion(nn.Module):
+    """Fuse atmosphere/ocean/land latents with a learned global gate."""
+
+    def __init__(self, latent_channels: int, branches: int = 3) -> None:
+        super().__init__()
+        total = latent_channels * branches
+        self.mix = nn.Sequential(
+            nn.Conv2d(total, latent_channels, kernel_size=1),
+            nn.GroupNorm(_group_count(latent_channels), latent_channels),
+            nn.GELU(),
+            ResidualConvBlock(latent_channels),
+        )
+        self.gate = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(latent_channels, latent_channels, kernel_size=1),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, features: Iterable[Tensor]) -> Tensor:
+        fused = self.mix(torch.cat(tuple(features), dim=1))
+        return fused * (1.0 + self.gate(fused))
+
+
+class BranchDecoder(nn.Module):
+    def __init__(self, latent_channels: int, out_channels: int, num_blocks: int) -> None:
+        super().__init__()
+        blocks = [ResidualConvBlock(latent_channels) for _ in range(max(1, int(num_blocks)))]
+        self.blocks = nn.Sequential(*blocks)
+        self.head = nn.Conv2d(latent_channels, out_channels, kernel_size=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.head(self.blocks(x))
+
+
+class PerturbationModule(nn.Module):
+    """Inject scale-controlled latent perturbations for ensemble members."""
+
+    def __init__(self, channels: int, scale: float = 0.01) -> None:
+        super().__init__()
+        self.scale = float(scale)
+        self.correlation = nn.Conv2d(channels, channels, kernel_size=1, bias=False)
+        self.mean_head = nn.Conv2d(channels, channels, kernel_size=1)
+        self.logvar_head = nn.Conv2d(channels, channels, kernel_size=1)
+        nn.init.eye_(self.correlation.weight.squeeze(-1).squeeze(-1))
+
+    def forward(self, latent: Tensor, enabled: bool, return_kl: bool = False):
+        if not enabled or self.scale == 0.0:
+            return (latent, latent.new_zeros(())) if return_kl else latent
+        noise = torch.randn_like(latent)
+        mean = self.mean_head(latent)
+        logvar = self.logvar_head(latent).clamp(-6.0, 2.0)
+        sample = mean + torch.exp(0.5 * logvar) * self.correlation(noise)
+        kl = 0.5 * torch.mean(torch.exp(logvar) + mean.square() - 1.0 - logvar)
+        perturbed = latent + self.scale * sample
+        return (perturbed, kl) if return_kl else perturbed
+
+
+class FengWuW2S(nn.Module):
+    """Coupled autoregressive FengWu-W2S model.
+
+    Inputs and outputs use ``[batch, channels, latitude, longitude]``.  The
+    channel groups can be changed without changing the model interface.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        group_indices: Mapping[str, Sequence[int]],
+        hidden_channels: int = 32,
+        latent_channels: int = 64,
+        patch_size: int = 4,
+        num_blocks: int = 2,
+        perturbation_scale: float = 0.01,
+    ) -> None:
+        super().__init__()
+        required = {"atmosphere", "ocean", "land"}
+        missing = required.difference(group_indices)
+        if missing:
+            raise ValueError(f"group_indices missing required groups: {sorted(missing)}")
+        self.in_channels = int(in_channels)
+        self.group_indices = {name: tuple(int(i) for i in indices) for name, indices in group_indices.items()}
+        all_indices = [i for values in self.group_indices.values() for i in values]
+        if sorted(all_indices) != list(range(self.in_channels)):
+            raise ValueError("group_indices must partition all input channels")
+
+        self.encoders = nn.ModuleDict(
+            {
+                name: BranchEncoder(
+                    len(indices), hidden_channels, latent_channels, patch_size, num_blocks
+                )
+                for name, indices in self.group_indices.items()
+            }
+        )
+        self.fusion = CoupledFusion(latent_channels, branches=3)
+        self.perturbation = PerturbationModule(latent_channels, perturbation_scale)
+        # A per-channel log scale makes the FengWu Gaussian probability loss
+        # explicit while keeping inference output equal to the predicted mean.
+        self.log_std = nn.Parameter(torch.zeros(self.in_channels, 1, 1))
+        self.decoders = nn.ModuleDict(
+            {
+                name: BranchDecoder(latent_channels, len(indices), num_blocks)
+                for name, indices in self.group_indices.items()
+            }
+        )
+
+    def _split(self, state: Tensor) -> Dict[str, Tensor]:
+        return {
+            name: state[:, indices, ...]
+            for name, indices in self.group_indices.items()
+        }
+
+    def forward(
+        self,
+        state: Tensor,
+        stochastic: Optional[bool] = None,
+        return_aux: bool = False,
+    ):
+        if state.ndim != 4 or state.shape[1] != self.in_channels:
+            raise ValueError(
+                f"expected [B,{self.in_channels},H,W], got {tuple(state.shape)}"
+            )
+        branches = self._split(state)
+        spatial = state.shape[-2:]
+        features = [self.encoders[name](branches[name], spatial) for name in ("atmosphere", "ocean", "land")]
+        latent = self.fusion(features)
+        latent, kl = self.perturbation(
+            latent,
+            self.training if stochastic is None else stochastic,
+            return_kl=True,
+        )
+        delta = state.new_zeros(state.shape)
+        for name, indices in self.group_indices.items():
+            delta[:, indices, ...] = self.decoders[name](latent)
+        # Predicting an increment stabilizes long autoregressive rollouts.
+        result = state + delta
+        if return_aux:
+            return result, {"kl": kl, "log_std": self.log_std}
+        return result
+
+    @torch.no_grad()
+    def rollout(
+        self,
+        initial_state: Tensor,
+        steps: int,
+        stochastic: bool = False,
+        seed: Optional[int] = None,
+    ) -> Tensor:
+        if initial_state.ndim == 5:
+            state = initial_state[:, -1]
+        elif initial_state.ndim == 4:
+            state = initial_state
+        else:
+            raise ValueError("initial_state must be [B,T,C,H,W] or [B,C,H,W]")
+        if seed is not None:
+            torch.manual_seed(int(seed))
+        predictions = []
+        was_training = self.training
+        self.eval()
+        for _ in range(max(0, int(steps))):
+            state = self.forward(state, stochastic=stochastic)
+            predictions.append(state)
+        if was_training:
+            self.train()
+        if not predictions:
+            return state.unsqueeze(1)[:, :0]
+        return torch.stack(predictions, dim=1)

scripts/data_loader.py

@@ -0,0 +1,152 @@
+"""ERA5-style HDF5 loader used by training and inference."""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Dict, List, Sequence, Tuple
+
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+
+
+def _decode(value: object) -> str:
+    return value.decode() if isinstance(value, bytes) else str(value)
+
+
+def resolve_data_dir(path: str | Path, project_root: Path | None = None) -> Path:
+    root = project_root or Path(__file__).resolve().parents[1]
+    candidate = Path(path).expanduser()
+    return candidate if candidate.is_absolute() else (root / candidate).resolve()
+
+
+def read_metadata(data_dir: str | Path, channels: Sequence[str]) -> Dict[str, np.ndarray | int | List[str]]:
+    data_dir = Path(data_dir)
+    files = sorted((data_dir / "data").glob("*.h5"))
+    if not files:
+        raise FileNotFoundError(f"No yearly HDF5 files found under {data_dir / 'data'}")
+    with h5py.File(files[0], "r") as source:
+        fields = source["fields"]
+        variables = [_decode(item) for item in fields.attrs["variables"]]
+        time_step = int(fields.attrs.get("time_step", 6))
+        means = np.asarray(source["global_means"][:], dtype=np.float32)
+        stds = np.asarray(source["global_stds"][:], dtype=np.float32)
+        shape = tuple(int(item) for item in fields.shape)
+    missing = [name for name in channels if name not in variables]
+    if missing:
+        raise ValueError(f"Variables missing from synthetic/ERA5 data: {missing}")
+    indices = np.asarray([variables.index(name) for name in channels], dtype=np.int64)
+    return {
+        "variables": variables,
+        "indices": indices,
+        "time_step": time_step,
+        "means": means[:, indices, :, :],
+        "stds": np.maximum(stds[:, indices, :, :], 1.0e-6),
+        "shape": shape,
+    }
+
+
+class ERA5WindowDataset(Dataset):
+    def __init__(
+        self,
+        data_dir: str | Path,
+        years: Sequence[int],
+        channels: Sequence[str],
+        input_steps: int = 2,
+        rollout_steps: int = 1,
+        normalize: bool = True,
+    ) -> None:
+        self.data_dir = Path(data_dir)
+        self.years = [int(year) for year in years]
+        self.channels = list(channels)
+        self.input_steps = max(1, int(input_steps))
+        self.rollout_steps = max(1, int(rollout_steps))
+        self.normalize = bool(normalize)
+        metadata = read_metadata(self.data_dir, self.channels)
+        self.channel_indices = metadata["indices"]
+        self.time_step = int(metadata["time_step"])
+        self.means = torch.from_numpy(metadata["means"])
+        self.stds = torch.from_numpy(metadata["stds"])
+        self.shape = metadata["shape"]
+        self.samples: List[Tuple[int, int]] = []
+        for year in self.years:
+            path = self.data_dir / "data" / f"{year}.h5"
+            if not path.exists():
+                raise FileNotFoundError(f"Missing year file: {path}")
+            with h5py.File(path, "r") as source:
+                timesteps = int(source["fields"].shape[0])
+            count = timesteps - self.input_steps - self.rollout_steps + 1
+            if count <= 0:
+                raise ValueError(
+                    f"Year {year} has {timesteps} steps, but input={self.input_steps} "
+                    f"and rollout={self.rollout_steps} require at least {self.input_steps + self.rollout_steps}"
+                )
+            self.samples.extend((year, index) for index in range(count))
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def _timestamp(self, year: int, index: int) -> str:
+        value = datetime(year, 1, 1) + timedelta(hours=index * self.time_step)
+        return value.strftime("%Y%m%d%H")
+
+    def __getitem__(self, item: int):
+        year, start = self.samples[item]
+        path = self.data_dir / "data" / f"{year}.h5"
+        with h5py.File(path, "r") as source:
+            fields = source["fields"]
+            input_data = np.asarray(
+                fields[start : start + self.input_steps, self.channel_indices, :, :], dtype=np.float32
+            )
+            target_data = np.asarray(
+                fields[
+                    start + self.input_steps : start + self.input_steps + self.rollout_steps,
+                    self.channel_indices,
+                    :,
+                    :,
+                ],
+                dtype=np.float32,
+            )
+        input_tensor = torch.from_numpy(input_data)
+        target_tensor = torch.from_numpy(target_data)
+        if self.normalize:
+            input_tensor = (input_tensor - self.means) / self.stds
+            target_tensor = (target_tensor - self.means) / self.stds
+        timestamp = self._timestamp(year, start + self.input_steps)
+        return input_tensor, target_tensor, timestamp
+
+
+def make_dataloader(
+    data_dir: str | Path,
+    years: Sequence[int],
+    channels: Sequence[str],
+    input_steps: int,
+    rollout_steps: int,
+    batch_size: int,
+    num_workers: int = 0,
+    distributed: bool = False,
+    train: bool = False,
+    pin_memory: bool = False,
+):
+    dataset = ERA5WindowDataset(
+        data_dir=data_dir,
+        years=years,
+        channels=channels,
+        input_steps=input_steps,
+        rollout_steps=rollout_steps,
+    )
+    sampler = DistributedSampler(dataset, shuffle=train) if distributed else None
+    loader = DataLoader(
+        dataset,
+        batch_size=max(1, int(batch_size)),
+        shuffle=train and sampler is None,
+        sampler=sampler,
+        num_workers=max(0, int(num_workers)),
+        pin_memory=bool(pin_memory),
+        drop_last=False,
+    )
+    return loader, sampler
+

scripts/fake_data.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""Generate a compact deterministic ERA5-style dataset for workflow checks."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import h5py
+import numpy as np
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--output-dir", help="Override data.data_dir from the config")
+    parser.add_argument("--height", type=int)
+    parser.add_argument("--width", type=int)
+    parser.add_argument("--timesteps", type=int)
+    parser.add_argument("--seed", type=int)
+    return parser.parse_args()
+
+
+def _resolve(path: str, root: Path) -> Path:
+    candidate = Path(path).expanduser()
+    return candidate if candidate.is_absolute() else (root / candidate).resolve()
+
+
+def generate_dataset(
+    output_dir: str | Path,
+    years: list[int],
+    channels: list[str],
+    height: int,
+    width: int,
+    timesteps: int,
+    time_step_hours: int = 6,
+    seed: int = 42,
+) -> Path:
+    """Write yearly HDF5 files and static fields without external data access."""
+    if min(height, width, timesteps) <= 0:
+        raise ValueError("height, width and timesteps must be positive")
+    output_dir = Path(output_dir)
+    data_dir = output_dir / "data"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    rng = np.random.default_rng(seed)
+    channels = list(channels)
+    channel_count = len(channels)
+    channel_mean = np.linspace(0.0, 0.9, channel_count, dtype=np.float32)
+    channel_std = np.linspace(0.8, 1.2, channel_count, dtype=np.float32)
+    latitude = np.linspace(1.0, -1.0, height, dtype=np.float32)[None, :, None]
+    longitude = np.linspace(0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32)[None, None, :]
+    base_pattern = 0.15 * latitude + 0.08 * np.sin(longitude)
+
+    for year_offset, year in enumerate(sorted(set(int(item) for item in years))):
+        path = data_dir / f"{year}.h5"
+        with h5py.File(path, "w") as output:
+            fields = output.create_dataset(
+                "fields",
+                shape=(timesteps, channel_count, height, width),
+                dtype="float32",
+                chunks=(1, channel_count, height, width),
+                compression="lzf",
+            )
+            fields.attrs["variables"] = np.asarray(channels, dtype="S")
+            fields.attrs["time_step"] = int(time_step_hours)
+            for timestep in range(timesteps):
+                values = np.empty((channel_count, height, width), dtype=np.float32)
+                for channel in range(channel_count):
+                    phase = 0.05 * (timestep + year_offset) + 0.2 * channel
+                    smooth = base_pattern + 0.03 * np.sin(phase)
+                    noise = rng.normal(0.0, 0.01, size=(height, width)).astype(np.float32)
+                    values[channel] = channel_mean[channel] + channel_std[channel] * (smooth + noise)
+                fields[timestep] = values
+            output.create_dataset("global_means", data=channel_mean.reshape(1, channel_count, 1, 1))
+            output.create_dataset("global_stds", data=channel_std.reshape(1, channel_count, 1, 1))
+        print(f"Generated {path} shape=({timesteps},{channel_count},{height},{width})")
+
+    static_dir = output_dir / "static"
+    static_dir.mkdir(parents=True, exist_ok=True)
+    latitude_2d = latitude.squeeze(0)
+    land_mask = np.broadcast_to((latitude_2d > 0).astype(np.float32), (height, width)).copy()
+    orography = (0.2 * np.cos(longitude) + 0.05 * latitude).astype(np.float32).squeeze(0)
+    sea_ice_mask = np.broadcast_to((latitude_2d < -0.55).astype(np.float32), (height, width)).copy()
+    np.save(static_dir / "land_mask.npy", land_mask)
+    np.save(static_dir / "orography.npy", orography)
+    np.save(static_dir / "sea_ice_mask.npy", sea_ice_mask)
+    metadata = {
+        "format": "ERA5-HDF5-window-v1",
+        "years": sorted(set(int(item) for item in years)),
+        "channels": channels,
+        "shape": [timesteps, channel_count, height, width],
+        "time_step_hours": int(time_step_hours),
+        "seed": int(seed),
+    }
+    (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
+    print(f"Generated static fields and metadata under {output_dir}")
+    return output_dir
+
+
+def main() -> None:
+    args = parse_args()
+    with Path(args.config).open(encoding="utf-8") as source:
+        config = yaml.safe_load(source)
+    data_cfg = config["data"]
+    fake_cfg = config.get("fake_data", {})
+    data_dir = _resolve(args.output_dir or data_cfg["data_dir"], PROJECT_ROOT)
+    years = list(data_cfg["train_years"]) + list(data_cfg["val_years"]) + list(data_cfg["test_years"])
+    generate_dataset(
+        data_dir,
+        years=years,
+        channels=list(data_cfg["channels"]),
+        height=int(args.height or fake_cfg.get("height", 32)),
+        width=int(args.width or fake_cfg.get("width", 64)),
+        timesteps=int(args.timesteps or fake_cfg.get("timesteps", 12)),
+        time_step_hours=int(data_cfg.get("time_step_hours", 6)),
+        seed=int(args.seed if args.seed is not None else fake_cfg.get("seed", 42)),
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+"""Run autoregressive FengWu-W2S inference and save forecast fields."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.fengwu_w2s import FengWuW2S
+from scripts.data_loader import ERA5WindowDataset, read_metadata, resolve_data_dir
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--data-dir")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--output-dir")
+    parser.add_argument("--split", choices=("train", "val", "test"), default="test")
+    parser.add_argument("--steps", type=int)
+    parser.add_argument("--limit", type=int, help="Limit windows for a quick smoke test")
+    parser.add_argument("--stochastic", action="store_true")
+    parser.add_argument("--seed", type=int, default=0)
+    return parser.parse_args()
+
+
+def _resolve(path: str | Path, root: Path = PROJECT_ROOT) -> Path:
+    candidate = Path(path).expanduser()
+    return candidate if candidate.is_absolute() else (root / candidate).resolve()
+
+
+def _build_model(config: dict) -> FengWuW2S:
+    model_cfg = config["model"]
+    data_cfg = config["data"]
+    return FengWuW2S(
+        in_channels=len(data_cfg["channels"]),
+        group_indices=data_cfg["groups"],
+        hidden_channels=int(model_cfg.get("hidden_channels", 32)),
+        latent_channels=int(model_cfg.get("latent_channels", 64)),
+        patch_size=int(model_cfg.get("patch_size", 4)),
+        num_blocks=int(model_cfg.get("num_blocks", 2)),
+        perturbation_scale=float(model_cfg.get("perturbation_scale", 0.01)),
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    with Path(args.config).open(encoding="utf-8") as source:
+        config = yaml.safe_load(source)
+    data_cfg = config["data"]
+    model_cfg = config["model"]
+    data_dir = resolve_data_dir(args.data_dir or data_cfg["data_dir"], PROJECT_ROOT)
+    checkpoint = _resolve(args.checkpoint or model_cfg.get("default_checkpoint", "./data/checkpoints/model_bak.pth"))
+    if not checkpoint.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint}; run scripts/train.py first")
+    split_years = {
+        "train": data_cfg["train_years"],
+        "val": data_cfg["val_years"],
+        "test": data_cfg["test_years"],
+    }[args.split]
+    steps = int(args.steps or model_cfg.get("inference_steps", 1))
+    dataset = ERA5WindowDataset(
+        data_dir=data_dir,
+        years=split_years,
+        channels=data_cfg["channels"],
+        input_steps=int(model_cfg.get("input_steps", 2)),
+        rollout_steps=max(1, steps),
+    )
+    metadata = read_metadata(data_dir, data_cfg["channels"])
+    means = torch.from_numpy(metadata["means"])
+    stds = torch.from_numpy(metadata["stds"])
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = _build_model(config).to(device)
+    state = torch.load(checkpoint, map_location=device, weights_only=False)
+    model.load_state_dict(state["model_state_dict"])
+    model.eval()
+    output_dir = _resolve(args.output_dir or "./result/output")
+    output_dir.mkdir(parents=True, exist_ok=True)
+    index = []
+    limit = len(dataset) if args.limit is None else min(len(dataset), max(0, int(args.limit)))
+    for sample_index in range(limit):
+        inputs, _, timestamp = dataset[sample_index]
+        initial = inputs.unsqueeze(0).to(device)
+        with torch.no_grad():
+            prediction = model.rollout(initial, steps=steps, stochastic=args.stochastic, seed=args.seed + sample_index)
+        prediction = prediction.squeeze(0).cpu() * stds + means
+        year = timestamp[:4]
+        year_dir = output_dir / year
+        year_dir.mkdir(parents=True, exist_ok=True)
+        start = datetime.strptime(timestamp, "%Y%m%d%H")
+        for lead, field in enumerate(prediction):
+            valid_time = start + timedelta(hours=lead * int(metadata["time_step"]))
+            valid_timestamp = valid_time.strftime("%Y%m%d%H")
+            path = year_dir / f"{timestamp}_lead{lead:03d}.npy"
+            np.save(path, field.numpy().astype(np.float32))
+            index.append({
+                "source_timestamp": timestamp,
+                "valid_timestamp": valid_timestamp,
+                "lead": lead,
+                # Keep the manifest relocatable when the project is copied to
+                # a compute node or downloaded from ModelScope.
+                "path": str(path.relative_to(output_dir)),
+            })
+    (output_dir / "index.json").write_text(json.dumps(index, indent=2), encoding="utf-8")
+    print(f"Saved {len(index)} forecast fields under {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+"""Evaluate forecasts and render meteorological diagnostic figures."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections import defaultdict
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any, Dict, List, Mapping, Sequence, Tuple
+
+import h5py
+import numpy as np
+import yaml
+
+try:
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    from matplotlib.colors import Normalize, TwoSlopeNorm
+except ImportError as exc:  # pragma: no cover
+    raise RuntimeError(
+        "scripts/result.py requires matplotlib for scientific figures"
+    ) from exc
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from scripts.data_loader import read_metadata, resolve_data_dir
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--data-dir")
+    parser.add_argument("--output-dir", default=str(PROJECT_ROOT / "result/output"))
+    parser.add_argument("--plot-count", type=int, default=3, help="Initialization times to visualize")
+    parser.add_argument("--plot-leads", type=int, default=3, help="Lead times per initialization to visualize")
+    parser.add_argument("--dpi", type=int, default=160)
+    return parser.parse_args()
+
+
+def _pretty_name(name: str) -> str:
+    replacements = {
+        "geopotential_": "Geopotential ",
+        "temperature_": "Temperature ",
+        "specific_humidity_": "Specific humidity ",
+        "2m_temperature": "2 m temperature",
+        "total_precipitation": "Total precipitation",
+        "mean_sea_level_pressure": "Mean sea-level pressure",
+        "10m_u_component_of_wind": "10 m U wind",
+        "10m_v_component_of_wind": "10 m V wind",
+        "100m_u_component_of_wind": "100 m U wind",
+        "100m_v_component_of_wind": "100 m V wind",
+        "u_component_of_wind_": "U wind ",
+        "v_component_of_wind_": "V wind ",
+        "outgoing_longwave_radiation": "Outgoing longwave radiation",
+        "sea_surface_temperature": "Sea-surface temperature",
+        "significant_wave_height": "Significant wave height",
+        "mean_wave_direction": "Mean wave direction",
+        "mean_wave_period": "Mean wave period",
+        "soil_moisture_0_7cm": "Soil moisture 0-7 cm",
+    }
+    for prefix, replacement in replacements.items():
+        if name.startswith(prefix):
+            suffix = name[len(prefix) :]
+            return f"{replacement}{suffix}" if suffix.isdigit() else replacement
+    return name.replace("_", " ").title()
+
+
+def _safe_percentile(values: np.ndarray, percentile: float, default: float = 1.0) -> float:
+    finite = np.asarray(values, dtype=np.float64)
+    finite = finite[np.isfinite(finite)]
+    if finite.size == 0:
+        return float(default)
+    result = float(np.percentile(finite, percentile))
+    return result if np.isfinite(result) else float(default)
+
+
+def _absolute_limits(values: np.ndarray) -> Tuple[float, float]:
+    lo, hi = _safe_percentile(values, 2.0, -1.0), _safe_percentile(values, 98.0, 1.0)
+    if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
+        center = float(np.nanmean(values)) if np.isfinite(values).any() else 0.0
+        span = max(abs(center) * 0.05, 1.0)
+        return center - span, center + span
+    return lo, hi
+
+
+def _error_limit(values: np.ndarray) -> float:
+    limit = _safe_percentile(np.abs(values), 98.0, 1.0)
+    if not np.isfinite(limit) or limit <= 1.0e-8:
+        limit = max(_safe_percentile(np.abs(values), 100.0, 1.0), 1.0)
+    return float(limit)
+
+
+def _field(field: np.ndarray, lookup: Mapping[str, int], name: str) -> np.ndarray | None:
+    index = lookup.get(name)
+    return None if index is None else np.asarray(field[index], dtype=np.float32)
+
+
+def _diagnostics(channels: Sequence[str], prediction: np.ndarray, truth: np.ndarray) -> List[Dict[str, Any]]:
+    """Create fields with meteorological meaning from a channel tensor."""
+    lookup = {name: index for index, name in enumerate(channels)}
+    diagnostics: List[Dict[str, Any]] = []
+
+    temperature_name = "2m_temperature" if "2m_temperature" in lookup else "temperature_500"
+    pred, true = _field(prediction, lookup, temperature_name), _field(truth, lookup, temperature_name)
+    if pred is not None and true is not None:
+        diagnostics.append({"key": "temperature", "title": _pretty_name(temperature_name), "unit": "data units", "prediction": pred, "truth": true, "cmap": "coolwarm"})
+
+    pred, true = _field(prediction, lookup, "geopotential_500"), _field(truth, lookup, "geopotential_500")
+    if pred is not None and true is not None:
+        diagnostics.append({"key": "geopotential", "title": "500 hPa geopotential", "unit": "data units", "prediction": pred, "truth": true, "cmap": "viridis"})
+
+    wind_pairs = (
+        ("10m_u_component_of_wind", "10m_v_component_of_wind", "10 m wind speed"),
+        ("u_component_of_wind_850", "v_component_of_wind_850", "850 hPa wind speed"),
+    )
+    for u_name, v_name, title in wind_pairs:
+        up, vp = _field(prediction, lookup, u_name), _field(prediction, lookup, v_name)
+        ut, vt = _field(truth, lookup, u_name), _field(truth, lookup, v_name)
+        if all(value is not None for value in (up, vp, ut, vt)):
+            diagnostics.append({"key": "wind_speed", "title": title, "unit": "data units", "prediction": np.hypot(up, vp), "truth": np.hypot(ut, vt), "cmap": "magma", "prediction_uv": (up, vp), "truth_uv": (ut, vt)})
+            break
+
+    pred, true = _field(prediction, lookup, "total_precipitation"), _field(truth, lookup, "total_precipitation")
+    if pred is not None and true is not None:
+        diagnostics.append({"key": "precipitation", "title": "Total precipitation", "unit": "data units", "prediction": pred, "truth": true, "cmap": "YlGnBu"})
+
+    if len(diagnostics) < 2:
+        used = {item["title"] for item in diagnostics}
+        for index, name in enumerate(channels):
+            if _pretty_name(name) in used:
+                continue
+            diagnostics.append({"key": f"channel_{index}", "title": _pretty_name(name), "unit": "data units", "prediction": prediction[index], "truth": truth[index], "cmap": "coolwarm"})
+            if len(diagnostics) >= 3:
+                break
+    return diagnostics
+
+
+def _draw_grid(
+    ax: Any,
+    field: np.ndarray,
+    title: str,
+    cmap: str,
+    unit: str,
+    error: bool = False,
+    overlay: Tuple[np.ndarray, np.ndarray] | None = None,
+    limits: Tuple[float, float] | None = None,
+) -> None:
+    """Draw a regular global latitude/longitude panel without cartopy."""
+    height, width = field.shape
+    longitude = np.linspace(0.0, 360.0, width, endpoint=False)
+    latitude = np.linspace(-90.0, 90.0, height)
+    if error:
+        limit = _error_limit(field) if limits is None else max(abs(float(limits[0])), abs(float(limits[1])))
+        norm = TwoSlopeNorm(vmin=-limit, vcenter=0.0, vmax=limit)
+        image = ax.imshow(field, extent=(0, 360, -90, 90), origin="lower", aspect="auto", cmap="RdBu_r", norm=norm, interpolation="nearest")
+    else:
+        lo, hi = _absolute_limits(field) if limits is None else limits
+        image = ax.imshow(field, extent=(0, 360, -90, 90), origin="lower", aspect="auto", cmap=cmap, norm=Normalize(vmin=lo, vmax=hi), interpolation="nearest")
+    ax.set_title(title, fontsize=10, pad=5)
+    ax.set_xlim(0, 360)
+    ax.set_ylim(-90, 90)
+    ax.set_xticks((0, 60, 120, 180, 240, 300, 360))
+    ax.set_yticks((-60, -30, 0, 30, 60))
+    ax.set_xlabel("Longitude (deg)", fontsize=8)
+    ax.set_ylabel("Latitude (deg)", fontsize=8)
+    ax.tick_params(labelsize=8)
+    ax.grid(color="white", alpha=0.35, linewidth=0.45)
+    ax.axhline(0, color="black", alpha=0.28, linewidth=0.55)
+    ax.axvline(180, color="black", alpha=0.28, linewidth=0.55)
+    if overlay is not None and not error:
+        u, v = overlay
+        step_y, step_x = max(1, height // 12), max(1, width // 24)
+        speed = np.hypot(u, v)
+        scale = max(_safe_percentile(speed, 95, 1.0) * 12, 1.0)
+        quiver = ax.quiver(longitude[::step_x], latitude[::step_y], u[::step_y, ::step_x], v[::step_y, ::step_x], color="black", alpha=0.68, scale=scale, width=0.0022, headwidth=3.5, minlength=0.1)
+        ax.quiverkey(quiver, 0.86, -0.16, _safe_percentile(speed, 75, 1.0), "75th pct", labelpos="E", coordinates="axes", fontproperties={"size": 7})
+    colorbar = ax.figure.colorbar(image, ax=ax, fraction=0.045, pad=0.025)
+    colorbar.ax.tick_params(labelsize=7)
+    colorbar.set_label(("Error (" + unit + ")") if error else unit, fontsize=8)
+
+
+def _plot_overview(path: Path, source: str, record: Mapping[str, Any], channels: Sequence[str], prediction: np.ndarray, truth: np.ndarray, dpi: int) -> None:
+    diagnostics = _diagnostics(channels, prediction, truth)
+    lead = int(record["lead"])
+    time_step = int(record.get("time_step", 6))
+    lead_hours = (lead + 1) * time_step
+    fig, axes = plt.subplots(len(diagnostics), 3, figsize=(16.5, max(4.0 * len(diagnostics), 8.0)), squeeze=False, facecolor="#f5f7fa")
+    for row, diagnostic in enumerate(diagnostics):
+        pred, true = diagnostic["prediction"], diagnostic["truth"]
+        error = pred - true
+        limits = _absolute_limits(np.concatenate((pred.reshape(-1), true.reshape(-1))))
+        _draw_grid(axes[row, 0], pred, f"{diagnostic['title']} | forecast", diagnostic["cmap"], diagnostic["unit"], overlay=diagnostic.get("prediction_uv"), limits=limits)
+        _draw_grid(axes[row, 1], true, f"{diagnostic['title']} | verifying field", diagnostic["cmap"], diagnostic["unit"], overlay=diagnostic.get("truth_uv"), limits=limits)
+        _draw_grid(axes[row, 2], error, f"{diagnostic['title']} | forecast error", diagnostic["cmap"], diagnostic["unit"], error=True)
+        axes[row, 2].text(0.02, 0.03, f"RMSE {np.sqrt(np.nanmean(error * error)):.3g}   bias {np.nanmean(error):+.3g}", transform=axes[row, 2].transAxes, fontsize=8, color="black", bbox={"facecolor": "white", "alpha": 0.82, "edgecolor": "none", "pad": 2.5})
+    valid = str(record["valid_timestamp"])
+    analysis_time = datetime.strptime(source, "%Y%m%d%H") - timedelta(hours=time_step)
+    analysis = analysis_time.strftime("%Y%m%d%H")
+    fig.suptitle(f"FengWu-W2S global forecast diagnostics\nInit {analysis[:8]} {analysis[8:]} UTC  |  Valid {valid[:8]} {valid[8:]} UTC  |  F{lead_hours:03d}", fontsize=15, fontweight="bold", y=0.995)
+    fig.subplots_adjust(top=0.93, left=0.04, right=0.98, bottom=0.045, hspace=0.56, wspace=0.20)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
+    plt.close(fig)
+
+
+def _plot_lead_summary(path: Path, lead_metrics: Sequence[Mapping[str, Any]], dpi: int) -> None:
+    if not lead_metrics:
+        return
+    leads = np.asarray([item["lead_hours"] for item in lead_metrics], dtype=float)
+    nrmse = np.asarray([item["mean_normalized_rmse"] for item in lead_metrics], dtype=float)
+    acc = np.asarray([item["mean_acc"] for item in lead_metrics], dtype=float)
+    fig, axes = plt.subplots(1, 2, figsize=(14, 4.8), facecolor="#f5f7fa")
+    for ax in axes:
+        ax.set_facecolor("white")
+        ax.grid(alpha=0.25, linewidth=0.7)
+        ax.set_xlabel("Forecast lead (hours)")
+    axes[0].plot(leads, nrmse, color="#b43f3f", marker="o", linewidth=2.2, markersize=5)
+    axes[0].fill_between(leads, 0, nrmse, color="#e8a1a1", alpha=0.28)
+    axes[0].set_title("Normalized error growth", fontweight="bold")
+    axes[0].set_ylabel("Mean RMSE / channel standard deviation")
+    axes[0].set_ylim(bottom=0)
+    axes[1].plot(leads, acc, color="#20639b", marker="o", linewidth=2.2, markersize=5)
+    axes[1].axhline(0, color="#555", linewidth=0.8)
+    axes[1].set_title("Anomaly correlation by lead", fontweight="bold")
+    axes[1].set_ylabel("Mean ACC")
+    axes[1].set_ylim(-1, 1)
+    fig.suptitle("FengWu-W2S forecast skill across available leads", fontsize=14, fontweight="bold")
+    fig.tight_layout(rect=(0, 0, 1, 0.93))
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
+    plt.close(fig)
+
+
+def _plot_channel_summary(path: Path, channels: Sequence[str], rmse: np.ndarray, normalized_rmse: np.ndarray, acc: np.ndarray, dpi: int) -> None:
+    count = min(12, len(channels))
+    ranking = np.argsort(normalized_rmse)[::-1][:count]
+    labels = [_pretty_name(channels[index]) for index in ranking][::-1]
+    values, acc_values = normalized_rmse[ranking][::-1], acc[ranking][::-1]
+    fig, axes = plt.subplots(1, 2, figsize=(15, max(5.5, count * 0.46)), facecolor="#f5f7fa")
+    for ax in axes:
+        ax.set_facecolor("white")
+        ax.grid(axis="x", alpha=0.25)
+    axes[0].barh(labels, values, color="#d46a6a", alpha=0.9)
+    axes[0].set_title("Largest normalized RMSE", fontweight="bold")
+    axes[0].set_xlabel("RMSE / channel standard deviation")
+    axes[1].barh(labels, acc_values, color="#4b88b8", alpha=0.9)
+    axes[1].set_title("ACC for the same channels", fontweight="bold")
+    axes[1].set_xlabel("Anomaly correlation coefficient")
+    axes[1].set_xlim(-1, 1)
+    fig.suptitle("Channel-level forecast diagnostics", fontsize=14, fontweight="bold")
+    fig.tight_layout(rect=(0, 0, 1, 0.94))
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
+    plt.close(fig)
+
+
+def _plot_training_history(path: Path, train: np.ndarray, valid: np.ndarray, dpi: int) -> None:
+    if train.size == 0 and valid.size == 0:
+        return
+    fig, ax = plt.subplots(figsize=(9.5, 4.8), facecolor="#f5f7fa")
+    ax.set_facecolor("white")
+    if train.size:
+        ax.plot(np.arange(1, train.size + 1), train, marker="o", color="#2166ac", linewidth=2, label="Training")
+    if valid.size:
+        ax.plot(np.arange(1, valid.size + 1), valid, marker="o", color="#b2182b", linewidth=2, label="Validation")
+    ax.set_xlabel("Epoch")
+    ax.set_ylabel("Reported non-negative probability loss")
+    ax.set_title("FengWu-W2S optimization history", fontweight="bold")
+    ax.grid(alpha=0.28)
+    ax.legend(frameon=False)
+    fig.tight_layout()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(path, dpi=dpi, facecolor=fig.get_facecolor())
+    plt.close(fig)
+
+
+def _truth_for_timestamp(data_dir: Path, timestamp: str, indices: np.ndarray, time_step: int) -> np.ndarray:
+    dt = datetime.strptime(timestamp, "%Y%m%d%H")
+    index = int((dt - datetime(dt.year, 1, 1)).total_seconds() // (3600 * time_step))
+    path = data_dir / "data" / f"{dt.year}.h5"
+    with h5py.File(path, "r") as source:
+        fields = source["fields"]
+        if index < 0 or index >= fields.shape[0]:
+            raise IndexError(f"Timestamp {timestamp} maps to index {index}, outside {path}")
+        return np.asarray(fields[index, indices, :, :], dtype=np.float32)
+
+
+def _load_prediction(output_dir: Path, record: Mapping[str, Any]) -> np.ndarray:
+    path = Path(record["path"])
+    if not path.is_absolute():
+        path = output_dir / path
+    if not path.exists():
+        raise FileNotFoundError(f"Forecast field listed in index.json does not exist: {path}")
+    field = np.load(path).astype(np.float32)
+    if field.ndim != 3:
+        raise ValueError(f"Forecast field must have [channel, lat, lon] shape, got {field.shape}")
+    return field
+
+
+def _selected_records(records: Sequence[Mapping[str, Any]], plot_count: int, plot_leads: int) -> List[Mapping[str, Any]]:
+    grouped: Dict[str, List[Mapping[str, Any]]] = defaultdict(list)
+    for record in records:
+        grouped[str(record["source_timestamp"])].append(record)
+    selected: List[Mapping[str, Any]] = []
+    for source in sorted(grouped)[: max(0, int(plot_count))]:
+        source_records = sorted(grouped[source], key=lambda item: int(item["lead"]))
+        number = max(1, min(int(plot_leads), len(source_records)))
+        positions = np.linspace(0, len(source_records) - 1, number, dtype=int)
+        selected.extend(source_records[int(position)] for position in np.unique(positions))
+    return selected
+
+
+def main() -> None:
+    args = parse_args()
+    with Path(args.config).open(encoding="utf-8") as source:
+        config = yaml.safe_load(source)
+    data_cfg = config["data"]
+    data_dir = resolve_data_dir(args.data_dir or data_cfg["data_dir"], PROJECT_ROOT)
+    output_dir = Path(args.output_dir).expanduser()
+    if not output_dir.is_absolute():
+        output_dir = (PROJECT_ROOT / output_dir).resolve()
+    index_path = output_dir / "index.json"
+    if not index_path.exists():
+        raise FileNotFoundError(f"Forecast index not found: {index_path}; run scripts/inference.py first")
+    records = json.loads(index_path.read_text(encoding="utf-8"))
+    if not isinstance(records, list):
+        raise ValueError("Forecast index must contain a list")
+    channels = list(data_cfg["channels"])
+    metadata = read_metadata(data_dir, channels)
+    indices, time_step = metadata["indices"], int(metadata["time_step"])
+    means = np.asarray(metadata["means"], dtype=np.float64).reshape(-1)
+    stds = np.asarray(metadata["stds"], dtype=np.float64).reshape(-1)
+    sum_squared = np.zeros(len(channels), dtype=np.float64)
+    sum_pred_anom_sq = np.zeros(len(channels), dtype=np.float64)
+    sum_truth_anom_sq = np.zeros(len(channels), dtype=np.float64)
+    sum_cross = np.zeros(len(channels), dtype=np.float64)
+    lead_sum_squared = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
+    lead_sum_pred_anom_sq = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
+    lead_sum_truth_anom_sq = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
+    lead_sum_cross = defaultdict(lambda: np.zeros(len(channels), dtype=np.float64))
+    lead_counts: Dict[int, int] = defaultdict(int)
+    count = 0
+    for record in records:
+        prediction = _load_prediction(output_dir, record)
+        truth = _truth_for_timestamp(data_dir, str(record["valid_timestamp"]), indices, time_step)
+        if prediction.shape != truth.shape:
+            raise ValueError(f"Prediction/truth shape mismatch: {prediction.shape} vs {truth.shape}")
+        if not np.isfinite(prediction).all() or not np.isfinite(truth).all():
+            raise ValueError(f"Non-finite forecast or truth field for {record}")
+        error = prediction - truth
+        pred_anom, truth_anom = prediction - means[:, None, None], truth - means[:, None, None]
+        pixels = int(error.shape[1] * error.shape[2])
+        sum_squared += np.sum(error * error, axis=(1, 2))
+        sum_pred_anom_sq += np.sum(pred_anom * pred_anom, axis=(1, 2))
+        sum_truth_anom_sq += np.sum(truth_anom * truth_anom, axis=(1, 2))
+        sum_cross += np.sum(pred_anom * truth_anom, axis=(1, 2))
+        lead = int(record["lead"])
+        lead_sum_squared[lead] += np.sum(error * error, axis=(1, 2))
+        lead_sum_pred_anom_sq[lead] += np.sum(pred_anom * pred_anom, axis=(1, 2))
+        lead_sum_truth_anom_sq[lead] += np.sum(truth_anom * truth_anom, axis=(1, 2))
+        lead_sum_cross[lead] += np.sum(pred_anom * truth_anom, axis=(1, 2))
+        lead_counts[lead] += pixels
+        count += pixels
+    rmse = np.sqrt(sum_squared / max(1, count))
+    acc = sum_cross / (np.sqrt(sum_pred_anom_sq * sum_truth_anom_sq) + 1.0e-8)
+    normalized_rmse = rmse / np.maximum(stds, 1.0e-8)
+    lead_metrics: List[Dict[str, Any]] = []
+    pixels_per_field = int(metadata["shape"][-2] * metadata["shape"][-1])
+    for lead in sorted(lead_counts):
+        denominator = max(1, lead_counts[lead])
+        lead_rmse = np.sqrt(lead_sum_squared[lead] / denominator)
+        lead_acc = lead_sum_cross[lead] / (np.sqrt(lead_sum_pred_anom_sq[lead] * lead_sum_truth_anom_sq[lead]) + 1.0e-8)
+        lead_metrics.append({"lead": lead, "lead_hours": (lead + 1) * time_step, "mean_rmse": float(np.mean(lead_rmse)), "mean_normalized_rmse": float(np.mean(lead_rmse / np.maximum(stds, 1.0e-8))), "mean_acc": float(np.mean(lead_acc)), "records": int(lead_counts[lead] // max(1, pixels_per_field))})
+    result_root, plots_dir = output_dir.parent, output_dir.parent / "plots"
+    plots_dir.mkdir(parents=True, exist_ok=True)
+    for old_plot in plots_dir.glob("*.png"):
+        old_plot.unlink()
+    metrics = {"count": len(records), "channels": channels, "rmse": rmse.tolist(), "normalized_rmse": normalized_rmse.tolist(), "acc": acc.tolist(), "mean_rmse": float(np.mean(rmse)), "mean_normalized_rmse": float(np.mean(normalized_rmse)), "mean_acc": float(np.mean(acc)), "lead_metrics": lead_metrics, "visualization": {"domain": "meteorology", "coordinate_convention": "regular grid, longitude 0-360 degrees, latitude -90 to 90 degrees", "figures": ["plots/*_overview.png", "plots/lead_skill.png", "plots/channel_skill.png", "plots/training_history.png"]}}
+    (result_root / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+    np.save(result_root / "rmse.npy", rmse.astype(np.float32))
+    np.save(result_root / "acc.npy", acc.astype(np.float32))
+    np.save(result_root / "normalized_rmse.npy", normalized_rmse.astype(np.float32))
+    for record in _selected_records(records, args.plot_count, args.plot_leads):
+        prediction = _load_prediction(output_dir, record)
+        truth = _truth_for_timestamp(data_dir, str(record["valid_timestamp"]), indices, time_step)
+        source, lead = str(record["source_timestamp"]), int(record["lead"])
+        _plot_overview(plots_dir / f"{source}_lead{lead:03d}_overview.png", source, {**record, "time_step": time_step}, channels, prediction, truth, max(80, int(args.dpi)))
+    _plot_lead_summary(plots_dir / "lead_skill.png", lead_metrics, max(80, int(args.dpi)))
+    _plot_channel_summary(plots_dir / "channel_skill.png", channels, rmse, normalized_rmse, acc, max(80, int(args.dpi)))
+    checkpoint_dir = PROJECT_ROOT / "data" / "checkpoints"
+    train_path, valid_path = checkpoint_dir / "trloss.npy", checkpoint_dir / "valoss.npy"
+    if train_path.exists() and valid_path.exists():
+        _plot_training_history(plots_dir / "training_history.png", np.load(train_path), np.load(valid_path), max(80, int(args.dpi)))
+    print(json.dumps(metrics, indent=2))
+    print(f"Saved scientific diagnostics under {plots_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+"""Train FengWu-W2S on ERA5-style HDF5 windows (single or torchrun DDP)."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import sys
+from pathlib import Path
+from typing import Any, Dict
+
+import numpy as np
+import torch
+import yaml
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.fengwu_w2s import FengWuW2S
+from scripts.data_loader import make_dataloader, resolve_data_dir
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--data-dir")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--max-epoch", type=int)
+    parser.add_argument("--rollout-steps", type=int)
+    parser.add_argument("--batch-size", type=int)
+    parser.add_argument("--num-workers", type=int)
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), help="Execution device; auto prefers DCU/GPU")
+    parser.add_argument("--finetune", action="store_true", help="Resume with the lower finetune learning rate")
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args()
+
+
+def _device(config: Dict[str, Any], local_rank: int, override: str | None = None) -> torch.device:
+    requested = str(override or config.get("runtime", {}).get("device", "auto")).lower()
+    if requested not in {"auto", "cpu", "cuda"}:
+        raise ValueError(f"unsupported device {requested!r}; expected auto, cpu, or cuda")
+    cuda_available = bool(torch.cuda.is_available())
+    cuda_count = int(torch.cuda.device_count()) if cuda_available else 0
+    if requested == "cpu":
+        return torch.device("cpu")
+    if not cuda_available or cuda_count < 1:
+        if requested == "cuda":
+            raise RuntimeError(
+                "CUDA/DCU was explicitly requested but torch.cuda is unavailable; "
+                "load the accelerator module before starting training"
+            )
+        return torch.device("cpu")
+    index = local_rank
+    if index >= cuda_count:
+        raise RuntimeError(f"local rank {index} exceeds visible accelerator count {cuda_count}")
+    torch.cuda.set_device(index)
+    return torch.device("cuda", index)
+
+
+def _distributed_context() -> tuple[bool, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    distributed = world_size > 1
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    return distributed, rank, local_rank
+
+
+def _model_from_config(config: Dict[str, Any]) -> FengWuW2S:
+    model_cfg = config["model"]
+    data_cfg = config["data"]
+    groups = {name: tuple(values) for name, values in data_cfg["groups"].items()}
+    return FengWuW2S(
+        in_channels=len(data_cfg["channels"]),
+        group_indices=groups,
+        hidden_channels=int(model_cfg.get("hidden_channels", 32)),
+        latent_channels=int(model_cfg.get("latent_channels", 64)),
+        patch_size=int(model_cfg.get("patch_size", 4)),
+        num_blocks=int(model_cfg.get("num_blocks", 2)),
+        perturbation_scale=float(model_cfg.get("perturbation_scale", 0.01)),
+    )
+
+
+def _load_checkpoint(path: Path, model: nn.Module, optimizer: torch.optim.Optimizer | None, device: torch.device) -> Dict[str, Any]:
+    state = torch.load(path, map_location=device, weights_only=False)
+    model.load_state_dict(state["model_state_dict"])
+    if optimizer is not None and "optimizer_state_dict" in state:
+        optimizer.load_state_dict(state["optimizer_state_dict"])
+    return state
+
+
+def _epoch(
+    model: nn.Module,
+    loader,
+    device: torch.device,
+    optimizer: torch.optim.Optimizer | None,
+    rollout_steps: int,
+    distributed: bool,
+    probability_loss: bool,
+    kl_weight: float,
+    perturbation_enabled: bool,
+) -> float:
+    training = optimizer is not None
+    model.train(training)
+    total = 0.0
+    batches = 0
+    for inputs, targets, _ in loader:
+        with torch.set_grad_enabled(training):
+            state = inputs[:, -1].to(device=device, dtype=torch.float32, non_blocking=True)
+            targets = targets.to(device=device, dtype=torch.float32, non_blocking=True)
+            if optimizer is not None:
+                optimizer.zero_grad(set_to_none=True)
+            loss = state.new_zeros(())
+            steps = max(1, min(int(rollout_steps), targets.shape[1]))
+            for step in range(steps):
+                prediction, aux = model(
+                    state,
+                    stochastic=training and perturbation_enabled,
+                    return_aux=True,
+                )
+                error = prediction - targets[:, step]
+                if probability_loss:
+                    min_log_std = -6.0
+                    log_std = aux["log_std"].clamp(min_log_std, 2.0)
+                    # The offset is constant and therefore preserves Gaussian
+                    # NLL gradients while making its logged value non-negative.
+                    nll = (
+                        0.5 * torch.exp(-2.0 * log_std) * error.square()
+                        + log_std
+                        - min_log_std
+                    )
+                    loss = loss + torch.mean(nll)
+                    if training and perturbation_enabled:
+                        loss = loss + float(kl_weight) * aux["kl"]
+                else:
+                    loss = loss + 0.8 * torch.mean(torch.abs(error)) + 0.2 * torch.mean(error.square())
+                state = prediction
+            loss = loss / steps
+            loss_value = float(loss.detach().cpu())
+            if not math.isfinite(loss_value) or loss_value < -1.0e-7:
+                raise RuntimeError(f"invalid loss {loss_value} on {device}")
+            if optimizer is not None:
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                optimizer.step()
+        total += loss_value
+        batches += 1
+    if batches == 0:
+        return float("nan")
+    value = torch.tensor(total / batches, device=device)
+    if distributed:
+        torch.distributed.all_reduce(value, op=torch.distributed.ReduceOp.SUM)
+        value = value / torch.distributed.get_world_size()
+    return float(value.cpu())
+
+
+def main() -> None:
+    args = parse_args()
+    with Path(args.config).open(encoding="utf-8") as source:
+        config = yaml.safe_load(source)
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    distributed, rank, local_rank = _distributed_context()
+    device = _device(config, local_rank, args.device)
+    if distributed:
+        backend = "nccl" if device.type == "cuda" else "gloo"
+        torch.distributed.init_process_group(backend=backend, init_method="env://")
+    model_cfg = config["model"]
+    data_cfg = config["data"]
+    data_dir = resolve_data_dir(args.data_dir or data_cfg["data_dir"], PROJECT_ROOT)
+    input_steps = int(model_cfg.get("input_steps", 2))
+    rollout_steps = int(args.rollout_steps or model_cfg.get("rollout_steps", 1))
+    batch_size = int(args.batch_size or data_cfg.get("dataloader", {}).get("batch_size", 1))
+    num_workers = int(args.num_workers if args.num_workers is not None else data_cfg.get("dataloader", {}).get("num_workers", 0))
+    train_loader, train_sampler = make_dataloader(
+        data_dir, data_cfg["train_years"], data_cfg["channels"], input_steps, rollout_steps,
+        batch_size, num_workers, distributed, True, data_cfg.get("dataloader", {}).get("pin_memory", False),
+    )
+    val_loader, val_sampler = make_dataloader(
+        data_dir, data_cfg["val_years"], data_cfg["channels"], input_steps, rollout_steps,
+        batch_size, num_workers, distributed, False, data_cfg.get("dataloader", {}).get("pin_memory", False),
+    )
+    model = _model_from_config(config).to(device)
+    if not args.finetune:
+        for parameter in model.perturbation.parameters():
+            parameter.requires_grad_(False)
+    parameter_device = next(model.parameters()).device
+    if parameter_device != device:
+        raise RuntimeError(f"model parameters are on {parameter_device}, expected {device}")
+    learning_rate = float(model_cfg.get("learning_rate", 2e-4))
+    if args.finetune:
+        learning_rate = float(model_cfg.get("finetune_learning_rate", learning_rate * 0.1))
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=learning_rate,
+        betas=(
+            float(model_cfg.get("adam_beta1", 0.9)),
+            float(model_cfg.get("adam_beta2", 0.999)),
+        ),
+        weight_decay=float(model_cfg.get("weight_decay", 1e-5)),
+    )
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=3)
+    checkpoint_dir = resolve_data_dir(model_cfg.get("checkpoint_dir", "./data/checkpoints"), PROJECT_ROOT)
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_path = Path(args.checkpoint) if args.checkpoint else checkpoint_dir / "model_bak.pth"
+    if not checkpoint_path.is_absolute():
+        checkpoint_path = (PROJECT_ROOT / checkpoint_path).resolve()
+    start_epoch = 0
+    best_valid = float("inf")
+    train_losses: list[float] = []
+    valid_losses: list[float] = []
+    if args.finetune:
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(
+                f"Finetune checkpoint not found: {checkpoint_path}; run base training first"
+            )
+        state = _load_checkpoint(checkpoint_path, model, optimizer, device)
+        start_epoch = int(state.get("epoch", -1)) + 1
+        best_valid = float(state.get("best_valid_loss", best_valid))
+        # Loading AdamW state also restores its previous learning rate; apply
+        # the explicit finetune rate after restoration.
+        for group in optimizer.param_groups:
+            group["lr"] = learning_rate
+    if distributed:
+        model = DDP(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    requested_epochs = int(args.max_epoch or (model_cfg.get("finetune_epoch") if args.finetune else model_cfg.get("max_epoch", 1)))
+    # ``--max-epoch`` is the number of epochs for this invocation.  A
+    # finetune run therefore always performs the requested extra epochs after
+    # the checkpoint epoch rather than silently skipping when it is already 1.
+    max_epoch = start_epoch + requested_epochs if args.finetune else requested_epochs
+    patience = int(model_cfg.get("patience", 10))
+    probability_loss = bool(model_cfg.get("probability_loss", True))
+    kl_weight = float(model_cfg.get("kl_weight", 1.0e-4))
+    stale = 0
+    if rank == 0:
+        cuda_state = {
+            "available": bool(torch.cuda.is_available()),
+            "count": int(torch.cuda.device_count()) if torch.cuda.is_available() else 0,
+            "name": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
+            "allocated_bytes": torch.cuda.memory_allocated(device) if device.type == "cuda" else 0,
+        }
+        print(
+            f"FengWu-W2S training on {device}; model_device={parameter_device}; "
+            f"cuda={cuda_state}; samples={len(train_loader.dataset)}; rollout={rollout_steps}",
+            flush=True,
+        )
+        print(f"Parameters: {sum(parameter.numel() for parameter in model.parameters()):,}", flush=True)
+    for epoch in range(start_epoch, max_epoch):
+        if distributed:
+            train_sampler.set_epoch(epoch)
+            val_sampler.set_epoch(epoch)
+        train_loss = _epoch(
+            model, train_loader, device, optimizer, rollout_steps, distributed,
+            probability_loss, kl_weight, args.finetune,
+        )
+        valid_loss = _epoch(
+            model, val_loader, device, None, rollout_steps, distributed,
+            probability_loss, kl_weight, False,
+        )
+        scheduler.step(valid_loss)
+        train_losses.append(train_loss)
+        valid_losses.append(valid_loss)
+        improved = valid_loss < best_valid
+        if improved:
+            best_valid = valid_loss
+            stale = 0
+        else:
+            stale += 1
+        if rank == 0:
+            model_to_save = model.module if hasattr(model, "module") else model
+            state = {
+                "model_state_dict": model_to_save.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "scheduler_state_dict": scheduler.state_dict(),
+                "epoch": epoch,
+                "best_valid_loss": best_valid,
+                "channels": data_cfg["channels"],
+                "group_indices": data_cfg["groups"],
+                "config": config,
+            }
+            if improved or epoch == max_epoch - 1:
+                torch.save(state, checkpoint_dir / "model_bak.pth")
+            np.save(checkpoint_dir / "trloss.npy", np.asarray(train_losses, dtype=np.float32))
+            np.save(checkpoint_dir / "valoss.npy", np.asarray(valid_losses, dtype=np.float32))
+            print(
+                f"Epoch {epoch + 1}/{max_epoch}: train={train_loss:.6f} valid={valid_loss:.6f}",
+                flush=True,
+            )
+        if stale > patience:
+            break
+    if distributed:
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+    if rank == 0:
+        (checkpoint_dir / "training_summary.json").write_text(
+            json.dumps({"best_valid_loss": best_valid, "epochs": len(train_losses), "device": str(device)}, indent=2),
+            encoding="utf-8",
+        )
+        print(f"Saved checkpoint to {checkpoint_dir / 'model_bak.pth'}")
+
+
+if __name__ == "__main__":
+    main()