Upload folder using huggingface_hub

README.md

@@ -0,0 +1,150 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecast
+- Medium-Range Weather Forecast
+- ERA5
+- FuXi
+tasks: []
+datasets:
+- OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">FuXi_v21</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+FuXi 2.1 is a global deterministic machine-learning weather forecast model developed by Fudan University in collaboration with the Shanghai Artificial Intelligence Laboratory (SAIS). Its theoretical basis remains the original FuXi paper.
+
+Paper: FuXi: A cascade machine learning forecasting system for 15-day global weather forecast
+
+https://arxiv.org/abs/2306.12873
+
+# Model Description
+
+The model addresses the excessive smoothing often observed in AI weather forecasts. It aims to produce clearer and more detailed forecast fields, improving the detection of extreme events such as heavy precipitation and strong winds without degrading conventional metrics such as root mean square error (RMSE).
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global weather forecast training | Train FuXi v2.1 with C85 ERA5 data in HDF5 format. |
+| Local quick validation | Use synthetic HDF5 data to check data loading, training, inference, and visualization of inference results. |
+| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Run distributed data-parallel training 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/FuXi_v21 --local-dir ./FuXi_v21
+cd FuXi_v21
+```
+
+### 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 training entry point uses the OneScience `ERA5Dataset`. The data root is specified by `paths.data_root` in `conf/config.yaml`. The OneScience community provides a data slice for interface validation and training:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+### Generate Synthetic Data
+
+Real HDF5 files must contain a `fields` dataset, C85 variable attributes, six-hour intervals, and normalization statistics. When real data is unavailable, generate protocol-compatible synthetic files:
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 scripts/train.py
+```
+
+Training checkpoints are saved to `data/checkpoint/model_bak.pth` by default, and metrics are saved to `output/training/metrics.json`.
+
+### Training Weights
+
+This repository provides weights trained on ERA5 reanalysis data in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference results are saved to `output/inference/forecast.nc` by default.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The default output is `figures/fuxi21_t2m.png`.
+
+# 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
+
+- This project is an unofficial forward-graph reproduction of FuXi v2.1. It does not represent official weights or training recipes released by Fudan University.
+- This adapted repository is distributed under Apache License 2.0 metadata. ERA5 data, OneScience, and the upstream FuXi implementation remain subject to their respective official licenses and terms of use.

conf/config.yaml

@@ -0,0 +1,101 @@
+protocol: non_official_protocol
+seed: 42
+
+paths:
+  data_root: data
+  static_root: data/static
+  checkpoint: data/checkpoint/model_bak.pth
+  training_metrics: output/training/metrics.json
+  inference_output: output/inference/forecast.nc
+  visualization_output: figures/fuxi21_t2m.png
+
+data:
+  input_steps: 2
+  output_steps: 1
+  time_step_hours: 6
+  grid_size: [721, 1440]
+  crop_size: null
+  num_workers: 0
+  splits:
+    train:
+      years: [2021,2022]
+      time_steps: 10
+    val:
+      years: [2023]
+      time_steps: 10
+    inference:
+      years: [2024]
+      time_steps: 10
+  hdf5:
+    chunks: [1, 1, 64, 64]
+    compression: gzip
+    compression_level: 1
+
+variables:
+  pressure_levels: [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
+  pressure: [z, t, u, v, q]
+  surface: [msl, t2m, d2m, sst, ws10m, ws100m, u10m, v10m, u100m, v100m, lcc, mcc, hcc, tcc, ssr, ssrd, fdir, ttr, tcw, tp]
+  diagnostic: [ssr, ssrd, fdir, ttr, tp]
+  mapping:
+    pressure_order: variable_major
+    pressure_name: "{variable}{level}"
+    surface_name: "{variable}"
+    source_dataset: fields
+    units: source_native
+    transform: identity
+
+model:
+  profile: full
+  static_fields_file: data/static/static_fields.npy
+  channel_mask_file: data/static/channel_mask.npy
+  activation_checkpointing: true
+  profiles:
+    full:
+      grid_size: [721, 1440]
+      patch_size: 6
+      embed_dim: 1536
+      depth: 30
+      num_heads: 24
+      mlp_dim: 4096
+      window_size: 20
+    smoke:
+      grid_size: [13, 12]
+      patch_size: 6
+      embed_dim: 32
+      depth: 2
+      num_heads: 4
+      mlp_dim: 64
+      window_size: 2
+
+training:
+  device: auto
+  epochs: 5
+  batch_size: 1
+  precision: fp32
+  distributed_strategy: ddp
+  fsdp_sharding: full_shard
+  gradient_accumulation_steps: 1
+  optimizer: AdamW
+  learning_rate: 0.0001
+  min_learning_rate: 0.000001
+  weight_decay: 0.01
+  scheduler: CosineAnnealingLR
+  gradient_clip_norm: 1.0
+  channel_weights: null
+  checkpoint_mode: scratch
+  load_checkpoint: null
+  save_checkpoint: data/checkpoint/model_bak.pth
+
+inference:
+  device: auto
+  checkpoint: data/checkpoint/model_bak.pth
+  split: inference
+  steps: 1
+  zero_diagnostic_feedback: false
+  output_file: output/inference/forecast.nc
+
+visualization:
+  input_file: output/inference/forecast.nc
+  output_file: figures/fuxi21_t2m.png
+  channel: t2m
+  cmap: coolwarm

config.json

@@ -0,0 +1,57 @@
+{
+  "model_name": "FuXi v2.1",
+  "model_type": "fuxi_v21",
+  "architectures": [
+    "FuXi21"
+  ],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "global-medium-range-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/FuXi21.py",
+    "scope": "trainable forward-graph reconstruction with full-resolution and smoke profiles"
+  },
+  "architecture": {
+    "family": "windowed Transformer with rotary attention and PixelShuffle decoder",
+    "grid_shape": [
+      721,
+      1440
+    ],
+    "input_channels": 85,
+    "output_channels": 85,
+    "static_channels": 6,
+    "patch_size": 6,
+    "embedding_size": 1536,
+    "depth": 30,
+    "attention_heads": 24,
+    "mlp_size": 4096,
+    "window_size": 20,
+    "activation": "SiLU/GELU"
+  },
+  "data": {
+    "dataset": "ERA5",
+    "spatial_resolution_degrees": 0.25,
+    "time_step_hours": 6,
+    "input_steps": 2,
+    "pressure_levels_hpa": [
+      50,
+      100,
+      150,
+      200,
+      250,
+      300,
+      400,
+      500,
+      600,
+      700,
+      850,
+      925,
+      1000
+    ],
+    "protocol": "non_official_protocol"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/FuXi21.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "weather_forecasting",
+  "model": "FuXi_v21",
+  "input_format": "BTCHW",
+  "protocol": "non_official_protocol",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/FuXi21.py

@@ -0,0 +1,298 @@
+"""Trainable FuXi 2.1 forward-graph reconstruction."""
+
+from __future__ import annotations
+
+import math
+from torch.utils.checkpoint import checkpoint as activation_checkpoint
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+DIAGNOSTIC_INDICES = (79, 80, 81, 82, 84)
+
+
+class UnbiasedNorm(nn.Module):
+    """Layer normalization matching the PT2 graph's unbiased variance."""
+
+    def __init__(self, dim: int, conditioned: bool = False, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.conditioned = conditioned
+        self.scale_shift = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim)) if conditioned else None
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor:
+        variance, mean = torch.var_mean(x, dim=-1, correction=1, keepdim=True)
+        x = (x - mean) * torch.rsqrt(variance + self.eps) * self.weight
+        if self.scale_shift is not None:
+            if condition is None:
+                raise ValueError("condition is required by conditioned normalization")
+            scale, shift = self.scale_shift(condition).chunk(2, dim=-1)
+            x = x * (1 + scale[:, None, :]) + shift[:, None, :]
+        return x
+
+
+def _rope_frequencies(height: int, width: int, head_dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+    if head_dim % 2:
+        raise ValueError("head_dim must be even for rotary embeddings")
+    y, x = torch.meshgrid(torch.arange(height), torch.arange(width), indexing="ij")
+    positions = (y * width + x).flatten().float()
+    frequencies = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
+    angles = positions[:, None] * frequencies[None, :]
+    return angles.cos(), angles.sin()
+
+
+def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    even, odd = x[..., 0::2], x[..., 1::2]
+    cos = cos[None, :, None, :].to(dtype=x.dtype, device=x.device)
+    sin = sin[None, :, None, :].to(dtype=x.dtype, device=x.device)
+    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
+
+
+def _window_partition(x: torch.Tensor, window: int) -> torch.Tensor:
+    batch, height, width, channels = x.shape
+    return x.view(batch, height // window, window, width // window, window, channels).permute(0, 1, 3, 2, 4, 5).reshape(-1, window * window, channels)
+
+
+def _window_reverse(x: torch.Tensor, batch: int, height: int, width: int, window: int) -> torch.Tensor:
+    return x.view(batch, height // window, width // window, window, window, -1).permute(0, 1, 3, 2, 4, 5).reshape(batch, height, width, -1)
+
+
+def _shift_mask(height: int, width: int, window: int) -> torch.Tensor:
+    shift = window // 2
+    labels = torch.zeros(1, height, width, 1)
+    h_slices = (slice(0, -window), slice(-window, -shift), slice(-shift, None))
+    w_slices = (slice(0, -window), slice(-window, -shift), slice(-shift, None))
+    index = 0
+    for h_slice in h_slices:
+        for w_slice in w_slices:
+            labels[:, h_slice, w_slice] = index
+            index += 1
+    labels = _window_partition(labels, window).squeeze(-1)
+    mask = labels[:, None, :] - labels[:, :, None]
+    return mask.masked_fill(mask != 0, float("-inf")).masked_fill(mask == 0, 0.0)
+
+
+class HeadGatedWindowAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, window: int, grid_size: tuple[int, int], shifted: bool) -> None:
+        super().__init__()
+        if dim % num_heads:
+            raise ValueError("dim must be divisible by num_heads")
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.window = window
+        self.grid_size = grid_size
+        self.shifted = shifted
+        self.wq = nn.Linear(dim, num_heads * (self.head_dim + 1), bias=False)
+        self.wk = nn.Linear(dim, dim, bias=False)
+        self.wv = nn.Linear(dim, dim, bias=False)
+        self.wo = nn.Linear(dim, dim, bias=False)
+        cos, sin = _rope_frequencies(*grid_size, self.head_dim)
+        self.register_buffer("freqs_cos", cos, persistent=False)
+        self.register_buffer("freqs_sin", sin, persistent=False)
+        self.register_buffer("attention_mask", _shift_mask(*grid_size, window) if shifted else None, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, tokens, channels = x.shape
+        height, width = self.grid_size
+        qg = self.wq(x).view(batch, tokens, self.num_heads, self.head_dim + 1)
+        q, gate = qg[..., : self.head_dim], qg[..., -1:].sigmoid()
+        k = self.wk(x).view(batch, tokens, self.num_heads, self.head_dim)
+        v = self.wv(x).view(batch, tokens, self.num_heads, self.head_dim)
+        q = _apply_rope(q, self.freqs_cos, self.freqs_sin).reshape(batch, height, width, channels)
+        k = _apply_rope(k, self.freqs_cos, self.freqs_sin).reshape(batch, height, width, channels)
+        v = v.reshape(batch, height, width, channels)
+        gate = gate.reshape(batch, height, width, self.num_heads, 1)
+
+        if self.shifted:
+            shift = self.window // 2
+            q, k, v, gate = [torch.roll(item, shifts=(-shift, -shift), dims=(1, 2)) for item in (q, k, v, gate)]
+        q, k, v = [_window_partition(item, self.window).view(-1, self.window**2, self.num_heads, self.head_dim).transpose(1, 2) for item in (q, k, v)]
+        gate = _window_partition(gate.flatten(-2), self.window).view(-1, self.window**2, self.num_heads, 1).transpose(1, 2)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if self.attention_mask is not None:
+            windows = self.attention_mask.shape[0]
+            scores = scores.view(batch, windows, self.num_heads, self.window**2, self.window**2)
+            scores = scores + self.attention_mask[None, :, None].to(scores)
+            scores = scores.flatten(0, 1)
+        output = torch.matmul(scores.softmax(dim=-1), v) * gate
+        output = output.transpose(1, 2).reshape(-1, self.window**2, channels)
+        output = _window_reverse(output, batch, height, width, self.window)
+        if self.shifted:
+            output = torch.roll(output, shifts=(self.window // 2, self.window // 2), dims=(1, 2))
+        return self.wo(output.reshape(batch, tokens, channels))
+
+
+class FuXi21Block(nn.Module):
+    def __init__(self, dim: int, mlp_dim: int, num_heads: int, window: int, grid_size: tuple[int, int], shifted: bool) -> None:
+        super().__init__()
+        self.adaln = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.norm1 = UnbiasedNorm(dim)
+        self.attn = HeadGatedWindowAttention(dim, num_heads, window, grid_size, shifted)
+        self.norm2 = UnbiasedNorm(dim)
+        self.w1 = nn.Linear(dim, mlp_dim, bias=False)
+        self.w2 = nn.Linear(mlp_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, mlp_dim, bias=False)
+
+    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
+        attn_scale, attn_shift, attn_gate, mlp_scale, mlp_shift, mlp_gate = self.adaln(condition).chunk(6, dim=-1)
+        normalized = self.norm1(x) * (1 + attn_scale[:, None]) + attn_shift[:, None]
+        x = x + attn_gate[:, None] * self.attn(normalized)
+        normalized = self.norm2(x) * (1 + mlp_scale[:, None]) + mlp_shift[:, None]
+        mlp = self.w2(F.silu(self.w1(normalized)) * self.w3(normalized))
+        return x + mlp_gate[:, None] * mlp
+
+
+class PixelShuffleHead(nn.Module):
+    def __init__(self, dim: int, output_channels: int) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(dim, 2 * dim, 3, padding=1)
+        self.conv2 = nn.Conv2d(dim // 2, output_channels * 9, 3, padding=1)
+
+    def forward(self, x: torch.Tensor, output_size: tuple[int, int]) -> torch.Tensor:
+        x = F.pad(x, (0, 0, 0, 1), mode="replicate")
+        x = F.gelu(F.pixel_shuffle(self.conv1(x), 2))
+        x = F.pixel_shuffle(self.conv2(x), 3)
+        return x[..., : output_size[0], : output_size[1]]
+
+
+class FuXi21(nn.Module):
+    """Randomly initialized, trainable reconstruction of the FuXi 2.1 PT2 forward graph.
+
+    The defaults reproduce the recovered architecture; reduced dimensions and grids
+    are intended for smoke tests.
+    """
+
+    def __init__(
+        self,
+        static_fields: torch.Tensor,
+        channel_mask: torch.Tensor,
+        grid_size: tuple[int, int] = (721, 1440),
+        embed_dim: int = 1536,
+        depth: int = 30,
+        num_heads: int = 24,
+        mlp_dim: int = 4096,
+        patch_size: int = 6,
+        window_size: int = 20,
+        activation_checkpointing: bool = False,
+    ) -> None:
+        super().__init__()
+        height, width = grid_size
+        token_grid = (height // patch_size, width // patch_size)
+        if patch_size != 6:
+            raise ValueError("The recovered PixelShuffle decoder requires patch_size=6")
+        if any(size % window_size for size in token_grid):
+            raise ValueError(f"token grid {token_grid} must be divisible by window_size={window_size}")
+        if static_fields.shape != (6, height, width):
+            raise ValueError(f"static_fields must have shape {(6, height, width)}, got {tuple(static_fields.shape)}")
+        if channel_mask.shape != (85, height, width):
+            raise ValueError(f"channel_mask must have shape {(85, height, width)}, got {tuple(channel_mask.shape)}")
+        if embed_dim % 4:
+            raise ValueError("embed_dim must be divisible by 4 for the PixelShuffle heads")
+
+        self.grid_size = grid_size
+        self.token_grid = token_grid
+        self.activation_checkpointing = activation_checkpointing
+        self.register_buffer("static_fields", static_fields.detach().float())
+        self.register_buffer("channel_mask", channel_mask.detach().float())
+        self.patch_embed = nn.Conv2d(170, embed_dim, patch_size, stride=patch_size)
+        self.patch_norm = UnbiasedNorm(embed_dim)
+        self.const_embed = nn.Conv2d(6, embed_dim, patch_size, stride=patch_size)
+        self.const_norm = UnbiasedNorm(embed_dim)
+        self.joint_embed_layer = nn.Sequential(nn.Linear(384, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim))
+        self.layers = nn.ModuleList(
+            FuXi21Block(embed_dim, mlp_dim, num_heads, window_size, token_grid, bool(index % 2))
+            for index in range(depth)
+        )
+        self.norm_layer = UnbiasedNorm(embed_dim, conditioned=True)
+        self.pressure_head = nn.ConvTranspose2d(embed_dim, 65, 9, stride=6, padding=1)
+        self.surface_head = PixelShuffleHead(embed_dim, 15)
+        self.derived_head = PixelShuffleHead(embed_dim, 5)
+        self.register_buffer("scatter_idx", torch.tensor([*range(79), 83, 79, 80, 81, 82, 84]), persistent=False)
+        self.reset_parameters()
+
+    @classmethod
+    def smoke(
+        cls,
+        static_fields: torch.Tensor | None = None,
+        channel_mask: torch.Tensor | None = None,
+    ) -> "FuXi21":
+        static_fields = torch.zeros(6, 13, 12) if static_fields is None else static_fields
+        channel_mask = torch.ones(85, 13, 12) if channel_mask is None else channel_mask
+        return cls(
+            static_fields,
+            channel_mask,
+            grid_size=(13, 12),
+            embed_dim=32,
+            depth=2,
+            num_heads=4,
+            mlp_dim=64,
+            window_size=2,
+        )
+
+    def reset_parameters(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.trunc_normal_(module.weight, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, (nn.Conv2d, nn.ConvTranspose2d)):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+        for block in self.layers:
+            nn.init.zeros_(block.adaln[-1].weight)
+            nn.init.zeros_(block.adaln[-1].bias)
+        nn.init.zeros_(self.norm_layer.scale_shift[-1].weight)
+        nn.init.zeros_(self.norm_layer.scale_shift[-1].bias)
+        for head in (self.pressure_head, self.surface_head.conv2, self.derived_head.conv2):
+            nn.init.trunc_normal_(head.weight, std=1e-3)
+
+    @staticmethod
+    def _time_embedding(value: torch.Tensor, periodic: bool) -> torch.Tensor:
+        frequency = torch.arange(64, device=value.device, dtype=value.dtype)
+        if periodic:
+            angles = 2 * math.pi * value.reshape(-1, 1) * frequency
+        else:
+            angles = value.reshape(-1, 1) / (10000 ** (frequency / 64))
+        return torch.cat((angles.sin(), angles.cos()), dim=-1)
+
+    def forward(self, state: torch.Tensor, step: torch.Tensor, hour: torch.Tensor, doy: torch.Tensor) -> torch.Tensor:
+        expected = (2, 85, *self.grid_size)
+        if tuple(state.shape[1:]) != expected:
+            raise ValueError(f"state must have shape (B, {expected}), got {tuple(state.shape)}")
+        state = torch.nan_to_num(state)
+        state = state.clone()
+        state[:, :, DIAGNOSTIC_INDICES] = 0
+        state = state * self.channel_mask
+        previous = state[:, -1]
+        batch = state.shape[0]
+
+        x = self.patch_embed(state.reshape(batch, 170, *self.grid_size)).flatten(2).transpose(1, 2)
+        x = self.patch_norm(x)
+        const = self.const_embed(self.static_fields[None].expand(batch, -1, -1, -1)).flatten(2).transpose(1, 2)
+        x = x + self.const_norm(const)
+        time_features = torch.cat(
+            (self._time_embedding(step, False), self._time_embedding(hour, True), self._time_embedding(doy, True)), dim=-1
+        )
+        condition = self.joint_embed_layer(time_features)
+        for layer in self.layers:
+            if self.training and self.activation_checkpointing:
+                x = activation_checkpoint(layer, x, condition, use_reentrant=False)
+            else:
+                x = layer(x, condition)
+        x = self.norm_layer(x, condition).transpose(1, 2).reshape(batch, -1, *self.token_grid)
+
+        pressure = self.pressure_head(x)[..., : self.grid_size[0], : self.grid_size[1]]
+        surface = self.surface_head(x, self.grid_size)
+        derived = self.derived_head(x, self.grid_size)
+        grouped = torch.cat((pressure, surface, derived), dim=1)
+        prediction = torch.empty_like(grouped)
+        prediction[:, self.scatter_idx] = grouped
+        return torch.stack((previous, prediction), dim=1)
+
+    @property
+    def trainable(self) -> bool:
+        return True

scripts/build_static.py

@@ -0,0 +1,53 @@
+"""Build project-defined FuXi static fields and C85 validity mask from scratch."""
+
+from __future__ import annotations
+
+import argparse
+
+import numpy as np
+
+from common import load_config, resolve_path
+from variables import c85_from_config
+
+
+def build_static_fields(height: int, width: int) -> np.ndarray:
+    latitude = np.deg2rad(np.linspace(90.0, -90.0, height, dtype=np.float32))[:, None]
+    longitude = np.deg2rad(np.arange(width, dtype=np.float32) * (360.0 / width))[None, :]
+    lat = np.broadcast_to(latitude, (height, width))
+    lon = np.broadcast_to(longitude, (height, width))
+    geopotential = np.zeros_like(lat)
+    land_sea = np.ones_like(lat)
+    return np.stack((geopotential, land_sea, np.cos(lat), np.sin(lat), np.cos(lon), np.sin(lon)))
+
+
+def build_channel_mask(channels: list[str], height: int, width: int) -> np.ndarray:
+    # The reconstruction has no external missing-channel metadata; all C85 fields are valid.
+    return np.ones((len(channels), height, width), dtype=np.float32)
+
+
+def ensure_static_resources(cfg: dict, force: bool = False) -> tuple:
+    channels, _ = c85_from_config(cfg)
+    profile = cfg["model"]["profiles"]["full"]
+    height, width = profile["grid_size"]
+    static_path = resolve_path(cfg["model"]["static_fields_file"], cfg)
+    mask_path = resolve_path(cfg["model"]["channel_mask_file"], cfg)
+    static_path.parent.mkdir(parents=True, exist_ok=True)
+    mask_path.parent.mkdir(parents=True, exist_ok=True)
+    if force or not static_path.is_file():
+        np.save(static_path, build_static_fields(height, width))
+    if force or not mask_path.is_file():
+        np.save(mask_path, build_channel_mask(channels, height, width))
+    return static_path, mask_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="conf/config.yaml")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    static_path, mask_path = ensure_static_resources(cfg, force=True)
+    print(f"Saved from-scratch static resources to {static_path} and {mask_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/common.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+
+def load_config(path: str | Path) -> dict:
+    path = Path(path)
+    with path.open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    config["_config_dir"] = str(path.resolve().parent)
+    return config
+
+
+def resolve_path(value: str, config: dict) -> Path:
+    path = Path(value).expanduser()
+    if path.is_absolute():
+        return path
+    return (Path(config["_config_dir"]).parent / path).resolve()

scripts/fake_data.py

@@ -0,0 +1,77 @@
+"""Create an ERA5-style HDF5 dataset consumable by OneScience ERA5Dataset."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import h5py
+import numpy as np
+
+from common import load_config, resolve_path
+from variables import c85_from_config
+
+
+def create_year(
+    path: Path,
+    year: int,
+    steps: int,
+    height: int,
+    width: int,
+    channels: list[str],
+    time_step_hours: int,
+    hdf5_config: dict,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with h5py.File(path, "w") as handle:
+        fields = handle.create_dataset(
+            "fields",
+            shape=(steps, len(channels), height, width),
+            dtype="float32",
+            chunks=tuple(min(size, limit) for size, limit in zip((steps, len(channels), height, width), hdf5_config["chunks"])),
+            fillvalue=0.0,
+            compression=hdf5_config["compression"],
+            compression_opts=hdf5_config["compression_level"],
+        )
+        fields.attrs["variables"] = np.asarray(channels, dtype=h5py.string_dtype("utf-8"))
+        fields.attrs["time_step"] = time_step_hours
+        fields.attrs["year"] = year
+        lat = np.linspace(90.0, -90.0, height, dtype=np.float32)[:, None]
+        lon = np.arange(width, dtype=np.float32)[None, :] * (360.0 / width)
+        t2m_index = channels.index("t2m")
+        for step in range(steps):
+            fields[step, t2m_index] = (
+                np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(lon + step * 15.0))
+            )
+        handle.create_dataset("global_means", data=np.zeros((1, len(channels), 1, 1), dtype=np.float32))
+        handle.create_dataset("global_stds", data=np.ones((1, len(channels), 1, 1), dtype=np.float32))
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="conf/config.yaml")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    data_cfg = cfg["data"]
+    channels, _ = c85_from_config(cfg)
+    root = resolve_path(cfg["paths"]["data_root"], cfg)
+    generated_years = set()
+    for split, split_cfg in data_cfg["splits"].items():
+        for year in split_cfg["years"]:
+            if year in generated_years:
+                raise ValueError(f"Year {year} is assigned to more than one data split")
+            generated_years.add(year)
+            create_year(
+                root / "data" / f"{year}.h5",
+                year,
+                split_cfg["time_steps"],
+                *data_cfg["grid_size"],
+                channels,
+                data_cfg["time_step_hours"],
+                data_cfg["hdf5"],
+            )
+    print(f"Created ERA5-compatible yearly datasets at {root / 'data'}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,146 @@
+"""Autoregressive inference using a project-produced FuXi 2.1 checkpoint."""
+
+from __future__ import annotations
+
+import argparse
+from datetime import datetime, timedelta
+
+import numpy as np
+import torch
+import xarray as xr
+from onescience.datapipes.climate.era5 import ERA5Dataset
+
+from common import load_config, resolve_path
+from model.FuXi21 import FuXi21
+from variables import c85_from_config
+
+
+CHECKPOINT_FORMAT = "fuxi21_reconstructed_checkpoint_v1"
+
+
+def select_device(requested: str) -> torch.device:
+    if requested not in {"auto", "cpu", "cuda"}:
+        raise ValueError("inference.device must be auto, cpu, or cuda")
+    if requested == "cuda" or (requested == "auto" and torch.cuda.is_available()):
+        if not torch.cuda.is_available():
+            raise RuntimeError("inference.device=cuda, but no CUDA/HIP device is available")
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def load_array(path_value: str | None, cfg: dict, shape: tuple[int, ...], name: str) -> torch.Tensor:
+    if path_value is None:
+        raise ValueError(f"model.{name}_file is required outside the smoke profile")
+    path = resolve_path(path_value, cfg)
+    value = torch.from_numpy(np.load(path)).float()
+    if tuple(value.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(value.shape)}")
+    return value
+
+
+def build_model(cfg: dict) -> FuXi21:
+    model_cfg = cfg["model"]
+    profile_name = model_cfg["profile"]
+    profile = model_cfg["profiles"][profile_name]
+    height, width = profile["grid_size"]
+    if profile_name == "smoke":
+        static_fields = torch.zeros(6, height, width)
+        channel_mask = torch.ones(85, height, width)
+    else:
+        static_fields = load_array(model_cfg["static_fields_file"], cfg, (6, height, width), "static_fields")
+        channel_mask = load_array(model_cfg["channel_mask_file"], cfg, (85, height, width), "channel_mask")
+    return FuXi21(
+        static_fields,
+        channel_mask,
+        activation_checkpointing=False,
+        **profile,
+    )
+
+
+def temporal_features(valid_time: datetime, step: int, device: torch.device) -> tuple[torch.Tensor, ...]:
+    return (
+        torch.tensor([step], device=device, dtype=torch.float32),
+        torch.tensor([(valid_time.hour * 60 + valid_time.minute) / 1440], device=device),
+        torch.tensor([min(365, valid_time.timetuple().tm_yday) / 365], device=device),
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None)
+    parser.add_argument("--preflight-only", action="store_true")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    if cfg.get("protocol") != "non_official_protocol":
+        raise ValueError("Inference config must declare protocol: non_official_protocol")
+
+    infer_cfg = cfg["inference"]
+    channels, diagnostics = c85_from_config(cfg)
+    checkpoint_path = resolve_path(infer_cfg["checkpoint"], cfg)
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"Project checkpoint not found: {checkpoint_path}")
+    device = select_device(args.device or infer_cfg["device"])
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
+    if checkpoint.get("format") != CHECKPOINT_FORMAT:
+        raise ValueError(f"Checkpoint must use format {CHECKPOINT_FORMAT}")
+    if checkpoint.get("protocol") != "non_official_protocol":
+        raise ValueError("Checkpoint protocol must be non_official_protocol")
+    if checkpoint.get("model_profile") != cfg["model"]["profile"]:
+        raise ValueError("Checkpoint model profile does not match the configured model profile")
+    if args.preflight_only:
+        print(f"checkpoint={checkpoint_path}, profile={checkpoint['model_profile']}, device={device}")
+        return
+
+    model = build_model(cfg).to(device)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    split = infer_cfg["split"]
+    split_cfg = cfg["data"]["splits"][split]
+    dataset = ERA5Dataset(
+        dataset_dir=str(resolve_path(cfg["paths"]["data_root"], cfg)),
+        used_years=split_cfg["years"],
+        used_variables=channels,
+        input_steps=cfg["data"]["input_steps"],
+        output_steps=cfg["data"]["output_steps"],
+        normalize=True,
+    )
+    state, _, _, _, time_index = dataset[0]
+    crop_size = cfg["data"]["crop_size"]
+    if crop_size is not None:
+        state = state[..., : crop_size[0], : crop_size[1]]
+    state = state.unsqueeze(0).to(device)
+    valid_time = datetime.strptime(time_index[-1], "%Y%m%d%H")
+    interval = timedelta(hours=cfg["data"]["time_step_hours"])
+    diagnostic_indices = [channels.index(name) for name in diagnostics]
+    forecasts = []
+    valid_times = []
+    for step in range(infer_cfg["steps"]):
+        with torch.inference_mode():
+            state = model(state, *temporal_features(valid_time, step, device))
+        forecasts.append(state[:, -1].float().cpu().numpy()[0])
+        valid_times.append(np.datetime64(valid_time))
+        if infer_cfg["zero_diagnostic_feedback"]:
+            state[:, -1, diagnostic_indices] = 0
+        valid_time += interval
+
+    height, width = forecasts[0].shape[-2:]
+    output_path = resolve_path(infer_cfg["output_file"], cfg)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    xr.DataArray(
+        np.stack(forecasts),
+        dims=("time", "channel", "lat", "lon"),
+        coords={
+            "time": valid_times,
+            "channel": channels,
+            "lat": np.linspace(90, -90, height),
+            "lon": np.arange(width) * (360 / width),
+        },
+        attrs={"checkpoint_format": CHECKPOINT_FORMAT, "protocol": cfg["protocol"]},
+        name="forecast",
+    ).to_netcdf(output_path)
+    print(f"Saved {len(forecasts)} forecast step(s) to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import xarray as xr
+
+from common import load_config, resolve_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--input", default=None)
+    parser.add_argument("--channel", default=None)
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    viz = cfg["visualization"]
+    source = Path(args.input) if args.input else resolve_path(viz["input_file"], cfg)
+    channel = args.channel or viz["channel"]
+    data = xr.open_dataarray(source).sel(channel=channel)
+    if "time" in data.dims:
+        data = data.isel(time=-1)
+    output = resolve_path(viz["output_file"], cfg)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    fig, ax = plt.subplots(figsize=(12, 5), constrained_layout=True)
+    image = ax.pcolormesh(data.lon, data.lat, data, shading="auto", cmap=viz["cmap"])
+    valid_time = str(data.time.values) if "time" in data.coords else ""
+    ax.set(title=f"FuXi 2.1 {channel} | {valid_time}", xlabel="Longitude", ylabel="Latitude")
+    fig.colorbar(image, ax=ax, label=channel)
+    fig.savefig(output, dpi=160)
+    plt.close(fig)
+    print(f"Saved visualization to {output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,362 @@
+"""Non-official from-scratch training baseline for the FuXi 2.1 reconstruction."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+from contextlib import nullcontext
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from onescience.datapipes.climate.era5 import ERA5Dataset
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from common import load_config, resolve_path
+from build_static import ensure_static_resources
+from model.FuXi21 import FuXi21
+from variables import c85_from_config
+
+try:
+    from torch.distributed.fsdp import (
+        FullStateDictConfig,
+        FullyShardedDataParallel,
+        ShardingStrategy,
+        StateDictType,
+    )
+except ImportError:  # pragma: no cover - depends on the installed torch build
+    FullyShardedDataParallel = None
+    FullStateDictConfig = None
+    ShardingStrategy = None
+    StateDictType = None
+
+
+def setup_distributed(device: torch.device) -> tuple[bool, int, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    distributed = world_size > 1
+    if distributed:
+        dist.init_process_group(backend="nccl" if device.type == "cuda" else "gloo")
+    rank = dist.get_rank() if distributed else 0
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    return distributed, rank, local_rank, world_size
+
+
+def select_device(requested: str, local_rank: int) -> torch.device:
+    if requested not in {"auto", "cpu", "cuda"}:
+        raise ValueError("training.device must be auto, cpu, or cuda")
+    use_accelerator = requested == "cuda" or (requested == "auto" and torch.cuda.is_available())
+    if use_accelerator:
+        if not torch.cuda.is_available():
+            raise RuntimeError("training.device=cuda, but no CUDA/HIP device is available")
+        if local_rank >= torch.cuda.device_count():
+            raise RuntimeError(
+                f"LOCAL_RANK={local_rank} exceeds {torch.cuda.device_count()} visible CUDA/HIP device(s); "
+                "use one process per visible device or --device cpu for DDP logic testing"
+            )
+        torch.cuda.set_device(local_rank)
+        return torch.device("cuda", local_rank)
+    return torch.device("cpu")
+
+
+def seed_everything(seed: int, rank: int) -> None:
+    seed += rank
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def load_array(path_value: str | None, cfg: dict, expected_shape: tuple[int, ...], name: str) -> torch.Tensor:
+    if path_value is None:
+        raise ValueError(f"model.{name}_file is required outside the smoke profile")
+    path = resolve_path(path_value, cfg)
+    if not path.is_file():
+        raise FileNotFoundError(f"{name} file not found: {path}")
+    value = torch.from_numpy(np.load(path)).float()
+    if tuple(value.shape) != expected_shape:
+        raise ValueError(f"{name} must have shape {expected_shape}, got {tuple(value.shape)}")
+    return value
+
+
+def build_model(cfg: dict) -> FuXi21:
+    model_cfg = cfg["model"]
+    profile_name = model_cfg["profile"]
+    profile = model_cfg["profiles"][profile_name]
+    height, width = profile["grid_size"]
+    if profile_name == "smoke":
+        static_fields = torch.zeros(6, height, width)
+        channel_mask = torch.ones(85, height, width)
+    else:
+        static_fields = load_array(model_cfg["static_fields_file"], cfg, (6, height, width), "static_fields")
+        channel_mask = load_array(model_cfg["channel_mask_file"], cfg, (85, height, width), "channel_mask")
+    return FuXi21(
+        static_fields=static_fields,
+        channel_mask=channel_mask,
+        activation_checkpointing=cfg["model"].get("activation_checkpointing", False),
+        **profile,
+    )
+
+
+def wrap_distributed_model(model: FuXi21, cfg: dict, device: torch.device, local_rank: int):
+    strategy = cfg["training"].get("distributed_strategy", "ddp")
+    if strategy == "ddp":
+        return DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    if strategy != "fsdp":
+        raise ValueError("training.distributed_strategy must be ddp or fsdp")
+    if FullyShardedDataParallel is None:
+        raise RuntimeError("This PyTorch build does not provide torch.distributed.fsdp")
+    if device.type == "cpu":
+        raise RuntimeError("FSDP training requires an accelerator device")
+    sharding = cfg["training"].get("fsdp_sharding", "full_shard")
+    if sharding != "full_shard":
+        raise ValueError("Only fsdp_sharding=full_shard is currently supported")
+    return FullyShardedDataParallel(
+        model,
+        device_id=device,
+        sharding_strategy=ShardingStrategy.FULL_SHARD,
+        use_orig_params=True,
+    )
+
+
+def is_fsdp(model) -> bool:
+    return FullyShardedDataParallel is not None and isinstance(model, FullyShardedDataParallel)
+
+
+def unwrap_model(model):
+    if isinstance(model, DistributedDataParallel) or is_fsdp(model):
+        return model.module
+    return model
+
+
+def build_loader(cfg: dict, split: str, distributed: bool, train: bool):
+    data_cfg = cfg["data"]
+    channels, _ = c85_from_config(cfg)
+    split_cfg = data_cfg["splits"][split]
+    dataset_root = resolve_path(cfg["paths"]["data_root"], cfg)
+    dataset = ERA5Dataset(
+        dataset_dir=str(dataset_root),
+        used_years=split_cfg["years"],
+        used_variables=channels,
+        input_steps=data_cfg["input_steps"],
+        output_steps=data_cfg["output_steps"],
+        normalize=True,
+    )
+    sampler = DistributedSampler(dataset, shuffle=train) if distributed else None
+    loader = DataLoader(
+        dataset,
+        batch_size=cfg["training"]["batch_size"],
+        shuffle=train and sampler is None,
+        sampler=sampler,
+        num_workers=data_cfg["num_workers"],
+        pin_memory=torch.cuda.is_available(),
+        drop_last=train and distributed,
+    )
+    return loader, sampler
+
+
+def crop_batch(inputs: torch.Tensor, targets: torch.Tensor, crop_size: list[int] | None):
+    if crop_size is None:
+        return inputs, targets
+    height, width = crop_size
+    return inputs[..., :height, :width], targets[..., :height, :width]
+
+
+def temporal_features(time_index, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    target_times = time_index[-1]
+    if isinstance(target_times, str):
+        target_times = [target_times]
+    parsed = [datetime.strptime(value, "%Y%m%d%H") for value in target_times]
+    step = torch.zeros(len(parsed), device=device)
+    hour = torch.tensor([(value.hour * 60 + value.minute) / 1440 for value in parsed], device=device)
+    doy = torch.tensor([min(365, value.timetuple().tm_yday) / 365 for value in parsed], device=device)
+    return step, hour, doy
+
+
+def weighted_mse(prediction: torch.Tensor, target: torch.Tensor, channel_weights: torch.Tensor) -> torch.Tensor:
+    latitude = torch.linspace(90, -90, prediction.shape[-2], device=prediction.device, dtype=prediction.dtype)
+    area = latitude.deg2rad().cos().clamp_min(0)
+    area = area / area.mean()
+    weights = channel_weights.to(prediction).view(1, -1, 1, 1) * area.view(1, 1, -1, 1)
+    return ((prediction - target).square() * weights).mean()
+
+
+def autocast_context(device: torch.device, precision: str):
+    if precision == "fp32":
+        return nullcontext()
+    if precision != "bf16":
+        raise ValueError("training.precision must be fp32 or bf16")
+    return torch.autocast(device_type=device.type, dtype=torch.bfloat16)
+
+
+def run_epoch(model, loader, optimizer, device, cfg, channel_weights, train: bool) -> float:
+    model.train(train)
+    total_loss = torch.zeros((), device=device)
+    total_samples = torch.zeros((), device=device)
+    clip_norm = cfg["training"]["gradient_clip_norm"]
+    accumulation_steps = max(1, int(cfg["training"].get("gradient_accumulation_steps", 1)))
+    if train:
+        optimizer.zero_grad(set_to_none=True)
+    for batch_index, (inputs, targets, _, _, time_index) in enumerate(loader):
+        inputs, targets = crop_batch(inputs, targets, cfg["data"]["crop_size"])
+        inputs = inputs.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+        if targets.ndim == 5:
+            targets = targets[:, 0]
+        temporal = temporal_features(time_index, device)
+        with torch.set_grad_enabled(train), autocast_context(device, cfg["training"]["precision"]):
+            prediction = model(inputs, *temporal)[:, -1]
+            loss = weighted_mse(prediction, targets, channel_weights)
+        if train:
+            (loss / accumulation_steps).backward()
+            if (batch_index + 1) % accumulation_steps == 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+        batch_size = inputs.shape[0]
+        total_loss += loss.detach() * batch_size
+        total_samples += batch_size
+    if train and len(loader) % accumulation_steps:
+        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
+        optimizer.step()
+        optimizer.zero_grad(set_to_none=True)
+    if dist.is_initialized():
+        dist.all_reduce(total_loss)
+        dist.all_reduce(total_samples)
+    return (total_loss / total_samples).item()
+
+
+def checkpoint_state(model, optimizer, scheduler, epoch: int, cfg: dict) -> dict:
+    if is_fsdp(model):
+        state_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
+        with FullyShardedDataParallel.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_config):
+            model_state = model.state_dict()
+        optimizer_state = FullyShardedDataParallel.optim_state_dict(model, optimizer)
+    else:
+        model_state = unwrap_model(model).state_dict()
+        optimizer_state = optimizer.state_dict()
+    return {
+        "format": "fuxi21_reconstructed_checkpoint_v1",
+        "protocol": cfg["protocol"],
+        "model": model_state,
+        "optimizer": optimizer_state,
+        "scheduler": scheduler.state_dict(),
+        "epoch": epoch,
+        "model_profile": cfg["model"]["profile"],
+        "model_config": cfg["model"]["profiles"][cfg["model"]["profile"]],
+        "distributed_strategy": cfg["training"].get("distributed_strategy", "ddp"),
+    }
+
+
+def load_checkpoint(path: Path, mode: str, model, optimizer, scheduler, device: torch.device) -> int:
+    state = torch.load(path, map_location=device, weights_only=True)
+    if state.get("format") != "fuxi21_reconstructed_checkpoint_v1":
+        raise ValueError("Only checkpoints produced by this reconstruction can be loaded")
+    if is_fsdp(model):
+        state_config = FullStateDictConfig(offload_to_cpu=False, rank0_only=False)
+        with FullyShardedDataParallel.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_config):
+            model.load_state_dict(state["model"])
+    else:
+        unwrap_model(model).load_state_dict(state["model"])
+    if mode == "resume":
+        if is_fsdp(model):
+            optimizer_state = FullyShardedDataParallel.optim_state_dict_to_load(
+                model, optimizer, state["optimizer"]
+            )
+            optimizer.load_state_dict(optimizer_state)
+        else:
+            optimizer.load_state_dict(state["optimizer"])
+        scheduler.load_state_dict(state["scheduler"])
+        return int(state["epoch"]) + 1
+    if mode == "initialize":
+        return 0
+    raise ValueError("checkpoint_mode must be scratch, initialize, or resume")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None)
+    parser.add_argument("--dry-run", action="store_true", help="Run one train and validation batch without saving")
+    args = parser.parse_args()
+    cfg = load_config(args.config)
+    if cfg.get("protocol") != "non_official_protocol":
+        raise ValueError("Training config must declare protocol: non_official_protocol")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    device = select_device(args.device or cfg["training"]["device"], local_rank)
+    distributed, rank, local_rank, _ = setup_distributed(device)
+    seed_everything(cfg["seed"], rank)
+    if cfg["model"]["profile"] == "full":
+        if rank == 0:
+            ensure_static_resources(cfg)
+        if distributed:
+            dist.barrier()
+    model = build_model(cfg).to(device)
+    if distributed:
+        model = wrap_distributed_model(model, cfg, device, local_rank)
+    train_years = set(cfg["data"]["splits"]["train"]["years"])
+    val_years = set(cfg["data"]["splits"]["val"]["years"])
+    if train_years & val_years:
+        raise ValueError("train and val years must be disjoint")
+    train_loader, train_sampler = build_loader(cfg, "train", distributed, True)
+    val_loader, _ = build_loader(cfg, "val", distributed, False)
+
+    train_cfg = cfg["training"]
+    if train_cfg["optimizer"] != "AdamW" or train_cfg["scheduler"] != "CosineAnnealingLR":
+        raise ValueError("This baseline supports optimizer=AdamW and scheduler=CosineAnnealingLR")
+    optimizer = torch.optim.AdamW(model.parameters(), lr=train_cfg["learning_rate"], weight_decay=train_cfg["weight_decay"])
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+        optimizer, T_max=train_cfg["epochs"], eta_min=train_cfg["min_learning_rate"]
+    )
+    channel_weights = train_cfg["channel_weights"]
+    channel_weights = torch.ones(85) if channel_weights is None else torch.tensor(channel_weights, dtype=torch.float32)
+    if channel_weights.shape != (85,) or torch.any(channel_weights <= 0):
+        raise ValueError("training.channel_weights must contain 85 positive values")
+    channel_weights = channel_weights / channel_weights.mean()
+
+    start_epoch = 0
+    checkpoint = train_cfg["load_checkpoint"]
+    mode = train_cfg["checkpoint_mode"]
+    if checkpoint is not None:
+        path = resolve_path(checkpoint, cfg)
+        start_epoch = load_checkpoint(path, mode, model, optimizer, scheduler, device)
+    elif mode != "scratch":
+        raise ValueError(f"checkpoint is required for checkpoint_mode={mode}")
+
+    checkpoint_path = resolve_path(train_cfg["save_checkpoint"], cfg)
+    metrics_path = resolve_path(cfg["paths"]["training_metrics"], cfg)
+    if rank == 0 and not args.dry_run:
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+        metrics_path.parent.mkdir(parents=True, exist_ok=True)
+    history = []
+    end_epoch = min(train_cfg["epochs"], start_epoch + 1) if args.dry_run else train_cfg["epochs"]
+    for epoch in range(start_epoch, end_epoch):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        train_loss = run_epoch(model, train_loader, optimizer, device, cfg, channel_weights, True)
+        val_loss = run_epoch(model, val_loader, optimizer, device, cfg, channel_weights, False)
+        scheduler.step()
+        record = {"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "learning_rate": scheduler.get_last_lr()[0]}
+        history.append(record)
+        state = None
+        if not args.dry_run:
+            # FSDP state-dict collection is collective; every rank must enter it.
+            state = checkpoint_state(model, optimizer, scheduler, epoch, cfg)
+        if rank == 0:
+            print(json.dumps(record))
+            if not args.dry_run:
+                torch.save(state, checkpoint_path)
+                metrics_path.write_text(json.dumps(history, indent=2) + "\n", encoding="utf-8")
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

scripts/variables.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+PRESSURE_LEVELS = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]
+PRESSURE_VARIABLES = ["z", "t", "u", "v", "q"]
+SURFACE_VARIABLES = [
+    "msl", "t2m", "d2m", "sst", "ws10m", "ws100m", "u10m", "v10m",
+    "u100m", "v100m", "lcc", "mcc", "hcc", "tcc", "ssr", "ssrd",
+    "fdir", "ttr", "tcw", "tp",
+]
+C85_CHANNEL_NAMES = [f"{name}{level}" for name in PRESSURE_VARIABLES for level in PRESSURE_LEVELS]
+C85_CHANNEL_NAMES += SURFACE_VARIABLES
+DIAGNOSTIC_CHANNELS = ["ssr", "ssrd", "fdir", "ttr", "tp"]
+
+assert len(C85_CHANNEL_NAMES) == 85
+
+
+def c85_from_config(config: dict) -> tuple[list[str], list[str]]:
+    """Build and validate the fixed C85 channel contract from configuration."""
+    variables = config["variables"]
+    mapping = variables["mapping"]
+    if mapping != {
+        "pressure_order": "variable_major",
+        "pressure_name": "{variable}{level}",
+        "surface_name": "{variable}",
+        "source_dataset": "fields",
+        "units": "source_native",
+        "transform": "identity",
+    }:
+        raise ValueError("variables.mapping must preserve the C85 source and ordering contract")
+    channels = [
+        mapping["pressure_name"].format(variable=name, level=level)
+        for name in variables["pressure"]
+        for level in variables["pressure_levels"]
+    ]
+    channels.extend(mapping["surface_name"].format(variable=name) for name in variables["surface"])
+    diagnostics = list(variables["diagnostic"])
+    if channels != C85_CHANNEL_NAMES:
+        raise ValueError("Configured variables do not match the required C85 channel order")
+    if diagnostics != DIAGNOSTIC_CHANNELS:
+        raise ValueError("Configured diagnostic variables do not match the required C85 contract")
+    return channels, diagnostics