Upload folder using huggingface_hub

README.md

@@ -0,0 +1,132 @@
+---
+frameworks: PyTorch
+tasks: []
+tags:
+  - OneScience
+  - Earth Science
+  - Weather Forecast
+  - ERA5
+  - Deterministic Forecast
+language:
+  - zh
+  - en
+license: apache-2.0
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">AIFS_Single_v1</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+AIFS Single v1.1 is the deterministic version of the AI-powered weather forecasting system developed by the European Centre for Medium-Range Weather Forecasts (ECMWF).
+
+Paper: AIFS — ECMWF's data-driven forecasting system, arXiv:2406.01465
+
+https://arxiv.org/abs/2406.01465
+
+# Model Description
+
+AIFS is built on a Graph Neural Network (GNN), pre-trained on ERA5 data and fine-tuned with NWP operational analysis data.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Weather Forecast Training | Train AIFS from scratch using ERA5 HDF5 data |
+| Local Quick Validation | Use synthetic data to verify data loading, model training & inference, and inference result visualization. |
+| ModelScope / OneCode Execution | Download as a standalone model package, install dependencies, and run scripts directly. |
+
+# 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 is recommended.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/AIFS_Single_v1 --local_dir ./AIFS_Single_v1
+cd AIFS_Single_v1
+```
+
+### Install the Runtime Environment
+
+**DCU**
+```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**
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data Introduction
+
+The OneScience community provides ERA5 data for training (due to file size limits, the current repository contains a slice of the full dataset). Users can download it with the command below and confirm that the data path in `conf/config.yaml` is set correctly:
+
+```bash
+modelscope download --dataset OneScience/ERA5 --local_dir ./data
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+Training weights are saved to `weights/model_bak.ckpt`, and the normalization file computed before training is saved to `weights/era5_stats.npz`.
+
+### Training Weights
+
+This repository provides weights trained on ERA5 reanalysis data in the `weights/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+The number of forecast steps is controlled by `test_lead_time` in `conf/config.yaml` (in hours; default 24 = 1 day).
+Inference results will be saved to the `output` directory.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Computes ACC / RMSE metrics and generates plots. Metrics are saved to `metrics/` and figures are saved to `plots/`.
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation & License
+
+- This repository is a reproduction of the original AIFS Single v1.1 paper.

conf/config.yaml

@@ -0,0 +1,298 @@
+# =============================================================================
+# AIFS v1.1 — Configuration
+# =============================================================================
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Data  (most frequently modified)
+# ═════════════════════════════════════════════════════════════════════════════
+data:
+  data_dir: "./fake_era5"       # ERA5 H5 root (contains data/{year}.h5)
+  # data_dir: "/work2/share/sugonhpcapp01/ERA5"
+
+  train_years:                      # training set
+    - 2005
+    - 2006
+
+  val_years:                        # validation set
+    - 2007
+
+  test_years:                       # test / inference set
+    - 2008
+
+  test_lead_time: 24                # hours (240 = 10 days, official default)
+
+  fake_data_nt: 60                   # timesteps per fake-data year
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Hardware
+# ═════════════════════════════════════════════════════════════════════════════
+hardware:
+  device: "dcu"
+  device_ids: "0"
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Checkpoint
+# ═════════════════════════════════════════════════════════════════════════════
+checkpoint:
+  pretrained: "./weights/model_bak.ckpt"   # official AIFS v1.1 checkpoint (relative to aifs_v11/)
+  output_dir: "./weights"
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Model  (model/aifs.py — onescience-style wrapper)
+# ═════════════════════════════════════════════════════════════════════════════
+model:
+  module: "model.aifs.AIFS"          # Python import path
+  description: >-
+    Wraps AnemoiModelEncProcDec (encoder-processor-decoder GNN) loaded from
+    the pretrained checkpoint.  Exposes input/output variable ordering,
+    N320 grid coordinates, and node weights via public properties.
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Output
+# ═════════════════════════════════════════════════════════════════════════════
+output:
+  inference_dir: "./output"
+  plot_dir: "./plots"
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Training  (config_pretraining.yaml)
+# ═════════════════════════════════════════════════════════════════════════════
+training:
+  batch_size: 1
+  from_scratch: true               # true: random weights (like FengWu/Fuxi); false: load pretrained
+  accum_grad_batches: 1
+  max_steps: 5
+  warmup_steps: 0
+  peak_lr: 3.125e-5
+  min_lr: 3.0e-7
+  precision: "16-mixed"
+  amp_dtype: "float16"                # "float16" or "bfloat16" (DCU may need bfloat16)
+  gradient_clip_val: 32.0
+  gradient_clip_algorithm: "value"
+  multistep_input: 2
+  seed: 42
+
+  optimizer:
+    betas: [0.9, 0.95]
+    weight_decay: 0.01              # official AdamW weight_decay
+
+  lr_schedule:
+    type: "cosine"
+    warmup: 1000
+    iterations: 260000
+
+  log_interval: 100
+  val_interval: 500
+  save_interval: 5000
+
+# ════════════════════════════════════════════════════��════════════════════════
+# Normalizer  (config_pretraining.yaml data.normalizer)
+# ═════════════════════════════════════════════════════════════════════════════
+normalizer:
+  default: "mean-std"
+  statistics_path: "weights/era5_stats.npz"  # auto-computed if missing; set "" for identity (fake data)
+  remap:
+    cp: "tp"
+    sf: "tp"
+  std:
+    - tp
+    - cp
+    - sf
+    - ro
+    - tcw
+    - ssrd
+    - q_50
+    - q_100
+    - q_150
+    - q_200
+    - q_250
+    - q_300
+    - q_400
+    - q_500
+    - q_600
+    - q_700
+    - q_850
+    - q_925
+    - q_1000
+  max:
+    - sdor
+    - slor
+    - z
+  none:
+    - cos_latitude
+    - cos_longitude
+    - sin_latitude
+    - sin_longitude
+    - cos_julian_day
+    - cos_local_time
+    - sin_julian_day
+    - sin_local_time
+    - insolation
+    - lsm
+    - tcc
+    - mcc
+    - hcc
+    - lcc
+    - swvl1
+    - swvl2
+
+# ═════════════════════════════════════════════════════════════════════════════
+# Loss scaling  (config_pretraining.yaml training.variable_loss_scaling)
+# ═════════════════════════════════════════════════════════════════════════════
+loss_scaling:
+  default: 1
+  pl:
+    q: 0.6
+    t: 6
+    u: 0.8
+    v: 0.5
+    w: 0.001
+    z: 12
+  sfc:
+    sp: 10
+    10u: 0.5
+    10v: 0.5
+    100u: 0.1
+    100v: 0.1
+    2d: 0.5
+    tp: 0.025
+    cp: 0.0025
+    ro: 0.0025
+    sf: 0.025
+    tcc: 0.1
+    mcc: 0.1
+    lcc: 0.1
+    hcc: 0.1
+    swvl2: 2
+    swvl1: 1
+    stl2: 10
+    stl1: 1
+    ssrd: 0.05
+    strd: 0.1
+  pressure_level_scaler:
+    minimum: 0.2
+    slope: 0.001
+
+# ═════════════════════════════════════════════════════════════════════════════
+# AIFS variable definitions  (immutable — from official checkpoint)
+# ═════════════════════════════════════════════════════════════════════════════
+aifs_variables:
+  surface:
+    - 10u
+    - 10v
+    - 2d
+    - 2t
+    - msl
+    - skt
+    - sp
+    - tcw
+    - lsm
+    - z
+    - slor
+    - sdor
+
+  soil:
+    - stl1
+    - stl2
+    - swvl1
+    - swvl2
+
+  pressure_level:
+    - z
+    - t
+    - u
+    - v
+    - w
+    - q
+
+  pressure_levels:
+    - 50
+    - 100
+    - 150
+    - 200
+    - 250
+    - 300
+    - 400
+    - 500
+    - 600
+    - 700
+    - 850
+    - 925
+    - 1000
+
+  diagnostic:
+    - tp
+    - cp
+    - sf
+    - tcc
+    - hcc
+    - lcc
+    - mcc
+    - ro
+    - ssrd
+    - strd
+    - 100u
+    - 100v
+
+  computed_forcing:
+    - cos_latitude
+    - cos_longitude
+    - sin_latitude
+    - sin_longitude
+    - cos_julian_day
+    - cos_local_time
+    - sin_julian_day
+    - sin_local_time
+    - insolation
+
+# ═════════════════════════════════════════════════════════════════════════════
+# ERA5 H5 → AIFS variable-name mapping
+# ═════════════════════════════════════════════════════════════════════════════
+# Keys = canonical AIFS names (DO NOT CHANGE).
+# Values = ERA5 H5 variable names (update to match your H5).
+era5_mapping:
+  surface:
+    10u: "10m_u_component_of_wind"
+    10v: "10m_v_component_of_wind"
+    2d: "2m_dewpoint_temperature"
+    2t: "2m_temperature"
+    msl: "mean_sea_level_pressure"
+    skt: "skin_temperature"
+    sp: "surface_pressure"
+    tcw: "total_column_water"
+    lsm: "land_sea_mask"
+    z: "geopotential"
+    slor: "slope_of_sub_gridscale_orography"
+    sdor: "standard_deviation_of_orography"
+
+  soil:
+    stl1: "soil_temperature_level_1"
+    stl2: "soil_temperature_level_2"
+    swvl1: "volumetric_soil_water_layer_1"
+    swvl2: "volumetric_soil_water_layer_2"
+
+  pressure_level:
+    z: "geopotential_{level}"
+    t: "temperature_{level}"
+    u: "u_component_of_wind_{level}"
+    v: "v_component_of_wind_{level}"
+    w: "vertical_velocity_{level}"
+    q: "specific_humidity_{level}"
+
+  diagnostic:
+    tp: "total_precipitation"
+    cp: "convective_precipitation"
+    sf: "snowfall_water_equivalent"
+    tcc: "total_cloud_cover"
+    hcc: "high_cloud_cover"
+    lcc: "low_cloud_cover"
+    mcc: "medium_cloud_cover"
+    ro: "runoff"
+    ssrd: "surface_solar_radiation_downwards"
+    strd: "surface_thermal_radiation_downwards"
+    100u: "100m_u_component_of_wind"
+    100v: "100m_v_component_of_wind"
+
+  fuzzy_fixes:
+    surface_solar_radiation_downwards: "surface_solar_radiation"
+    total_column_water_vapour: "total_column_water"

configuration.json

@@ -0,0 +1 @@
+{"framework":"Anemoi","task":"other"}

model/aifs.py

@@ -0,0 +1,535 @@
+#!/usr/bin/env python
+# coding: utf-8
+"""
+AIFS v1.1 — Onescience-style Model  (真正从零构建)
+====================================================
+
+``class AIFS(nn.Module)`` — 兼容 onescience 的 ECMWF AIFS
+编码器-处理器-解码器图神经网络。
+
+两种构建模式
+------------
+1. **从零构建** — 无需 ``.ckpt`` 文件，与 FengWu/Fuxi 一致::
+
+       model = AIFS.from_scratch(device="cuda")
+
+   从项目内的静态文件加载架构配置（``aifs_config.json``, 47 KB）和
+   N320 网格坐标（``grid-n320.npz``, 4.2 MB），通过 ``anemoi-graphs``
+   构建 HeteroData 图，由 ``AnemoiModelEncProcDec`` 随机初始化所有权重。
+
+2. **加载预训练** — 向后兼容::
+
+       model = AIFS("aifs.ckpt", device="cuda", pretrained=True)
+
+架构 (anemoi-models 0.5.0)
+--------------------------
+- **Encoder**  : ``GraphTransformerForwardMapper`` — data nodes → hidden nodes
+- **Processor**: ``TransformerProcessor`` × 16 layers — 在 o96 隐藏图上做
+  滑动窗口注意力 (window_size=1120, flash_attention, 16 heads)
+- **Decoder**  : ``GraphTransformerBackwardMapper`` — hidden → data nodes
+- **Bounding** : ``ReLUBounding`` / ``HardtanhBounding`` / ``FractionBounding``
+  — 强制输出变量满足物理约束
+
+静态文件说明
+------------
+- ``aifs_config.json``: 从官方 checkpoint 的 ``ai-models.json`` 提取
+  (model_config + data_indices + dataset), 定义完整模型架构
+- ``grid-n320.npz``: N320 高斯网格经纬度坐标 (542,080 节点), 从官方
+  checkpoint 的 ``latitudes.numpy`` / ``longitudes.numpy`` 提取
+
+Reference
+---------
+- Lang et al., *AIFS — ECMWF's data-driven forecasting system*, arXiv:2406.01465
+- anemoi-models: https://github.com/ecmwf/anemoi-models
+- anemoi-graphs: https://github.com/ecmwf/anemoi-graphs
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import numpy as np
+import torch
+from torch import nn
+
+from onescience.models.meta import ModelMetaData
+
+LOG = logging.getLogger(__name__)
+
+# ============================================================================
+# 静态文件路径（相对于本文件所在目录）
+# ============================================================================
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_CONFIG_PATH = os.path.join(_HERE, "aifs_config.json")
+_GRID_PATH = os.path.join(_HERE, "grid-n320.npz")
+
+
+# ============================================================================
+# Metadata  (onescience convention)
+# ============================================================================
+
+@dataclass
+class MetaData(ModelMetaData):
+    name: str = "AIFS"
+    jit: bool = False
+    cuda_graphs: bool = False
+    amp: bool = True
+    amp_gpu: bool = True
+    bf16: bool = False
+    onnx: bool = False              # flash_attn 不兼容 ONNX
+    func_torch: bool = False
+    auto_grad: bool = False
+    var_dim: int = -1               # 每图节点的变量数
+
+
+# ============================================================================
+# Model
+# ============================================================================
+
+class AIFS(nn.Module):
+    """AIFS v1.1 编码器-处理器-解码器 GNN。
+
+    Parameters
+    ----------
+    checkpoint_path : str, optional
+        ``pretrained=True`` 时必传——预训练 ``.ckpt`` 文件路径。
+    device : str
+        PyTorch 设备。默认 ``"cuda"``。
+    pretrained : bool
+        ``False`` (默认): 从零构建——读取本地 ``aifs_config.json`` +
+        ``grid-n320.npz``, 不依赖 ``.ckpt`` 文件。
+        ``True``: 从 ``.ckpt`` 加载预训练权重（向后兼容）。
+    """
+
+    metadata = MetaData()
+
+    # ==================================================================
+    # Factory: 真正从零（无需 .ckpt）
+    # ==================================================================
+
+    @classmethod
+    def from_scratch(cls, device: str = "cuda") -> "AIFS":
+        """从零创建 AIFS 模型——无需 ``.ckpt`` 文件。
+
+        等价于 FengWu/Fuxi 的 ``model = Fengwu()`` 模式。
+        所有权重随机初始化，架构配置和网格坐标从项目内静态文件加载。
+        """
+        return cls(device=device, pretrained=False)
+
+    # ==================================================================
+    # Construction
+    # ==================================================================
+
+    def __init__(
+        self,
+        checkpoint_path: Optional[str] = None,
+        device: str = "cuda",
+        pretrained: bool = False,
+    ):
+        super().__init__()
+
+        if pretrained:
+            if not checkpoint_path:
+                raise ValueError(
+                    "checkpoint_path is required when pretrained=True"
+                )
+            if not os.path.exists(checkpoint_path):
+                raise FileNotFoundError(
+                    f"Checkpoint not found: {checkpoint_path}"
+                )
+            self._init_from_pretrained(checkpoint_path, device)
+        else:
+            self._init_from_scratch(device)
+
+    # ==================================================================
+    # Mode 1: 从零构建（无需 .ckpt）
+    # ==================================================================
+
+    def _init_from_scratch(self, device: str):
+        """用静态配置文件和标准网格坐标构建模型。"""
+        LOG.info("Building AIFS from scratch (no checkpoint) …")
+
+        # ---- 1. 加载静态架构配置 ----------------------------------------
+        if not os.path.exists(_CONFIG_PATH):
+            raise FileNotFoundError(
+                f"Static config not found: {_CONFIG_PATH}. "
+                "Run extract_static_config.py first."
+            )
+        if not os.path.exists(_GRID_PATH):
+            raise FileNotFoundError(
+                f"Grid file not found: {_GRID_PATH}. "
+                "Run extract_static_config.py first."
+            )
+
+        with open(_CONFIG_PATH) as f:
+            config_data = json.load(f)
+        LOG.info(
+            "Loaded static config (%d KB)",
+            os.path.getsize(_CONFIG_PATH) // 1024,
+        )
+
+        grid = np.load(_GRID_PATH)
+        lat = np.asarray(grid["latitudes"], dtype=np.float32)
+        lon = np.asarray(grid["longitudes"], dtype=np.float32)
+        LOG.info(
+            "Loaded N320 grid: %d nodes", len(lat),
+        )
+
+        # ---- 2. 构建 model_config + data_indices (官方类) ---------
+        # IndexCollection 需要 OmegaConf, AnemoiModelEncProcDec 需要 DotDict.
+        # 先从原始 JSON dict 创建 OmegaConf, 再转为 DotDict.
+        raw_mc = config_data["model_config"]
+
+        from anemoi.utils.config import DotDict
+
+        # ---- 兼容性: 0.5.0 config → 0.9.0 API ---------------------------
+        # 1. activation 参数已从 mapper/processor 移除（内部默认 GELU）
+        for section in ("encoder", "decoder", "processor"):
+            raw_mc["model"][section].pop("activation", None)
+        raw_mc["model"].pop("activation", None)
+
+        # 2. layer_kernels 从顶层移入各子模块 config
+        lk = raw_mc["model"].get("layer_kernels", {})
+        for section in ("encoder", "decoder", "processor"):
+            if section in lk:
+                raw_mc["model"][section]["layer_kernels"] = lk[section]
+
+        model_config = DotDict(raw_mc)
+
+        all_vars = config_data["dataset"]["variables"]
+        name_to_index = {name: i for i, name in enumerate(all_vars)}
+
+        from anemoi.models.data_indices.collection import IndexCollection
+        from omegaconf import OmegaConf
+        data_indices = IndexCollection(
+            config=OmegaConf.create(raw_mc),
+            name_to_index=name_to_index,
+        )
+        graph_data = self._build_graph(lat, lon)
+
+        # 提取 area_weight 作为 node_weights
+        area_wt = (
+            graph_data["data"]["area_weight"].cpu().numpy().squeeze()
+        )
+
+        # ---- 4. 实例化 AnemoiModelEncProcDec（随机权重）---------------
+        # 训练时直接用裸模型——AnemoiModelInterface 自带 normalizer
+        # 预处理，会导致训练时双重归一化。保存 checkpoint 时再包装。
+        LOG.info("Instantiating AnemoiModelEncProcDec (random weights) …")
+        from anemoi.models.models.encoder_processor_decoder import \
+            AnemoiModelEncProcDec
+
+        self._model = AnemoiModelEncProcDec(
+            model_config=model_config,
+            data_indices=data_indices,
+            statistics={},
+            graph_data=graph_data,
+            truncation_data={},
+        ).to(device)
+
+        # 保存接口构建所需素材（_save 中构建 AnemoiModelInterface 用）
+        self._interface_config = model_config
+        self._interface_graph_data = graph_data
+        self._interface_data_indices = data_indices
+
+        # ---- 5. 提取变量排序和坐标 -----------------------------------
+        self._meta = config_data
+        self._arrays = {
+            "latitudes": lat,
+            "longitudes": lon,
+            "area_weight": area_wt,
+        }
+        self._extract_metadata()
+
+        LOG.info(
+            "AIFS (from scratch) ready: %d → %d vars, %d grid pts, "
+            "%.1f M params",
+            len(self._input_vars),
+            len(self._output_vars),
+            len(lat),
+            sum(p.numel() for p in self.parameters()) / 1e6,
+        )
+
+    def _build_graph(
+        self,
+        lat: np.ndarray,
+        lon: np.ndarray,
+    ) -> "HeteroData":
+        """构建 AIFS 的 HeteroData 图（N320 ↔ o96）。
+
+        完全使用 ``anemoi-graphs`` 官方 API, 不依赖任何外部数据文件。
+        节点坐标来自标准 ECMWF 网格定义, 边由纯几何算法计算。
+        """
+        from torch_geometric.data import HeteroData
+
+        from anemoi.graphs.edges.builders.cutoff import CutOffEdges
+        from anemoi.graphs.edges.builders.knn import KNNEdges
+        from anemoi.graphs.nodes.builders.from_reduced_gaussian import \
+            ReducedGaussianGridNodes
+        from anemoi.graphs.nodes.builders.from_vectors import LatLonNodes
+
+        graph = HeteroData()
+
+        # -- 数据节点 (N320: 542,080 nodes) --
+        data_builder = LatLonNodes(
+            latitudes=lat, longitudes=lon, name="data",
+        )
+        data_attrs = {
+            "area_weight": {
+                "_target_": "anemoi.graphs.nodes.attributes.SphericalAreaWeights",
+                "fill_value": 0,
+                "norm": "unit-max",
+            },
+        }
+        graph = data_builder.update_graph(graph, attrs_config=data_attrs)
+
+        # -- 隐藏节点 (o96: ~8,000 nodes) --
+        hidden_builder = ReducedGaussianGridNodes(
+            grid="o96", name="hidden",
+        )
+        graph = hidden_builder.update_graph(graph, attrs_config={})
+
+        # -- data → hidden 边 (CutOffEdges, cutoff_factor=0.6) --
+        d2h_attrs = {
+            "edge_dirs": {
+                "_target_": "anemoi.graphs.edges.attributes.EdgeDirection",
+                "norm": "unit-std",
+            },
+            "edge_length": {
+                "_target_": "anemoi.graphs.edges.attributes.EdgeLength",
+                "norm": "unit-std",
+            },
+        }
+        d2h_builder = CutOffEdges(
+            source_name="data",
+            target_name="hidden",
+            cutoff_factor=0.6,
+        )
+        graph = d2h_builder.update_graph(graph, attrs_config=None)
+        graph = d2h_builder.register_attributes(graph, d2h_attrs)
+
+        # -- hidden → data 边 (KNNEdges, K=3) --
+        h2d_attrs = {
+            "edge_dirs": {
+                "_target_": "anemoi.graphs.edges.attributes.EdgeDirection",
+                "norm": "unit-std",
+            },
+            "edge_length": {
+                "_target_": "anemoi.graphs.edges.attributes.EdgeLength",
+                "norm": "unit-std",
+            },
+        }
+        h2d_builder = KNNEdges(
+            source_name="hidden",
+            target_name="data",
+            num_nearest_neighbours=3,
+        )
+        graph = h2d_builder.update_graph(graph, attrs_config=None)
+        graph = h2d_builder.register_attributes(graph, h2d_attrs)
+
+        LOG.info(
+            "Graph built: data(%d nodes) ↔ hidden(%d nodes), "
+            "d→h edges=%d, h→d edges=%d",
+            graph["data"].num_nodes,
+            graph["hidden"].num_nodes,
+            graph["data", "to", "hidden"].edge_index.shape[1],
+            graph["hidden", "to", "data"].edge_index.shape[1],
+        )
+        return graph
+
+    # ==================================================================
+    # Mode 2: 加载预训练 (pretrained=True, 向后兼容)
+    # ==================================================================
+
+    def _init_from_pretrained(self, checkpoint_path: str, device: str):
+        """从 checkpoint 加载完整序列化模型（图 + 权重 + 索引）。"""
+        LOG.info("Loading AIFS checkpoint from %s …", checkpoint_path)
+        _cp = torch.load(
+            checkpoint_path, map_location="cpu", weights_only=False,
+        )
+        self._model = _cp.to(device)
+
+        # 提取元数据
+        self._meta, self._arrays = self._read_checkpoint_metadata(
+            checkpoint_path,
+        )
+        self._extract_metadata()
+
+        LOG.info(
+            "AIFS (pretrained) ready: %d → %d vars, %d grid pts, "
+            "%.1f M params",
+            len(self._input_vars),
+            len(self._output_vars),
+            len(self.latitudes),
+            sum(p.numel() for p in self.parameters()) / 1e6,
+        )
+
+    # ==================================================================
+    # Metadata helpers
+    # ==================================================================
+
+    @staticmethod
+    def _read_checkpoint_metadata(checkpoint_path: str):
+        """从 checkpoint ZIP 中读取 JSON 元数据和支持数组。"""
+        import zipfile
+        from anemoi.utils.checkpoints import load_supporting_arrays
+
+        with zipfile.ZipFile(checkpoint_path, "r") as zf:
+            # 支持新旧两种元数据文件名
+            metadata = None
+            for name in ["anemoi.json", "ai-models.json"]:
+                for fname in zf.namelist():
+                    if fname.endswith(name):
+                        metadata = json.load(zf.open(fname, "r"))
+                        break
+                if metadata:
+                    break
+
+            if metadata is None:
+                raise FileNotFoundError(
+                    f"No metadata JSON found in {checkpoint_path}"
+                )
+
+            arrays = load_supporting_arrays(
+                zf, metadata.get("supporting_arrays_paths", {}),
+            )
+            return metadata, arrays
+
+    def _extract_metadata(self):
+        """从已加载的元数据中提取变量排序和网格坐标。"""
+        all_vars: List[str] = self._meta["dataset"]["variables"]
+        di = self._meta["data_indices"]["data"]
+        self._input_vars = [all_vars[i] for i in di["input"]["full"]]
+        self._output_vars = [all_vars[i] for i in di["output"]["full"]]
+        self._aifs_name_to_ds_idx: Dict[str, int] = {
+            n: i for i, n in enumerate(all_vars)
+        }
+
+        lat = np.asarray(self._arrays["latitudes"], dtype=np.float32)
+        lon = np.asarray(self._arrays["longitudes"], dtype=np.float32)
+        self.register_buffer("latitudes", torch.from_numpy(lat))
+        self.register_buffer("longitudes", torch.from_numpy(lon))
+
+    # ------------------------------------------------------------------
+    # Public accessors
+    # ------------------------------------------------------------------
+
+    @property
+    def num_input_vars(self) -> int:
+        """输入张量中的变量数 (103)。"""
+        return len(self._input_vars)
+
+    @property
+    def num_output_vars(self) -> int:
+        """输出张量中的变量数 (102)。"""
+        return len(self._output_vars)
+
+    @property
+    def num_grid_points(self) -> int:
+        """N320 网格节点数 (542,080)。"""
+        return len(self.latitudes)
+
+    @property
+    def node_weights(self) -> Optional[np.ndarray]:
+        """损失函数中每节点面积权重，无则为 None。"""
+        for k in ["node_weights", "area_weight"]:
+            if k in self._arrays:
+                return np.asarray(
+                    self._arrays[k], dtype=np.float32,
+                ).squeeze()
+        return None
+
+    @property
+    def input_variables(self) -> List[str]:
+        """输入张量中 AIFS 变量名的有序列表。"""
+        return self._input_vars
+
+    @property
+    def output_variables(self) -> List[str]:
+        """输出张量中 AIFS 变量名的有序列表。"""
+        return self._output_vars
+
+    def aifs_name_to_dataset_index(self, name: str) -> int:
+        """返回 *name* 在 115 变量数据集中的位置。"""
+        return self._aifs_name_to_ds_idx.get(name, -1)
+
+    def input_name_to_channel(self, name: str) -> int:
+        """返回 *name* 在输入张量中的通道索引。"""
+        try:
+            return self._input_vars.index(name)
+        except ValueError:
+            return -1
+
+    def output_name_to_channel(self, name: str) -> int:
+        """返回 *name* 在输出张量中的通道索引。"""
+        try:
+            return self._output_vars.index(name)
+        except ValueError:
+            return -1
+
+    # ------------------------------------------------------------------
+    # Forward  (匹配官方 AnemoiModelEncProcDec.forward)
+    # ------------------------------------------------------------------
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """前向传播。
+
+        等价于 ``AnemoiModelEncProcDec.forward(x, model_comm_group=None)``
+        在单 GPU 训练/推理时的行为。
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            归一化输入。
+            - 4-D: ``(B, T, G, V_in)`` — 自动提升为 5-D。
+            - 5-D: ``(B, T, E, G, V_in)`` — ensemble 维度保留。
+
+        Returns
+        -------
+        torch.Tensor
+            - 4-D 输入 → ``(B, G, V_out)``
+            - 5-D 输入 → ``(B, E, G, V_out)``
+        """
+        if x.ndim == 4:
+            x = x.unsqueeze(2)          # (B, T, G, V) → (B, T, 1, G, V)
+
+        out = self._model(x)            # (B, E, G, V_out)
+
+        if out.shape[1] == 1:
+            out = out.squeeze(1)        # (B, 1, G, V) → (B, G, V)
+        return out
+
+    # ------------------------------------------------------------------
+    # Inference helper
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """推理前向传播（无梯度，eval 模式，autocast fp16）。
+
+        FlashAttention 要求 fp16/bf16 输入，用 autocast 自动转换。
+        """
+        was_training = self.training
+        self.eval()
+        with torch.amp.autocast("cuda", dtype=torch.float16):
+            out = self.forward(x)
+        if was_training:
+            self.train()
+        return out.float()
+
+    # ------------------------------------------------------------------
+    # Train / eval propagation
+    # ------------------------------------------------------------------
+
+    def train(self, mode: bool = True):
+        """将 train/eval 模式传播到被封装的官方模型。"""
+        super().train(mode)
+        self._model.train(mode)
+        return self
+
+    def eval(self):
+        return self.train(False)

model/grid-n320.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83085553383f2f97c5653a22368b10fe4b2a379eed0501b03a547d582102fd34
+size 4337164

scripts/fake_data.py

@@ -0,0 +1,307 @@
+#!/usr/bin/env python
+# coding: utf-8
+"""
+Generate AIFS-only ERA5 fake H5 dataset for training & inference.
+
+Follows the onescience fake-data generation pattern (chunked H5 with
+fillvalue=0, embedded global_means/global_stds).  Contains exactly the
+106 ERA5 variables required by AIFS v1.1 — no more, no less.
+
+Each year file has 60 timesteps (6h × 60 = 360h = 15 days), matching
+AIFS's maximum inference lead time.
+
+Usage::
+
+    python fake_data_all.py                          # 2005, 60 steps
+    python fake_data_all.py --years 2005,2006        # two full years
+    python fake_data_all.py -y 2005 -o ./my_era5      # custom output dir
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from typing import List
+
+import h5py
+import numpy as np
+
+# Project root (scripts/ → aifs_v11/)
+ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(ROOT))
+
+# ============================================================================
+# AIFS-required ERA5 variables  (106 total, no extras)
+# ============================================================================
+# Order: surface (12) → soil (4) → pressure-levels (78) → diagnostic (12)
+# These exactly match the keys/values in train.py's ERA5_*_MAP dictionaries.
+
+AIFS_ERA5_VARIABLES: List[str] = [
+    # ── surface prognostic (12) ─────────────────────────────────────────
+    "10m_u_component_of_wind",
+    "10m_v_component_of_wind",
+    "2m_dewpoint_temperature",
+    "2m_temperature",
+    "mean_sea_level_pressure",
+    "skin_temperature",
+    "surface_pressure",
+    "total_column_water",
+    "land_sea_mask",
+    "geopotential",
+    "slope_of_sub_gridscale_orography",
+    "standard_deviation_of_orography",
+
+    # ── soil prognostic (4) ─────────────────────────────────────────────
+    "soil_temperature_level_1",
+    "soil_temperature_level_2",
+    "volumetric_soil_water_layer_1",
+    "volumetric_soil_water_layer_2",
+
+    # ── pressure levels: geopotential (13) ──────────────────────────────
+    "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",
+
+    # ── pressure levels: temperature (13) ───────────────────────────────
+    "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",
+
+    # ── pressure levels: u wind (13) ────────────────────────────────────
+    "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",
+
+    # ── pressure levels: v wind (13) ────────────────────────────────────
+    "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",
+
+    # ── pressure levels: vertical velocity (13) ─────────────────────────
+    "vertical_velocity_50",
+    "vertical_velocity_100",
+    "vertical_velocity_150",
+    "vertical_velocity_200",
+    "vertical_velocity_250",
+    "vertical_velocity_300",
+    "vertical_velocity_400",
+    "vertical_velocity_500",
+    "vertical_velocity_600",
+    "vertical_velocity_700",
+    "vertical_velocity_850",
+    "vertical_velocity_925",
+    "vertical_velocity_1000",
+
+    # ── pressure levels: specific humidity (13) ─────────────────────────
+    "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",
+
+    # ── diagnostic (12, output-only) ────────────────────────────────────
+    "total_precipitation",
+    "convective_precipitation",
+    "snowfall_water_equivalent",
+    "total_cloud_cover",
+    "high_cloud_cover",
+    "low_cloud_cover",
+    "medium_cloud_cover",
+    "runoff",
+    "surface_solar_radiation_downwards",
+    "surface_thermal_radiation_downwards",
+    "100m_u_component_of_wind",
+    "100m_v_component_of_wind",
+]
+
+# ===========================================================================
+# Dataset dimensions
+# ===========================================================================
+# 60 timesteps @ 6h = 360h = 15 days  (covers AIFS max inference lead time)
+_DIMS = {
+    "T": 60,
+    "H": 721,
+    "W": 1440,
+    "time_step": 6,
+}
+
+
+# ===========================================================================
+# Core generation
+# ===========================================================================
+
+def generate_fake_h5(
+    output_dir: str,
+    var_names: List[str],
+    years: List[int],
+    dims: dict,
+) -> None:
+    """Generate one H5 file per year with the correct schema.
+
+    Uses HDF5 chunked datasets with ``fillvalue=0.0`` — unwritten chunks
+    return zero, keeping files tiny while preserving the logical shape.
+    """
+    data_dir = os.path.join(output_dir, "data")
+    os.makedirs(data_dir, exist_ok=True)
+
+    T, C, H, W = dims["T"], len(var_names), dims["H"], dims["W"]
+    means = np.zeros((1, C, 1, 1), dtype=np.float32)
+    stds = np.ones((1, C, 1, 1), dtype=np.float32)
+    logical_gib = T * C * H * W * 4 / 1024**3
+
+    for year in years:
+        path = os.path.join(data_dir, f"{year}.h5")
+        with h5py.File(path, "w") as f:
+            ds = f.create_dataset(
+                "fields",
+                shape=(T, C, H, W),
+                dtype="float32",
+                chunks=(1, C, H, W),
+                fillvalue=0.0,
+            )
+            ds.attrs["variables"] = var_names
+            ds.attrs["time_step"] = dims["time_step"]
+            f.create_dataset("global_means", data=means)
+            f.create_dataset("global_stds", data=stds)
+
+        size_mb = os.path.getsize(path) / 1024**2
+        print(f"  {year}.h5  shape=({T},{C},{H},{W})  "
+              f"logical={logical_gib:.1f} GiB  actual={size_mb:.1f} MiB")
+
+
+# ===========================================================================
+# CLI
+# ===========================================================================
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="Generate AIFS-only ERA5 fake H5 files",
+    )
+    p.add_argument("--config", "-c", type=str,
+                   default=str(ROOT / "conf" / "config.yaml"),
+                   help="Path to config.yaml")
+    p.add_argument("--output_dir", "-o", type=str,
+                   default=str(ROOT / "fake_era5"),
+                   help="Output root directory")
+    p.add_argument("--years", "-y", type=str, default=None,
+                   help="Comma-separated years (overrides config).  "
+                        "Each year = 60 timesteps (15 days @ 6h).")
+    return p.parse_args()
+
+
+# ===========================================================================
+# Main
+# ===========================================================================
+
+def main() -> None:
+    args = parse_args()
+
+    # Resolve years: CLI > config train+val+test > default
+    if args.years:
+        years = [int(y.strip()) for y in args.years.replace("，", ",").split(",")
+                 if y.strip()]
+    elif os.path.exists(args.config):
+        import yaml
+        cfg = yaml.safe_load(open(args.config))
+        data = cfg.get("data", {})
+        raw = (data.get("train_years", [])
+               + data.get("val_years", [])
+               + data.get("test_years", []))
+        years = sorted(set(raw))
+        if not years:
+            years = [2005]
+    else:
+        years = [2005]
+
+    dims = _DIMS
+
+    print("=" * 60)
+    print("  AIFS v1.1 — Fake ERA5 Dataset Generator")
+    print("=" * 60)
+    print(f"  Output dir : {os.path.abspath(args.output_dir)}")
+    print(f"  Years      : {years}  ({dims['T']} steps = "
+          f"{dims['T'] * dims['time_step'] // 24} days each)")
+    print(f"  Variables  : {len(AIFS_ERA5_VARIABLES)}  (AIFS-only, no extras)")
+    print(f"  Shape      : ({dims['T']}, {len(AIFS_ERA5_VARIABLES)}, "
+          f"{dims['H']}, {dims['W']})")
+    print(f"  Timestep   : {dims['time_step']}h")
+    if not args.years:
+        print(f"  (years auto-collected from config: train+val+test)")
+    print()
+
+    print("[1/1] Generating per-year H5 files ...")
+    generate_fake_h5(args.output_dir, AIFS_ERA5_VARIABLES, years, dims)
+
+    # Quick verification
+    diag_count = sum(
+        1 for v in AIFS_ERA5_VARIABLES
+        if v in {"total_precipitation", "convective_precipitation",
+                 "snowfall_water_equivalent", "total_cloud_cover",
+                 "high_cloud_cover", "low_cloud_cover", "medium_cloud_cover",
+                 "runoff", "surface_solar_radiation_downwards",
+                 "surface_thermal_radiation_downwards",
+                 "100m_u_component_of_wind", "100m_v_component_of_wind"}
+    )
+    pl_count = sum(1 for v in AIFS_ERA5_VARIABLES
+                   if any(v.endswith(f"_{lvl}") for lvl in
+                          [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]))
+    sfc_count = len(AIFS_ERA5_VARIABLES) - pl_count - diag_count
+
+    print(f"\n  Variables: {sfc_count} surface/soil + {pl_count} PL + "
+          f"{diag_count} diagnostic = {len(AIFS_ERA5_VARIABLES)}")
+    print(f"\n✅ Done.  Use with train.py:")
+    print(f"   python train.py --dataset_path {os.path.abspath(args.output_dir)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,720 @@
+#!/usr/bin/env python
+# coding: utf-8
+"""
+AIFS v1.1 — Inference Script  (AIFS-powered)
+===============================================
+10-day deterministic forecast from ERA5 initial conditions.
+
+Uses ``model.aifs.AIFS`` for autoregressive forecast — variable ordering,
+normalisation, and model forward are guaranteed to match training.
+
+Usage:
+    python scripts/inference.py
+    python scripts/inference.py --checkpoint weights/aifs_pretrain_final.ckpt
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import math
+import os
+import sys
+import time
+import traceback
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pytz
+import torch
+import yaml
+from tqdm import tqdm
+
+# Project root (scripts/ → aifs_v11/)
+ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(ROOT))
+
+import earthkit.regrid as ekr
+from onescience.datapipes.climate.era5 import ERA5Dataset
+from model.aifs import AIFS
+
+
+# ============================================================================
+# Config loader
+# ============================================================================
+
+def load_config(path: str) -> dict:
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+# ============================================================================
+# ERA5 variable list builder  (matches train.py)
+# ============================================================================
+
+def build_era5_variable_list(cfg: dict) -> List[str]:
+    av = cfg["aifs_variables"]
+    em = cfg["era5_mapping"]
+    vars_: List[str] = []
+    for name in av["surface"]:
+        vars_.append(em["surface"][name])
+    for name in av["soil"]:
+        vars_.append(em["soil"][name])
+    for v in av["pressure_level"]:
+        tpl = em["pressure_level"][v]
+        for lvl in av["pressure_levels"]:
+            vars_.append(tpl.format(level=lvl))
+    for name in av["diagnostic"]:
+        vars_.append(em["diagnostic"][name])
+    return vars_
+
+
+def _auto_correct_variable_names(required: List[str], available: set, cfg: dict) -> List[str]:
+    fuzzy = cfg["era5_mapping"].get("fuzzy_fixes", {})
+    corrected = list(required)
+    for i, v in enumerate(corrected):
+        if v not in available and v in fuzzy:
+            alt = fuzzy[v]
+            if alt in available:
+                print(f"[INFO] Auto-corrected: '{v}' → '{alt}'")
+                corrected[i] = alt
+    return corrected
+
+
+# ============================================================================
+# Normalisation  (consistent with train.py)
+# ============================================================================
+
+def compute_normalisation_params(
+    cfg: dict, variable_names: List[str],
+    statistics_path: Optional[str] = None,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute per-variable norm_mul and norm_add from config.
+
+    If ``statistics_path`` provides a ``.npz`` with ``mean``/``stdev``
+    arrays, those values are used.  Otherwise identity.
+    """
+    nc = cfg["normalizer"]
+    default_method = nc["default"]
+    remap = nc.get("remap", {}) or {}
+    methods: Dict[str, str] = {}
+    for v in variable_names:
+        if v in (nc.get("none") or []):
+            methods[v] = "none"
+        elif v in (nc.get("max") or []):
+            methods[v] = "max"
+        elif v in (nc.get("min-max") or []):
+            methods[v] = "min-max"
+        elif v in (nc.get("std") or []):
+            methods[v] = "std"
+        else:
+            methods[v] = default_method
+
+    name_to_idx = {n: i for i, n in enumerate(variable_names)}
+    num_vars = len(variable_names)
+    _mean = np.zeros(num_vars, dtype=np.float32)
+    _stdev = np.ones(num_vars, dtype=np.float32)
+
+    if statistics_path and os.path.exists(statistics_path):
+        stats = np.load(statistics_path)
+        if "variables" in stats:
+            ds_vars = [str(v) for v in stats["variables"]]
+        else:
+            import json
+            ds_vars = json.load(
+                open(str(ROOT / "model" / "aifs_config.json"))
+            )["dataset"]["variables"]
+        ds_name_to_idx = {n: i for i, n in enumerate(ds_vars)}
+        for name, i in name_to_idx.items():
+            if name in ds_name_to_idx:
+                j = ds_name_to_idx[name]
+                _mean[i] = float(stats["mean"][j])
+                _stdev[i] = float(stats["stdev"][j])
+        print(f"[INFO] Loaded statistics from {statistics_path}")
+
+    for target_var, source_var in remap.items():
+        if target_var in name_to_idx and source_var in name_to_idx:
+            ti, si = name_to_idx[target_var], name_to_idx[source_var]
+            _mean[ti], _stdev[ti] = _mean[si], _stdev[si]
+
+    norm_mul = np.ones(num_vars, dtype=np.float32)
+    norm_add = np.zeros(num_vars, dtype=np.float32)
+    for name, i in name_to_idx.items():
+        method = methods.get(name, default_method)
+        if method == "mean-std":
+            norm_mul[i] = 1.0 / max(float(_stdev[i]), 1e-9)
+            norm_add[i] = -float(_mean[i]) / max(float(_stdev[i]), 1e-9)
+        elif method == "std":
+            norm_mul[i] = 1.0 / max(float(_stdev[i]), 1e-9)
+        elif method in ("max", "min-max", "none"):
+            norm_mul[i] = 1.0
+    return norm_mul, norm_add
+
+
+def build_denorm_params(norm_mul: np.ndarray, norm_add: np.ndarray):
+    """Invert normalisation:  x = (x_norm - add) / mul."""
+    eps = 1e-9
+    denorm_mul = np.where(np.abs(norm_mul) > eps, 1.0 / norm_mul, 1.0)
+    denorm_add = -norm_add * denorm_mul
+    return denorm_mul.astype(np.float32), denorm_add.astype(np.float32)
+
+
+# ============================================================================
+# Forcing features  (consistent with train.py)
+# ============================================================================
+
+def _compute_insolation(
+    ts: datetime.datetime, lat_rad: np.ndarray, lon_rad: np.ndarray
+) -> np.ndarray:
+    ref = datetime.datetime(2000, 1, 1, 12, 0, 0, tzinfo=pytz.utc)
+    mt = ts.timestamp()
+    days = (mt - ref.timestamp()) / (24.0 * 3600.0)
+    jc = days / 36525.0
+    theta = 67310.54841 + jc * (
+        876600.0 * 3600.0
+        + 8640184.812866
+        + jc * (0.093104 - jc * 6.2e-5)
+    )
+    gmst = np.fmod((theta / 240.0) * np.pi / 180.0, 2.0 * np.pi)
+    ma = np.deg2rad(
+        357.52910 + 35999.05030 * jc - 0.0001559 * jc**2 - 0.00000048 * jc**3
+    )
+    ml = np.deg2rad(280.46645 + 36000.76983 * jc + 0.0003032 * jc**2)
+    dl = np.deg2rad(
+        (1.914600 - 0.004817 * jc - 0.000014 * jc**2) * np.sin(ma)
+        + (0.019993 - 0.000101 * jc) * np.sin(2.0 * ma)
+        + 0.000290 * np.sin(3.0 * ma)
+    )
+    tl = ml + dl
+    eps = np.deg2rad(
+        23.0
+        + 26.0 / 60.0
+        + 21.406 / 3600.0
+        - (
+            46.836769 * jc
+            - 0.0001831 * jc**2
+            + 0.00200340 * jc**3
+            - 0.576e-6 * jc**4
+            - 4.34e-8 * jc**5
+        )
+        / 3600.0
+    )
+    x = np.cos(tl)
+    y = np.cos(eps) * np.sin(tl)
+    z = np.sin(eps) * np.sin(tl)
+    r = np.sqrt(1.0 - z * z)
+    dec = np.arctan2(z, r)
+    ra = 2.0 * np.arctan2(y, x + r)
+    ha = (gmst + lon_rad) - ra
+    cos_z = np.sin(lat_rad) * np.sin(dec) + np.cos(lat_rad) * np.cos(dec) * np.cos(ha)
+    return cos_z.astype(np.float32)
+
+
+def compute_forcing_features(
+    latitudes: np.ndarray,
+    longitudes: np.ndarray,
+    timestamps: List[datetime.datetime],
+    static_fields: Dict[str, np.ndarray],
+) -> Dict[str, np.ndarray]:
+    """Compute the 13 forcing channels for a 2-timestep window.
+
+    *static_fields* supplies lsm, z, slor, sdor (time-invariant).
+    """
+    lat_rad = np.deg2rad(latitudes)
+    lon_rad = np.deg2rad(longitudes)
+    multi_step = len(timestamps)
+    forcing: Dict[str, np.ndarray] = {}
+
+    forcing["cos_latitude"] = np.tile(
+        np.cos(lat_rad)[np.newaxis, :], (multi_step, 1)
+    )
+    forcing["sin_latitude"] = np.tile(
+        np.sin(lat_rad)[np.newaxis, :], (multi_step, 1)
+    )
+    forcing["cos_longitude"] = np.tile(
+        np.cos(lon_rad)[np.newaxis, :], (multi_step, 1)
+    )
+    forcing["sin_longitude"] = np.tile(
+        np.sin(lon_rad)[np.newaxis, :], (multi_step, 1)
+    )
+
+    cos_jd = np.zeros((multi_step, len(latitudes)), dtype=np.float32)
+    sin_jd = np.zeros_like(cos_jd)
+    cos_lt = np.zeros_like(cos_jd)
+    sin_lt = np.zeros_like(cos_jd)
+    insol = np.zeros_like(cos_jd)
+
+    for t_idx, ts in enumerate(timestamps):
+        doy = ts.timetuple().tm_yday
+        jd_angle = 2.0 * np.pi * doy / 365.25
+        cos_jd[t_idx] = np.cos(jd_angle)
+        sin_jd[t_idx] = np.sin(jd_angle)
+        hours = ts.hour + ts.minute / 60.0 + ts.second / 3600.0
+        lt_angle = 2.0 * np.pi * hours / 24.0 + lon_rad
+        cos_lt[t_idx] = np.cos(lt_angle)
+        sin_lt[t_idx] = np.sin(lt_angle)
+        insol[t_idx] = _compute_insolation(ts, lat_rad, lon_rad)
+
+    forcing["cos_julian_day"] = cos_jd
+    forcing["sin_julian_day"] = sin_jd
+    forcing["cos_local_time"] = cos_lt
+    forcing["sin_local_time"] = sin_lt
+    forcing["insolation"] = insol
+
+    # Static fields — copied from initial conditions, held invariant
+    for var_name in ["lsm", "z", "slor", "sdor"]:
+        if var_name in static_fields:
+            forcing[var_name] = np.tile(
+                static_fields[var_name][:1], (multi_step, 1)
+            )
+    return forcing
+
+
+# ============================================================================
+# N320 interpolation
+# ============================================================================
+
+def _interp_n320(arr: np.ndarray) -> np.ndarray:
+    return ekr.interpolate(arr, {"grid": (0.25, 0.25)}, {"grid": "N320"})
+
+
+# ============================================================================
+# ERA5 frame loading
+# ============================================================================
+
+def load_frames(era5_dir: str, year: int, step_idx: int, used_vars: List[str]):
+    """Load two consecutive frames for the initial condition."""
+    # ERA5Dataset expects dataset_dir to contain data/{year}.h5,
+    # so if our path is ./fake_era5, data is at ./fake_era5/data/2008.h5
+    ds = ERA5Dataset(
+        dataset_dir=era5_dir,
+        used_years=[year],
+        used_variables=used_vars,
+        input_steps=2,
+        output_steps=0,
+        normalize=False,
+    )
+    invar, _, _, _, _ = ds[step_idx]
+    return invar[0].numpy(), invar[1].numpy()
+
+
+# ============================================================================
+# ERA5 → AIFS field construction
+# ============================================================================
+
+def build_aifs_fields_from_frames(
+    frame_tm6: np.ndarray,
+    frame_t0: np.ndarray,
+    era5_var_list: List[str],
+    cfg: dict,
+    diag_map: Optional[Dict[str, str]] = None,
+) -> Dict[str, np.ndarray]:
+    """Convert raw ERA5 frames to per-variable N320-interpolated fields.
+
+    Returns
+    -------
+    dict  {aifs_name: np.array([frame_t-6h, frame_t0])}
+    """
+    av = cfg["aifs_variables"]
+    em = cfg["era5_mapping"]
+    name_to_ch = {n: i for i, n in enumerate(era5_var_list)}
+    fields: Dict[str, np.ndarray] = {}
+
+    for aifs_name in av["surface"]:
+        ch = name_to_ch[em["surface"][aifs_name]]
+        fields[aifs_name] = np.stack(
+            [_interp_n320(frame_tm6[ch]), _interp_n320(frame_t0[ch])]
+        )
+    for aifs_name in av["soil"]:
+        ch = name_to_ch[em["soil"][aifs_name]]
+        fields[aifs_name] = np.stack(
+            [_interp_n320(frame_tm6[ch]), _interp_n320(frame_t0[ch])]
+        )
+    for aifs_var in av["pressure_level"]:
+        tpl = em["pressure_level"][aifs_var]
+        for level in av["pressure_levels"]:
+            ch = name_to_ch[tpl.format(level=level)]
+            fields[f"{aifs_var}_{level}"] = np.stack(
+                [_interp_n320(frame_tm6[ch]), _interp_n320(frame_t0[ch])]
+            )
+    _diag_map = diag_map if diag_map is not None else em["diagnostic"]
+    for aifs_name in av["diagnostic"]:
+        era5_name = _diag_map.get(aifs_name)
+        if era5_name and era5_name in name_to_ch:
+            ch = name_to_ch[era5_name]
+            fields[aifs_name] = np.stack(
+                [_interp_n320(frame_tm6[ch]), _interp_n320(frame_t0[ch])]
+            )
+    return fields
+
+
+# ============================================================================
+# Autoregressive forecaster
+# ============================================================================
+
+class AIFSForecaster:
+    """Deterministic autoregressive forecaster using the AIFS model.
+
+    Parameters
+    ----------
+    model : AIFS
+        Loaded AIFS model wrapper.
+    norm_mul : np.ndarray  shape (V_in,)
+        Per-input-channel multiplicative normalisation.
+    norm_add : np.ndarray  shape (V_in,)
+        Per-input-channel additive normalisation.
+    denorm_mul : np.ndarray  shape (V_out,)
+        Per-output-channel multiplicative denormalisation.
+    denorm_add : np.ndarray  shape (V_out,)
+        Per-output-channel additive denormalisation.
+    device : str
+        PyTorch device.
+    """
+
+    # Variables that are time-invariant — copied from initial condition
+    _STATIC_NAMES = {"lsm", "z", "slor", "sdor"}
+
+    def __init__(
+        self,
+        model: AIFS,
+        norm_mul: np.ndarray,
+        norm_add: np.ndarray,
+        denorm_mul: np.ndarray,
+        denorm_add: np.ndarray,
+        device: str = "cuda",
+    ):
+        self.model = model
+        self.device = device
+
+        self._input_vars = model.input_variables
+        self._output_vars = model.output_variables
+        self._norm_mul = norm_mul.astype(np.float32)
+        self._norm_add = norm_add.astype(np.float32)
+        self._denorm_mul = denorm_mul.astype(np.float32)
+        self._denorm_add = denorm_add.astype(np.float32)
+
+        # Channel lookups
+        self._in_ch: Dict[str, int] = {
+            n: i for i, n in enumerate(self._input_vars)
+        }
+        self._out_ch: Dict[str, int] = {
+            n: i for i, n in enumerate(self._output_vars)
+        }
+
+        # Input variable → output variable channel mapping (for feedback)
+        # Only prognostic vars exist in both input and output.
+        self._in_to_out: Dict[int, int] = {}
+        for i_idx, name in enumerate(self._input_vars):
+            if name in self._out_ch:
+                self._in_to_out[i_idx] = self._out_ch[name]
+
+        # Identify which input vars are prognostic (fed back from output)
+        # vs forcing (computed each step).
+        self._prognostic_input_indices = sorted(self._in_to_out.keys())
+        self._forcing_input_indices = sorted(
+            set(range(len(self._input_vars))) - set(self._in_to_out.keys())
+        )
+
+        n_prog, n_forc = (
+            len(self._prognostic_input_indices),
+            len(self._forcing_input_indices),
+        )
+        print(
+            f"[INFO] Forecaster: {n_prog} prognostic + {n_forc} forcing "
+            f"= {len(self._input_vars)} input vars → {len(self._output_vars)} output vars"
+        )
+
+    # ------------------------------------------------------------------
+    def forecast(
+        self,
+        initial_fields: Dict[str, np.ndarray],
+        start_date: datetime.datetime,
+        lead_time_hours: int,
+        latitudes: np.ndarray,
+        longitudes: np.ndarray,
+    ) -> List[Dict]:
+        """Run an autoregressive forecast.
+
+        Parameters
+        ----------
+        initial_fields : dict  {aifs_name: np.array([frame_t-6h, frame_t0])}
+            Initial atmospheric state on the N320 grid.
+        start_date : datetime
+            Valid time of the second frame (t0).
+        lead_time_hours : int
+            Total forecast length in hours (multiple of 6).
+        latitudes, longitudes : np.ndarray
+            N320 node coordinates.
+
+        Returns
+        -------
+        list of dict
+            One dict per output step::
+                {"date": datetime, "latitudes": array, "longitudes": array,
+                 "fields": {var_name: array}}
+        """
+        num_steps = lead_time_hours // 6
+        states: List[Dict] = []
+        fields = initial_fields.copy()
+        current_date = start_date
+        num_nodes = len(latitudes)
+
+        t_start = time.time()
+        pbar = tqdm(
+            range(num_steps), desc="Forecast", unit="step",
+            dynamic_ncols=True,
+        )
+        for step in pbar:
+            # ---- 1.  Timestamps for this window ----------------------------
+            t0 = current_date
+            t_m6 = t0 - datetime.timedelta(hours=6)
+
+            # ---- 2.  Compute time-dependent forcing -----------------------
+            forcing = compute_forcing_features(
+                latitudes, longitudes, [t_m6, t0], fields
+            )
+
+            # ---- 3.  Assemble normalised input tensor (2, G, V_in) --------
+            x = np.zeros(
+                (2, num_nodes, len(self._input_vars)), dtype=np.float32
+            )
+            for var_name, ch in self._in_ch.items():
+                if var_name in fields:
+                    x[:, :, ch] = fields[var_name]
+                elif var_name in forcing:
+                    x[:, :, ch] = forcing[var_name]
+                # else: stays zero (should not happen for valid inputs)
+
+            x = x * self._norm_mul[np.newaxis, np.newaxis, :] + self._norm_add[
+                np.newaxis, np.newaxis, :
+            ]
+
+            # ---- 4.  Model forward ----------------------------------------
+            x_t = torch.from_numpy(x).unsqueeze(0).to(self.device)
+            with torch.no_grad():
+                pred = self.model.predict(x_t)  # (1, G, V_out)
+            pred_np = pred[0].cpu().numpy()  # (G, V_out)
+
+            # ---- 5.  Denormalise output -----------------------------------
+            pred_phys = (
+                pred_np * self._denorm_mul[np.newaxis, :]
+                + self._denorm_add[np.newaxis, :]
+            )
+
+            # ---- 6.  Parse into fields dict -------------------------------
+            next_fields: Dict[str, np.ndarray] = {}
+            for var_name, ch in self._out_ch.items():
+                next_fields[var_name] = pred_phys[:, ch].copy()
+
+            # ---- 7.  Save state -------------------------------------------
+            next_date = t0 + datetime.timedelta(hours=6)
+            states.append(
+                {
+                    "date": next_date,
+                    "latitudes": latitudes.copy(),
+                    "longitudes": longitudes.copy(),
+                    "fields": {k: v.copy() for k, v in next_fields.items()},
+                }
+            )
+
+            # ---- 8.  Shift for next autoregressive step -------------------
+            # New window: [old t0, prediction]
+            shifted: Dict[str, np.ndarray] = {}
+            for var_name in fields:
+                if var_name in next_fields:
+                    shifted[var_name] = np.stack(
+                        [fields[var_name][1], next_fields[var_name]], axis=0
+                    )
+                elif var_name in self._STATIC_NAMES:
+                    # Static fields: keep invariant
+                    shifted[var_name] = fields[var_name].copy()
+            fields = shifted
+            current_date = next_date
+
+            # 更新进度条
+            elapsed = time.time() - t_start
+            pbar.set_postfix({
+                "date": next_date.strftime("%m-%d %H:%M"),
+                "spd": f"{elapsed / (step + 1):.1f}s",
+            })
+
+        elapsed = time.time() - t_start
+        print(
+            f"[INFO] Forecast done: {num_steps} steps in {elapsed:.1f}s "
+            f"({elapsed / max(num_steps, 1):.2f}s/step)"
+        )
+        return states
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def main():
+    p = argparse.ArgumentParser(description="AIFS v1.1 — 10-day Inference")
+    p.add_argument("-c", "--config", default=str(ROOT / "conf" / "config.yaml"))
+    p.add_argument("--checkpoint", default=None, help="override checkpoint path")
+    p.add_argument("--year", type=int, default=None, help="override test year")
+    p.add_argument(
+        "--lead_time", type=int, default=None, help="override lead time (h)"
+    )
+    p.add_argument("-o", "--output_dir", default=None)
+    p.add_argument("--device", default=None)
+    p.add_argument("--device_ids", default=None)
+    args = p.parse_args()
+
+    cfg = load_config(args.config) if os.path.exists(args.config) else {}
+    hw = cfg.get("hardware", {})
+    ck = cfg.get("checkpoint", {})
+    data = cfg.get("data", {})
+    out = cfg.get("output", {})
+
+    device_str = args.device or hw.get("device", "dcu")
+    device_ids = args.device_ids or str(hw.get("device_ids", "0"))
+    ckpt = args.checkpoint or ck.get("pretrained", "")
+    test_years = data.get("test_years", [2008])
+    test_year = args.year or test_years[0]
+    lead_time = args.lead_time or data.get("test_lead_time", 240)
+    era5_dir = data.get("data_dir", str(ROOT / "data" / "fake_era5"))
+    out_dir = args.output_dir or str(ROOT / "output")
+
+    os.environ["CUDA_VISIBLE_DEVICES"] = device_ids
+    if device_str == "dcu":
+        os.environ["HIP_VISIBLE_DEVICES"] = device_ids
+    os.environ.setdefault(
+        "PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True"
+    )
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if device_str == "cpu":
+        print("warning: CPU running")
+        device = "cpu"
+
+    # ---- Analysis date: Jan 1 06:00 of test_year -------------------------
+    step_idx = 1  # second 6h step → Jan 1 06:00
+    analysis_date = datetime.datetime(
+        test_year, 1, 1, 6, 0, 0, tzinfo=pytz.utc
+    )
+    date_str = analysis_date.strftime("%Y%m%d%H")
+
+    print(f"\n{'='*60}\n  AIFS v1.1 — 10-day Forecast\n{'='*60}")
+    print(f"  Test year : {test_year}  |  analysis: {date_str}")
+    print(
+        f"  Lead time : {lead_time}h ({lead_time // 6} steps)"
+    )
+    print(f"  Device    : {device}  |  ckpt: {ckpt}\n{'='*60}\n")
+
+    try:
+        # ---- 1.  Load AIFS model ----------------------------------------
+        t0 = time.time()
+        model = AIFS(ckpt, device=device, pretrained=True)
+        print(f"[INFO] Model loaded ({time.time() - t0:.1f}s)")
+
+        input_vars = model.input_variables
+        output_vars = model.output_variables
+        grid_lat = model.latitudes.cpu().numpy()
+        grid_lon = model.longitudes.cpu().numpy()
+
+        # ---- 2.  Compute normalisation params ---------------------------
+        stats_path = cfg["normalizer"].get("statistics_path") or None
+        norm_mul_in, norm_add_in = compute_normalisation_params(
+            cfg, input_vars, statistics_path=stats_path,
+        )
+        norm_mul_out, norm_add_out = compute_normalisation_params(
+            cfg, output_vars, statistics_path=stats_path,
+        )
+        denorm_mul, denorm_add = build_denorm_params(norm_mul_out, norm_add_out)
+
+        # ---- 3.  Load ERA5 initial frames -------------------------------
+        era5_list = build_era5_variable_list(cfg)
+
+        # Auto-correct names against available H5 variables
+        data_dir = os.path.join(era5_dir, "data")
+        if os.path.isdir(data_dir):
+            h5_files = sorted(
+                [f for f in os.listdir(data_dir) if f.endswith(".h5")]
+            )
+            available_set: set = set()
+            if h5_files:
+                import h5py
+
+                with h5py.File(
+                    os.path.join(data_dir, h5_files[0]), "r"
+                ) as f:
+                    available_vars = [
+                        v.decode() if isinstance(v, bytes) else v
+                        for v in f["fields"].attrs["variables"]
+                    ]
+                available_set = set(available_vars)
+            corrected = _auto_correct_variable_names(
+                era5_list, available_set, cfg
+            )
+            used_vars = [v for v in corrected if v in available_set]
+            missing = [v for v in corrected if v not in available_set]
+            if missing:
+                print(
+                    f"[WARN] {len(missing)} vars missing from H5: {missing[:10]}"
+                )
+        else:
+            used_vars = era5_list
+
+        # Build diagnostic map with corrected names
+        diag_map: Dict[str, str] = {}
+        for aifs_name, orig in cfg["era5_mapping"]["diagnostic"].items():
+            cn = _auto_correct_variable_names([orig], available_set, cfg)[0]
+            diag_map[aifs_name] = (
+                cn if cn in available_set else orig
+            )
+
+        ft6, ft0 = load_frames(era5_dir, test_year, step_idx, used_vars)
+        fields = build_aifs_fields_from_frames(
+            ft6, ft0, used_vars, cfg, diag_map=diag_map
+        )
+
+        # ---- 4.  Autoregressive forecast --------------------------------
+        forecaster = AIFSForecaster(
+            model,
+            norm_mul_in,
+            norm_add_in,
+            denorm_mul,
+            denorm_add,
+            device=device,
+        )
+
+        initial_state = {
+            "date": analysis_date,
+            "fields": fields,
+        }
+
+        states = forecaster.forecast(
+            initial_fields=fields,
+            start_date=analysis_date,
+            lead_time_hours=lead_time,
+            latitudes=grid_lat,
+            longitudes=grid_lon,
+        )
+
+        # ---- 5.  Save to .npz -------------------------------------------
+        os.makedirs(out_dir, exist_ok=True)
+        for s in states:
+            vt = s["date"].strftime("%Y%m%d%H")
+            path = os.path.join(out_dir, f"aifs_forecast_{vt}.npz")
+            d = dict(
+                latitudes=s["latitudes"], longitudes=s["longitudes"]
+            )
+            d.update(s["fields"])
+            np.savez_compressed(
+                path, **d, date=np.array(str(s["date"]))
+            )
+            print(f"[INFO] Saved: {path}")
+
+        print(f"[INFO] Done — {len(states)} files → {out_dir}")
+
+    except Exception:
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,396 @@
+#!/usr/bin/env python
+# coding: utf-8
+"""
+AIFS v1.1 — Evaluation & Visualisation
+========================================
+Computes RMSE and ACC (Anomaly Correlation Coefficient) on inference
+results, then plots selected variables on the N320 Gaussian grid.
+
+Usage:
+    python scripts/result.py
+    python scripts/result.py -v 2t,z_500,tp
+    python scripts/result.py -v all
+    python scripts/result.py --no-metrics           # plot only
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import yaml
+from tqdm import tqdm
+
+# Project root (scripts/ → aifs_v11/)
+ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(ROOT))
+
+
+# ============================================================================
+# Variable metadata
+# ============================================================================
+_VAR_META = {
+    "10u": ("10m U-wind", "RdBu_r", "m/s"),
+    "10v": ("10m V-wind", "RdBu_r", "m/s"),
+    "2t": ("2m Temperature", "RdYlBu_r", "K"),
+    "2d": ("2m Dewpoint", "RdYlBu_r", "K"),
+    "msl": ("MSLP", "RdYlBu_r", "Pa"),
+    "skt": ("Skin Temp", "RdYlBu_r", "K"),
+    "sp": ("Surface Pressure", "RdYlBu_r", "Pa"),
+    "tcw": ("Total Column Water", "Blues", "kg/m²"),
+    "z": ("Surface Geopotential", "terrain", "m²/s²"),
+    "tp": ("Total Precip", "YlGnBu", "m"),
+    "cp": ("Convective Precip", "YlGnBu", "m"),
+    "sf": ("Snowfall", "PuBu", "m"),
+    "tcc": ("Total Cloud Cover", "Greys", "0-1"),
+    "lcc": ("Low Cloud", "Greys", "0-1"),
+    "mcc": ("Mid Cloud", "Greys", "0-1"),
+    "hcc": ("High Cloud", "Greys", "0-1"),
+    "ro": ("Runoff", "YlGnBu", "m"),
+    "ssrd": ("Solar Radiation", "YlOrRd", "J/m²"),
+    "strd": ("Thermal Radiation", "YlOrRd", "J/m²"),
+    "100u": ("100m U-wind", "RdBu_r", "m/s"),
+    "100v": ("100m V-wind", "RdBu_r", "m/s"),
+    "stl1": ("Soil Temp L1", "RdYlBu_r", "K"),
+    "stl2": ("Soil Temp L2", "RdYlBu_r", "K"),
+    "swvl1": ("Soil Moisture L1", "Blues", "m³/m³"),
+    "swvl2": ("Soil Moisture L2", "Blues", "m³/m³"),
+}
+_PL_META = {
+    "z": ("Geopotential", "RdYlBu_r", "m²/s²"),
+    "t": ("Temperature", "RdYlBu_r", "K"),
+    "u": ("U-wind", "RdBu_r", "m/s"),
+    "v": ("V-wind", "RdBu_r", "m/s"),
+    "w": ("Vertical Velocity", "RdBu_r", "Pa/s"),
+    "q": ("Specific Humidity", "Blues", "kg/kg"),
+}
+for _a, (_desc, _cmap, _unit) in list(_PL_META.items()):
+    for _l in [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50]:
+        _VAR_META[f"{_a}_{_l}"] = (f"{_desc} {_l}hPa", _cmap, _unit)
+
+DEFAULT_VARS = ["10u", "2t", "msl", "tp", "z_500", "t_850"]
+
+
+# ============================================================================
+# Helpers
+# ============================================================================
+
+def _load_config():
+    config_path = ROOT / "conf" / "config.yaml"
+    if config_path.exists():
+        return yaml.safe_load(open(config_path))
+    return {}
+
+# ============================================================================
+# ACC / RMSE computation
+# ============================================================================
+
+def compute_metrics(pred_dir: str, label_dir: str, test_years: list):
+    """Compute per-variable RMSE and ACC vs ERA5 ground truth.
+
+    Parameters
+    ----------
+    pred_dir : str
+        Directory of ``aifs_forecast_*.npz`` files.
+    label_dir : str
+        Root of ERA5 H5 data (contains ``data/{year}.h5``).
+    test_years : list[int]
+        Years to evaluate.
+    """
+    import h5py
+
+    files = sorted(glob.glob(os.path.join(pred_dir, "aifs_forecast_*.npz")))
+    if not files:
+        raise FileNotFoundError(f"No forecast .npz in {pred_dir}")
+
+    # Discover variables from first prediction
+    sample = np.load(files[0])
+    var_names = sorted(
+        k for k in sample.keys()
+        if k not in ("latitudes", "longitudes", "date")
+    )
+    num_vars = len(var_names)
+    sample.close()
+
+    # Build AIFS → ERA5 H5 channel mapping from config
+    cfg = _load_config()
+    em = cfg.get("era5_mapping", {})
+    av = cfg.get("aifs_variables", {})
+    aifs_to_era5 = {}
+    for cat in ["surface", "soil", "diagnostic"]:
+        for aifs_name, era5_name in em.get(cat, {}).items():
+            aifs_to_era5[aifs_name] = era5_name
+    for vv in av.get("pressure_level", []):
+        tpl = em.get("pressure_level", {}).get(vv, "")
+        for lvl in av.get("pressure_levels", []):
+            aifs_to_era5[f"{vv}_{lvl}"] = tpl.format(level=lvl)
+
+    # Accumulators
+    num_samples = 0
+    numerator = np.zeros(num_vars, dtype=np.float64)
+    pred_sq = np.zeros(num_vars, dtype=np.float64)
+    label_sq = np.zeros(num_vars, dtype=np.float64)
+    rmse_sum = np.zeros(num_vars, dtype=np.float64)
+
+    pbar = tqdm(files, desc="Computing metrics", unit="file")
+    for fp in pbar:
+        data = np.load(fp)
+        fname = os.path.splitext(os.path.basename(fp))[0]
+        date_str = fname.replace("aifs_forecast_", "")
+        year = int(date_str[:4])
+        if year not in test_years:
+            data.close()
+            continue
+
+        # Load corresponding ERA5 label
+        h5_path = os.path.join(label_dir, "data", f"{year}.h5")
+        if not os.path.exists(h5_path):
+            data.close()
+            continue
+
+        with h5py.File(h5_path, "r") as hf:
+            ds = hf["fields"]
+            h5_vars = [v.decode() if isinstance(v, bytes) else v
+                       for v in ds.attrs["variables"]]
+            h5_time_step = int(ds.attrs.get("time_step", 6))
+            # Parse date → time index
+            from datetime import datetime
+            dt = datetime.strptime(date_str, "%Y%m%d%H")
+            year_start = datetime(dt.year, 1, 1)
+            hours = (dt - year_start).total_seconds() / 3600
+            t_idx = int(hours / h5_time_step)
+            if t_idx >= ds.shape[0]:
+                data.close()
+                continue
+            label_full = hf["fields"][t_idx]  # (C, 721, 1440)
+
+        # Build per-file H5 channel index
+        h5_ch_map = {}
+        for vi, vname in enumerate(var_names):
+            era5_name = aifs_to_era5.get(vname)
+            if era5_name and era5_name in h5_vars:
+                h5_ch_map[vi] = h5_vars.index(era5_name)
+
+        if not h5_ch_map:
+            data.close()
+            continue
+
+        # Regrid labels to N320
+        label_n320 = {}
+        for vi, h5_ch in h5_ch_map.items():
+            arr_2d = label_full[h5_ch]
+            arr_1d = _interp_to_n320(arr_2d)
+            if arr_1d is not None:
+                label_n320[vi] = arr_1d
+
+        # Compute metrics per variable
+        for vi, label_1d in label_n320.items():
+            pred = data[var_names[vi]]
+
+            # RMSE
+            sq_err = (pred - label_1d) ** 2
+            rmse_sum[vi] += np.sqrt(sq_err.mean())
+
+            # ACC: anomaly correlation (anomaly = deviation from spatial mean)
+            pred_anom = pred - pred.mean()
+            label_anom = label_1d - label_1d.mean()
+
+            numerator[vi] += (pred_anom * label_anom).sum()
+            pred_sq[vi] += (pred_anom ** 2).sum()
+            label_sq[vi] += (label_anom ** 2).sum()
+
+        num_samples += 1
+        data.close()
+
+    if num_samples == 0:
+        raise RuntimeError("No matching label data found")
+
+    rmse = rmse_sum / num_samples
+    denom = np.sqrt(pred_sq * label_sq)
+    denom = np.where(denom > 1e-8, denom, 1.0)
+    acc = numerator / denom
+
+    return var_names, rmse, acc
+
+
+def _interp_to_n320(field_2d: np.ndarray) -> np.ndarray | None:
+    """Interpolate (721, 1440) → N320 (542080,)."""
+    try:
+        import earthkit.regrid as ekr
+        result = ekr.interpolate(
+            field_2d[np.newaxis, ...].astype(np.float64),
+            {"grid": (0.25, 0.25)},
+            {"grid": "N320"},
+        )
+        return result.flatten().astype(np.float64)
+    except Exception:
+        return None
+
+
+# ============================================================================
+# Display
+# ============================================================================
+
+def print_metrics(var_names, rmse, acc):
+    w = 24
+    print(f"\n┌{'─' * (w + 2)}┬{'─' * 14}┬{'─' * 14}┐")
+    print(f"│ {'Variable':<{w}} │ {'RMSE':>12} │ {'ACC':>12} │")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    for i in range(len(var_names)):
+        print(f"│ {var_names[i]:<{w}} │ {rmse[i]:>12.6f} │ {acc[i]:>12.6f} │")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    print(f"│ {'Average':<{w}} │ {np.mean(rmse):>12.6f} │ {np.mean(acc):>12.6f} │")
+    print(f"└{'─' * (w + 2)}┴{'─' * 14}┴{'─' * 14}┘")
+
+
+# ============================================================================
+# Plotting
+# ============================================================================
+
+def plot_field(state, var_name, out_dir):
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    import cartopy.crs as ccrs
+    import cartopy.feature as cfeat
+    import matplotlib.tri as tri
+
+    lats = state.get("latitudes")
+    lons = state.get("longitudes")
+    values = state.get("fields", {}).get(var_name)
+    if lats is None or values is None:
+        return None
+
+    desc, cmap, unit = _VAR_META.get(var_name, (var_name, "viridis", ""))
+
+    fig, ax = plt.subplots(
+        figsize=(11, 6),
+        subplot_kw={"projection": ccrs.PlateCarree()},
+    )
+    ax.coastlines(linewidth=0.5)
+    ax.add_feature(cfeat.BORDERS, linestyle=":", linewidth=0.3)
+    ax.set_global()
+
+    lons_adj = np.where(lons > 180, lons - 360, lons)
+    t = tri.Triangulation(lons_adj, lats)
+    c = ax.tricontourf(t, values, levels=20, transform=ccrs.PlateCarree(),
+                       cmap=cmap)
+    cb = fig.colorbar(c, ax=ax, orientation="vertical", shrink=0.7, pad=0.02)
+    if unit:
+        cb.set_label(unit)
+
+    vs = str(state.get("date", "unknown")).replace("T", " ")[:16]
+    plt.title(f"{desc}  |  valid {vs}", fontsize=12)
+
+    os.makedirs(out_dir, exist_ok=True)
+    path = os.path.join(out_dir, f"{var_name}_{vs.replace(' ', '_').replace(':','')}.png")
+    fig.savefig(path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    return path
+
+
+# ============================================================================
+# Main
+# ============================================================================
+
+def main():
+    p = argparse.ArgumentParser(description="AIFS v1.1 — Evaluate & Visualise")
+    p.add_argument("-c", "--config",
+                   default=str(ROOT / "conf" / "config.yaml"))
+    p.add_argument("-i", "--input_dir", default=None)
+    p.add_argument("-p", "--plot_dir", default=None)
+    p.add_argument("-v", "--variables", default="",
+                   help="comma-separated vars; empty=default, 'all'=all")
+    p.add_argument("-l", "--list", action="store_true",
+                   help="list available vars and exit")
+    p.add_argument("--no-metrics", action="store_true",
+                   help="skip ACC/RMSE (plot only)")
+    args = p.parse_args()
+
+    cfg = _load_config()
+    out_cfg = cfg.get("output", {})
+    data_cfg = cfg.get("data", {})
+
+    in_dir = args.input_dir or str(ROOT / "output")
+    plt_dir = args.plot_dir or str(ROOT / "plots")
+    test_years = data_cfg.get("test_years", [2008])
+
+    files = sorted(glob.glob(os.path.join(in_dir, "aifs_forecast_*.npz")))
+    if not files:
+        print(f"[ERROR] No aifs_forecast_*.npz in {in_dir}")
+        sys.exit(1)
+
+    s = np.load(files[0])
+    avail = [k for k in s.keys()
+             if k not in ("latitudes", "longitudes", "date")]
+    s.close()
+
+    if args.list:
+        print(f"\nAvailable variables ({len(avail)}):")
+        for v in avail:
+            d, _, u = _VAR_META.get(v, (v, "", ""))
+            print(f"  {v:12s}  {d}")
+        return
+
+    if args.variables == "":
+        vars_ = [v for v in DEFAULT_VARS if v in avail]
+    elif args.variables == "all":
+        vars_ = avail
+    else:
+        vars_ = [v.strip() for v in args.variables.split(",")
+                 if v.strip() in avail]
+
+    # ---- ACC / RMSE ----
+    if not args.no_metrics:
+        label_dir = data_cfg.get("data_dir", "")
+        if label_dir and os.path.isdir(os.path.join(label_dir, "data")):
+            try:
+                var_names, rmse, acc = compute_metrics(
+                    in_dir, label_dir, test_years,
+                )
+                print_metrics(var_names, rmse, acc)
+                # Save metrics
+                metrics_dir = ROOT / "metrics"
+                os.makedirs(str(metrics_dir), exist_ok=True)
+                np.save(str(metrics_dir / "rmse.npy"), rmse)
+                np.save(str(metrics_dir / "acc.npy"), acc)
+                with open(str(metrics_dir / "metrics.txt"), "w") as f:
+                    f.write(f"{'Variable':<24s} {'RMSE':>12s} {'ACC':>12s}\n")
+                    for i, v in enumerate(var_names):
+                        f.write(f"{v:<24s} {rmse[i]:>12.6f} {acc[i]:>12.6f}\n")
+                print(f"[INFO] Metrics saved to metrics/")
+            except Exception as e:
+                print(f"[WARN] ACC/RMSE skipped: {e}")
+        else:
+            print("[INFO] No label data found — skipping ACC/RMSE")
+
+    # ---- Plotting ----
+    print(f"[INFO] {len(files)} file(s), {len(vars_)} var(s): {', '.join(vars_[:6])}")
+    total = 0
+    for fp in files:
+        d = np.load(fp)
+        st = dict(
+            date=str(d.get("date", "unknown")),
+            latitudes=d["latitudes"],
+            longitudes=d["longitudes"],
+            fields={k: d[k] for k in avail if k in d},
+        )
+        for v in vars_:
+            out_path = plot_field(st, v, plt_dir)
+            if out_path:
+                total += 1
+        d.close()
+
+    if total > 0:
+        print(f"[INFO] Done: {total} plot(s) → {plt_dir}")
+        # Print only first few to avoid flooding
+    else:
+        print("[WARN] No plots generated")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,894 @@
+#!/usr/bin/env python
+# coding: utf-8
+"""
+AIFS v1.1 — Pre-training Script
+=================================
+Training logic strictly follows config_pretraining.yaml and anemoi-training 0.4.0.
+
+Usage:
+    python scripts/train.py
+    python scripts/train.py --config conf/config.yaml
+"""
+
+from __future__ import annotations
+
+import argparse
+import datetime
+import math
+import os
+import sys
+import time
+import traceback
+import warnings
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import pytz
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+import yaml
+
+# Project root (scripts/ → aifs_v11/)
+ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(ROOT))
+
+import onescience.datapipes.climate.era5 as onescience_era5
+import earthkit.regrid as ekr
+
+from model.aifs import AIFS
+
+# ============================================================================
+# Config loader
+# ============================================================================
+
+def load_config(path: str) -> dict:
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+# ============================================================================
+# Variable list builder
+# ============================================================================
+
+def build_era5_variable_list(cfg: dict) -> List[str]:
+    av = cfg["aifs_variables"]
+    em = cfg["era5_mapping"]
+    vars_: List[str] = []
+    for name in av["surface"]:
+        vars_.append(em["surface"][name])
+    for name in av["soil"]:
+        vars_.append(em["soil"][name])
+    for v in av["pressure_level"]:
+        tpl = em["pressure_level"][v]
+        for lvl in av["pressure_levels"]:
+            vars_.append(tpl.format(level=lvl))
+    for name in av["diagnostic"]:
+        vars_.append(em["diagnostic"][name])
+    return vars_
+
+
+def _auto_correct_variable_names(required: List[str], available: set, cfg: dict) -> List[str]:
+    fuzzy = cfg["era5_mapping"].get("fuzzy_fixes", {})
+    corrected = list(required)
+    for i, v in enumerate(corrected):
+        if v not in available and v in fuzzy:
+            alt = fuzzy[v]
+            if alt in available:
+                print(f"[INFO] Auto-corrected: '{v}' → '{alt}'")
+                corrected[i] = alt
+    return corrected
+
+
+def _era5_to_aifs_name(era5_name: str, cfg: dict) -> Optional[str]:
+    em = cfg["era5_mapping"]
+    for aifs, ename in em["surface"].items():
+        if ename == era5_name: return aifs
+    for aifs, ename in em["soil"].items():
+        if ename == era5_name: return aifs
+    for v in cfg["aifs_variables"]["pressure_level"]:
+        tpl = em["pressure_level"][v]
+        for lvl in cfg["aifs_variables"]["pressure_levels"]:
+            if tpl.format(level=lvl) == era5_name: return f"{v}_{lvl}"
+    for aifs, ename in em["diagnostic"].items():
+        if ename == era5_name: return aifs
+    return None
+
+# ============================================================================
+# N320 interpolation
+# ============================================================================
+
+def _interp_n320(arr: np.ndarray) -> np.ndarray:
+    return ekr.interpolate(arr, {"grid": (0.25, 0.25)}, {"grid": "N320"})
+
+# ============================================================================
+# ERA5 → AIFS field conversion
+# ============================================================================
+
+def build_aifs_fields_from_frames(
+    frame_tm6: np.ndarray, frame_t0: np.ndarray, frame_tp6: np.ndarray,
+    era5_var_list: List[str], cfg: dict,
+    diag_map: Optional[Dict[str, str]] = None,
+) -> Dict[str, np.ndarray]:
+    av = cfg["aifs_variables"]
+    em = cfg["era5_mapping"]
+    name_to_ch = {n: i for i, n in enumerate(era5_var_list)}
+    fields: Dict[str, np.ndarray] = {}
+
+    for aifs_name in av["surface"]:
+        ch = name_to_ch[em["surface"][aifs_name]]
+        fields[aifs_name] = np.stack([_interp_n320(frame_tm6[ch]),
+                                       _interp_n320(frame_t0[ch]),
+                                       _interp_n320(frame_tp6[ch])])
+    for aifs_name in av["soil"]:
+        ch = name_to_ch[em["soil"][aifs_name]]
+        fields[aifs_name] = np.stack([_interp_n320(frame_tm6[ch]),
+                                       _interp_n320(frame_t0[ch]),
+                                       _interp_n320(frame_tp6[ch])])
+    for aifs_var in av["pressure_level"]:
+        tpl = em["pressure_level"][aifs_var]
+        for level in av["pressure_levels"]:
+            ch = name_to_ch[tpl.format(level=level)]
+            fields[f"{aifs_var}_{level}"] = np.stack([_interp_n320(frame_tm6[ch]),
+                                                       _interp_n320(frame_t0[ch]),
+                                                       _interp_n320(frame_tp6[ch])])
+    _diag_map = diag_map if diag_map is not None else em["diagnostic"]
+    for aifs_name in av["diagnostic"]:
+        era5_name = _diag_map.get(aifs_name)
+        if era5_name and era5_name in name_to_ch:
+            ch = name_to_ch[era5_name]
+            fields[aifs_name] = np.stack([_interp_n320(frame_tm6[ch]),
+                                           _interp_n320(frame_t0[ch]),
+                                           _interp_n320(frame_tp6[ch])])
+    return fields
+
+# ============================================================================
+# Forcing feature computation
+# ============================================================================
+
+def compute_forcing_features(
+    latitudes: np.ndarray, longitudes: np.ndarray,
+    timestamps: List[datetime.datetime],
+    fields: Dict[str, np.ndarray],
+) -> Dict[str, np.ndarray]:
+    lat_rad = np.deg2rad(latitudes)
+    lon_rad = np.deg2rad(longitudes)
+    multi_step = len(timestamps)
+    forcing: Dict[str, np.ndarray] = {}
+
+    forcing["cos_latitude"] = np.tile(np.cos(lat_rad)[np.newaxis, :], (multi_step, 1))
+    forcing["sin_latitude"] = np.tile(np.sin(lat_rad)[np.newaxis, :], (multi_step, 1))
+    forcing["cos_longitude"] = np.tile(np.cos(lon_rad)[np.newaxis, :], (multi_step, 1))
+    forcing["sin_longitude"] = np.tile(np.sin(lon_rad)[np.newaxis, :], (multi_step, 1))
+
+    cos_jd = np.zeros((multi_step, len(latitudes)), dtype=np.float32)
+    sin_jd = np.zeros_like(cos_jd)
+    cos_lt = np.zeros_like(cos_jd)
+    sin_lt = np.zeros_like(cos_jd)
+    insol = np.zeros_like(cos_jd)
+
+    for t_idx, ts in enumerate(timestamps):
+        doy = ts.timetuple().tm_yday
+        jd_angle = 2.0 * np.pi * doy / 365.25
+        cos_jd[t_idx] = np.cos(jd_angle)
+        sin_jd[t_idx] = np.sin(jd_angle)
+        hours = ts.hour + ts.minute / 60.0 + ts.second / 3600.0
+        lt_angle = 2.0 * np.pi * hours / 24.0 + lon_rad
+        cos_lt[t_idx] = np.cos(lt_angle)
+        sin_lt[t_idx] = np.sin(lt_angle)
+        insol[t_idx] = _compute_insolation(ts, lat_rad, lon_rad)
+
+    forcing["cos_julian_day"] = cos_jd
+    forcing["sin_julian_day"] = sin_jd
+    forcing["cos_local_time"] = cos_lt
+    forcing["sin_local_time"] = sin_lt
+    forcing["insolation"] = insol
+
+    for var_name in ["lsm", "z", "slor", "sdor"]:
+        if var_name in fields:
+            forcing[var_name] = np.tile(fields[var_name][:1], (multi_step, 1))
+    return forcing
+
+
+def _compute_insolation(ts: datetime.datetime, lat_rad: np.ndarray, lon_rad: np.ndarray) -> np.ndarray:
+    ref = datetime.datetime(2000, 1, 1, 12, 0, 0, tzinfo=pytz.utc)
+    mt = ts.timestamp()
+    days = (mt - ref.timestamp()) / (24.0 * 3600.0)
+    jc = days / 36525.0
+    theta = (67310.54841 + jc * (876600.0 * 3600.0 + 8640184.812866
+             + jc * (0.093104 - jc * 6.2e-5)))
+    gmst = np.fmod((theta / 240.0) * np.pi / 180.0, 2.0 * np.pi)
+    ma = np.deg2rad(357.52910 + 35999.05030 * jc - 0.0001559 * jc**2 - 0.00000048 * jc**3)
+    ml = np.deg2rad(280.46645 + 36000.76983 * jc + 0.0003032 * jc**2)
+    dl = np.deg2rad((1.914600 - 0.004817 * jc - 0.000014 * jc**2) * np.sin(ma)
+                     + (0.019993 - 0.000101 * jc) * np.sin(2.0 * ma)
+                     + 0.000290 * np.sin(3.0 * ma))
+    tl = ml + dl
+    eps = np.deg2rad(23.0 + 26.0 / 60.0 + 21.406 / 3600.0
+                     - (46.836769 * jc - 0.0001831 * jc**2 + 0.00200340 * jc**3
+                        - 0.576e-6 * jc**4 - 4.34e-8 * jc**5) / 3600.0)
+    x = np.cos(tl); y = np.cos(eps) * np.sin(tl); z = np.sin(eps) * np.sin(tl)
+    r = np.sqrt(1.0 - z * z)
+    dec = np.arctan2(z, r)
+    ra = 2.0 * np.arctan2(y, x + r)
+    ha = (gmst + lon_rad) - ra
+    cos_z = np.sin(lat_rad) * np.sin(dec) + np.cos(lat_rad) * np.cos(dec) * np.cos(ha)
+    return cos_z.astype(np.float32)
+
+# ============================================================================
+# Normalisation
+# ============================================================================
+
+def compute_normalisation_params(
+    cfg: dict, variable_names: List[str],
+    statistics_path: Optional[str] = None,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute per-variable norm_mul and norm_add from config.
+
+    If ``statistics_path`` points to a ``.npz`` file with keys
+    ``mean`` and ``stdev`` (per-variable arrays matching the full
+    dataset variable order), those values are used.  Otherwise
+    identity (mean=0, stdev=1) is assumed — correct for fake data
+    but NOT for real ERA5.
+    """
+    nc = cfg["normalizer"]
+    default_method = nc["default"]
+    remap = nc.get("remap", {}) or {}
+    methods: Dict[str, str] = {}
+    for v in variable_names:
+        if v in (nc.get("none") or []): methods[v] = "none"
+        elif v in (nc.get("max") or []): methods[v] = "max"
+        elif v in (nc.get("min-max") or []): methods[v] = "min-max"
+        elif v in (nc.get("std") or []): methods[v] = "std"
+        else: methods[v] = default_method
+
+    name_to_idx = {n: i for i, n in enumerate(variable_names)}
+    num_vars = len(variable_names)
+
+    # Load pre-computed statistics if available, else identity
+    _mean = np.zeros(num_vars, dtype=np.float32)
+    _stdev = np.ones(num_vars, dtype=np.float32)
+
+    if statistics_path and os.path.exists(statistics_path):
+        stats = np.load(statistics_path)
+        # statistics file stores arrays indexed by its own variable list
+        if "variables" in stats:
+            ds_vars = [str(v) for v in stats["variables"]]
+        else:
+            # Fallback: assume 115-variable dataset order from aifs_config
+            import json
+            cfg_path = ROOT / "model" / "aifs_config.json"
+            ds_vars = json.load(open(cfg_path))["dataset"]["variables"]
+        ds_name_to_idx = {n: i for i, n in enumerate(ds_vars)}
+        for name, i in name_to_idx.items():
+            if name in ds_name_to_idx:
+                j = ds_name_to_idx[name]
+                _mean[i] = float(stats["mean"][j])
+                _stdev[i] = float(stats["stdev"][j])
+        print(f"[INFO] Loaded statistics from {statistics_path}")
+    else:
+        print("[INFO] No statistics file — using identity normalisation "
+              "(mean=0, stdev=1).  Fine for fake data, "
+              "WRONG for real ERA5.")
+
+    for target_var, source_var in remap.items():
+        if target_var in name_to_idx and source_var in name_to_idx:
+            ti, si = name_to_idx[target_var], name_to_idx[source_var]
+            _mean[ti], _stdev[ti] = _mean[si], _stdev[si]
+
+    norm_mul = np.ones(num_vars, dtype=np.float32)
+    norm_add = np.zeros(num_vars, dtype=np.float32)
+    for name, i in name_to_idx.items():
+        method = methods.get(name, default_method)
+        if method == "mean-std":
+            norm_mul[i] = 1.0 / max(float(_stdev[i]), 1e-9)
+            norm_add[i] = -float(_mean[i]) / max(float(_stdev[i]), 1e-9)
+        elif method == "std":
+            norm_mul[i] = 1.0 / max(float(_stdev[i]), 1e-9)
+        elif method in ("max", "min-max", "none"):
+            norm_mul[i] = 1.0
+    return norm_mul, norm_add
+
+
+def compute_era5_statistics(
+    data_dir: str, years: List[int],
+    variable_names: List[str],
+    output_path: str,
+) -> None:
+    """Compute per-variable mean and stdev from ERA5 H5 files.
+
+    Saves ``mean`` and ``stdev`` arrays (indexed by *variable_names*
+    order) to *output_path* as a ``.npz`` file.
+
+    This is a one-time pre-computation step before training with
+    real ERA5 data.  The resulting file is consumed by
+    ``compute_normalisation_params(statistics_path=...)``.
+    """
+    import h5py
+
+    num_vars = len(variable_names)
+    sum_x = np.zeros(num_vars, dtype=np.float64)
+    sum_x2 = np.zeros(num_vars, dtype=np.float64)
+    count = 0
+
+    for year in years:
+        h5_path = os.path.join(data_dir, "data", f"{year}.h5")
+        if not os.path.exists(h5_path):
+            print(f"[WARN] Missing {h5_path}, skipping")
+            continue
+
+        with h5py.File(h5_path, "r") as f:
+            ds = f["fields"]
+            h5_vars = [v.decode() if isinstance(v, bytes) else v
+                       for v in ds.attrs["variables"]]
+            T = ds.shape[0]
+
+            # Index mapping for each timestep
+            for t in range(T):
+                frame = ds[t]  # (C, H, W)
+                for i, vname in enumerate(variable_names):
+                    if vname in h5_vars:
+                        ch = h5_vars.index(vname)
+                        vals = frame[ch].astype(np.float64)
+                        sum_x[i] += vals.mean()
+                        sum_x2[i] += (vals ** 2).mean()
+                count += 1
+
+    mean = (sum_x / count).astype(np.float32)
+    # stdev = sqrt(E[X²] - E[X]²)
+    var = np.maximum(sum_x2 / count - mean.astype(np.float64) ** 2, 0)
+    stdev = np.sqrt(var).astype(np.float32)
+
+    np.savez(output_path, mean=mean, stdev=stdev)
+    print(f"[INFO] Statistics saved to {output_path} "
+          f"({count} frames, {num_vars} vars)")
+
+# ============================================================================
+# Loss scaling
+# ============================================================================
+
+def build_loss_scale_factors(cfg: dict, output_var_names: List[str]) -> np.ndarray:
+    sc = cfg["loss_scaling"]
+    default = sc["default"]; pl = sc["pl"]; sfc = sc["sfc"]
+    av = cfg["aifs_variables"]
+    name_to_scale: Dict[str, float] = {}
+    for var_short in av["pressure_level"]:
+        base = pl.get(var_short, default)
+        for lvl in av["pressure_levels"]:
+            name_to_scale[f"{var_short}_{lvl}"] = base
+    for var_name in av["surface"] + av["soil"] + av["diagnostic"]:
+        name_to_scale[var_name] = sfc.get(var_name, default)
+    for k, v in sfc.items():
+        if k not in name_to_scale: name_to_scale[k] = v
+    return np.array([name_to_scale.get(n, default) for n in output_var_names], dtype=np.float32)
+
+
+def build_pressure_level_scaler(cfg: dict, output_var_names: List[str]) -> np.ndarray:
+    pls = cfg["loss_scaling"]["pressure_level_scaler"]
+    minimum, slope = pls["minimum"], pls["slope"]
+    av = cfg["aifs_variables"]
+    factors = np.ones(len(output_var_names), dtype=np.float32)
+    for i, name in enumerate(output_var_names):
+        for aifs_var in av["pressure_level"]:
+            if name.startswith(f"{aifs_var}_"):
+                plev_hpa = int(name[len(aifs_var) + 1:])
+                factors[i] = max(minimum, slope * plev_hpa)
+                break
+    return factors
+
+# ============================================================================
+# Dataset
+# ============================================================================
+
+class AIFSTrainingDataset:
+    def __init__(self, dataset_path: str, years: List[int], cfg: dict,
+                 model_input_order: List[str], model_output_order: List[str],
+                 norm_mul: np.ndarray, norm_add: np.ndarray,
+                 input_steps=2, output_steps=1,
+                 grid_lat=None, grid_lon=None):
+        dp = dataset_path.rstrip("/")
+        if os.path.isdir(os.path.join(dp, "data")): era5_root = dp
+        elif os.path.isdir(dp) and any(f.endswith(".h5") for f in os.listdir(dp)): era5_root = os.path.dirname(dp)
+        else: era5_root = dp
+        self._era5_root = era5_root
+
+        data_dir = os.path.join(era5_root, "data")
+        h5_files = sorted([f for f in os.listdir(data_dir) if f.endswith(".h5")])
+        era5_var_list = build_era5_variable_list(cfg)
+        available_set: set = set()
+
+        if h5_files:
+            import h5py
+            with h5py.File(os.path.join(data_dir, h5_files[0]), "r") as f:
+                available_vars = [v.decode() if isinstance(v, bytes) else v for v in f["fields"].attrs["variables"]]
+            available_set = set(available_vars)
+            corrected = _auto_correct_variable_names(era5_var_list, available_set, cfg)
+            used_vars = [v for v in corrected if v in available_set]
+            missing = [v for v in corrected if v not in available_set]
+            if missing:
+                print(f"[WARN] {len(missing)}/{len(era5_var_list)} vars missing from H5: {missing[:10]}")
+        else:
+            used_vars = era5_var_list
+
+        self._cfg = cfg
+        self.era5_var_list = used_vars
+        self.model_input_order = model_input_order
+        self.model_output_order = model_output_order
+        self.norm_mul = norm_mul; self.norm_add = norm_add
+        self.input_steps = input_steps; self.output_steps = output_steps
+        self._name_to_input_idx = {n: i for i, n in enumerate(model_input_order)}
+        self._name_to_output_idx = {n: i for i, n in enumerate(model_output_order)}
+
+        # corrected diagnostic map
+        self._diag_map: Dict[str, str] = {}
+        for aifs_name, orig in cfg["era5_mapping"]["diagnostic"].items():
+            cn = _auto_correct_variable_names([orig], available_set, cfg)[0] if h5_files else orig
+            self._diag_map[aifs_name] = cn if (h5_files and cn in available_set) else orig
+
+        self._dataset = onescience_era5.ERA5Dataset(
+            dataset_dir=era5_root, used_years=years,
+            used_variables=used_vars, input_steps=input_steps,
+            output_steps=output_steps, normalize=False)
+
+        self._grid_lat = grid_lat; self._grid_lon = grid_lon
+        self.n_samples = len(self._dataset)
+
+        # audit
+        self._print_audit(cfg, used_vars)
+
+    def _print_audit(self, cfg, used_vars):
+        av = cfg["aifs_variables"]; em = cfg["era5_mapping"]
+        fields_ok = set()
+        for aifs in av["surface"]:
+            if em["surface"][aifs] in used_vars: fields_ok.add(aifs)
+        for aifs in av["soil"]:
+            if em["soil"][aifs] in used_vars: fields_ok.add(aifs)
+        for vv in av["pressure_level"]:
+            for lvl in av["pressure_levels"]:
+                if em["pressure_level"][vv].format(level=lvl) in used_vars: fields_ok.add(f"{vv}_{lvl}")
+        for aifs in av["diagnostic"]:
+            if self._diag_map[aifs] in used_vars: fields_ok.add(aifs)
+        forcing_ok = set(av["computed_forcing"] + ["lsm", "z", "slor", "sdor"])
+        mi = [v for v in self.model_input_order if v not in fields_ok and v not in forcing_ok]
+        mo = [v for v in self.model_output_order if v not in fields_ok]
+        print(f"[INFO] Variables: {len(fields_ok)} fields + {len(forcing_ok)} forcing = all available")
+        if mi: print(f"[WARN] {len(mi)} input vars zero-filled: {mi}")
+        if mo: print(f"[WARN] {len(mo)} output vars no target data: {mo}")
+
+    def __len__(self): return self.n_samples
+    def set_grid(self, lat, lon):
+        self._grid_lat = lat.astype(np.float32); self._grid_lon = lon.astype(np.float32)
+
+    def _get_raw_sample(self, idx: int):
+        result = self._dataset[int(idx)]
+        return result[0], result[1], int(result[3]), list(result[4])
+
+    def process_sample(self, invar, outvar, step_idx, time_index):
+        if invar.ndim == 3: invar = invar[np.newaxis, ...]
+        if outvar.ndim == 3: outvar = outvar[np.newaxis, ...]
+        total_steps = self.input_steps + self.output_steps
+        all_frames = np.concatenate([invar, outvar], axis=0)
+        timestamps = [datetime.datetime.strptime(t, "%Y%m%d%H").replace(tzinfo=pytz.utc) for t in time_index]
+
+        frame_tm6, frame_t0, frame_tp6 = all_frames[0], all_frames[1], all_frames[2] if total_steps >= 3 else all_frames[-1]
+        fields = build_aifs_fields_from_frames(frame_tm6, frame_t0, frame_tp6, self.era5_var_list, self._cfg, diag_map=self._diag_map)
+
+        num_pts = next(iter(fields.values())).shape[-1]
+        if self._grid_lat is None:
+            self._grid_lat = np.linspace(90, -90, 542080, dtype=np.float32)[:num_pts]
+            self._grid_lon = np.linspace(0, 360, 542080, dtype=np.float32)[:num_pts]
+
+        forcing = compute_forcing_features(self._grid_lat, self._grid_lon, timestamps, fields)
+
+        num_input_vars = len(self.model_input_order)
+        input_tensor = np.zeros((self.input_steps, num_pts, num_input_vars), dtype=np.float32)
+        for var_name, var_idx in self._name_to_input_idx.items():
+            for t in range(self.input_steps):
+                if var_name in fields: input_tensor[t, :, var_idx] = fields[var_name][t]
+                elif var_name in forcing: input_tensor[t, :, var_idx] = forcing[var_name][t]
+
+        input_tensor = input_tensor * self.norm_mul[np.newaxis, np.newaxis, :] + self.norm_add[np.newaxis, np.newaxis, :]
+
+        num_output_vars = len(self.model_output_order)
+        target_tensor = np.zeros((1, num_pts, num_output_vars), dtype=np.float32)
+        t_target = self.input_steps
+        for var_name, var_idx in self._name_to_output_idx.items():
+            if var_name in fields: target_tensor[0, :, var_idx] = fields[var_name][t_target]
+        return input_tensor, target_tensor
+
+# ============================================================================
+# Model loading
+# ============================================================================
+
+def load_aifs_model(checkpoint_path: str, device: str, pretrained: bool = False):
+    """Load the AIFS model through the onescience wrapper.
+
+    Parameters
+    ----------
+    pretrained : bool
+        False (default): Build model from scratch via ``AIFS.from_scratch()``.
+        Uses local static config + grid files — NO checkpoint needed.
+        Equivalent to FengWu/Fuxi ``model = Fengwu()`` pattern.
+        True: Load serialised model with pretrained weights from .ckpt.
+    """
+    t0 = time.time()
+    if pretrained:
+        model = AIFS(checkpoint_path, device=device, pretrained=True)
+    else:
+        model = AIFS.from_scratch(device=device)
+    mode = "pretrained" if pretrained else "from scratch"
+    print(f"[INFO] Model loaded ({mode}, {time.time() - t0:.1f}s)")
+
+    input_vars  = model.input_variables
+    output_vars = model.output_variables
+    grid_lat    = model.latitudes.cpu().numpy()
+    grid_lon    = model.longitudes.cpu().numpy()
+    node_wts    = model.node_weights   # np.ndarray or None
+    return model, input_vars, output_vars, grid_lat, grid_lon, node_wts
+
+# ============================================================================
+# Loss
+# ============================================================================
+
+class WeightedMSELoss(nn.Module):
+    def __init__(self, var_scale: np.ndarray, pl_scale: np.ndarray, node_weights=None):
+        super().__init__()
+        self.register_buffer("combined_scale", torch.from_numpy((var_scale * pl_scale).astype(np.float32)))
+        if node_weights is not None:
+            self.register_buffer("node_weights", torch.from_numpy(node_weights.astype(np.float32)))
+        else: self.node_weights = None
+
+    def forward(self, pred, target):
+        sq_error = (pred - target) ** 2
+        scaled = sq_error * self.combined_scale[None, None, :]
+        if self.node_weights is not None:
+            x = scaled.mean(dim=-1) * self.node_weights[None, :]
+            return (x / self.node_weights.sum()).sum()
+        return scaled.mean()
+
+# ============================================================================
+# LR schedule
+# ============================================================================
+
+def cosine_lr_schedule(step, warmup_steps, total_steps, peak_lr, min_lr):
+    """Cosine LR with linear warmup.
+
+    Manual implementation equivalent to ``timm.scheduler.CosineLRScheduler``
+    for a single decay cycle (no restarts).  Same formula as the official
+    anemoi-training configuration.
+    """
+    if step < warmup_steps:
+        return peak_lr * (step / max(warmup_steps, 1))
+    progress = min(
+        (step - warmup_steps) / max(total_steps - warmup_steps, 1), 1.0,
+    )
+    return min_lr + (peak_lr - min_lr) * 0.5 * (1.0 + math.cos(math.pi * progress))
+
+# ============================================================================
+# Training
+# ============================================================================
+
+def train(cfg: dict):
+    hw = cfg["hardware"]; data = cfg["data"]; ck = cfg["checkpoint"]; tr = cfg["training"]
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if hw["device"] == "cpu": device = "cpu"
+    print(f"[INFO] Device: {device}")
+
+    train_years = data["train_years"]
+    val_years = data["val_years"]
+
+    from_scratch = tr.get("from_scratch", True)
+    model, model_input_vars, model_output_vars, grid_lat, grid_lon, node_wts = \
+        load_aifs_model(ck.get("pretrained", ""), device,
+                        pretrained=not from_scratch)
+
+    # ---- 自动计算归一化统计量（如需要）-------------------------------
+    stats_path = cfg["normalizer"].get("statistics_path") or None
+    if stats_path:
+        stats_path = os.path.join(str(ROOT), stats_path)
+        if not os.path.exists(stats_path):
+            # 自动计算 —— 对标 anemoi-training DataModule 的统计量阶段
+            print("[INFO] Statistics file not found — auto-computing ...")
+            ds_vars = None
+            if hasattr(model, '_meta') and 'dataset' in model._meta:
+                ds_vars = model._meta['dataset']['variables']
+            elif os.path.exists(str(ROOT / "model" / "aifs_config.json")):
+                import json
+                ds_vars = json.load(
+                    open(str(ROOT / "model" / "aifs_config.json"))
+                )['dataset']['variables']
+            if ds_vars is None:
+                print("[WARN] Cannot determine dataset variables — "
+                      "falling back to identity normalisation")
+                stats_path = None
+            else:
+                os.makedirs(os.path.dirname(stats_path), exist_ok=True)
+                compute_era5_statistics(
+                    data["data_dir"], train_years + val_years,
+                    ds_vars, stats_path,
+                )
+
+    norm_mul, norm_add = compute_normalisation_params(
+        cfg, model_input_vars,
+        statistics_path=stats_path,
+    )
+
+    var_scale = build_loss_scale_factors(cfg, model_output_vars)
+    pl_scale = build_pressure_level_scaler(cfg, model_output_vars)
+
+    criterion = WeightedMSELoss(var_scale, pl_scale, node_wts).to(device)
+    val_criterion = WeightedMSELoss(
+        np.ones(len(model_output_vars), dtype=np.float32),
+        np.ones(len(model_output_vars), dtype=np.float32), node_wts).to(device)
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=0.0,
+        betas=tr["optimizer"]["betas"],
+        weight_decay=tr["optimizer"].get("weight_decay", 0.01),
+    )
+
+    train_dataset = AIFSTrainingDataset(
+        data["data_dir"], train_years, cfg, model_input_vars, model_output_vars,
+        norm_mul, norm_add, input_steps=2, output_steps=1,
+        grid_lat=grid_lat, grid_lon=grid_lon)
+    train_dataset.set_grid(grid_lat, grid_lon)
+    val_dataset = AIFSTrainingDataset(
+        data["data_dir"], val_years, cfg, model_input_vars, model_output_vars,
+        norm_mul, norm_add, input_steps=2, output_steps=1,
+        grid_lat=grid_lat, grid_lon=grid_lon)
+    val_dataset.set_grid(grid_lat, grid_lon)
+
+    os.makedirs(ck["output_dir"], exist_ok=True)
+    use_amp = (device == "cuda")
+    amp_dtype = getattr(torch, tr.get("amp_dtype", "float16"))
+    scaler = torch.amp.GradScaler("cuda") if use_amp else None
+
+    model.train()
+    global_step, epoch, best_val = 0, 0, float("inf")
+    t_start = time.time()
+    pbar = tqdm(total=tr["max_steps"], desc="Training", unit="step", dynamic_ncols=True)
+
+    while global_step < tr["max_steps"]:
+        epoch += 1
+        indices = np.random.permutation(len(train_dataset))
+        for batch_start in range(0, len(train_dataset), tr["batch_size"]):
+            batch_indices = indices[batch_start:batch_start + tr["batch_size"]]
+            batch_inputs, batch_targets = [], []
+            for idx in batch_indices:
+                try:
+                    invar, outvar, si, ti = train_dataset._get_raw_sample(idx)
+                    invar_np = invar.numpy() if hasattr(invar, "numpy") else np.asarray(invar)
+                    outvar_np = outvar.numpy() if hasattr(outvar, "numpy") else np.asarray(outvar)
+                    inp, tgt = train_dataset.process_sample(invar_np, outvar_np, si, ti)
+                    batch_inputs.append(inp); batch_targets.append(tgt)
+                except Exception as e:
+                    print(f"[WARN] Skip sample {idx}: {e}")
+            if not batch_inputs: continue
+
+            x = torch.from_numpy(np.stack(batch_inputs, axis=0)).unsqueeze(2).to(device)
+            y = torch.from_numpy(np.stack(batch_targets, axis=0)).squeeze(1).to(device)
+
+            with (torch.amp.autocast("cuda", dtype=amp_dtype) if use_amp
+                  else torch.no_grad()):
+                pred = model(x)                 # AIFS.forward already squeezes dim
+
+            loss = criterion(pred, y) / tr["accum_grad_batches"]
+            (scaler.scale(loss) if scaler else loss).backward()
+
+            # optimizer step
+            if (batch_start // tr["batch_size"] + 1) % tr["accum_grad_batches"] == 0:
+                if scaler: scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_value_(model.parameters(), tr["gradient_clip_val"])
+                (scaler.step(optimizer) if scaler else optimizer.step())
+                if scaler: scaler.update()
+                optimizer.zero_grad()
+                global_step += 1
+
+                lr = cosine_lr_schedule(global_step, tr["warmup_steps"], tr["max_steps"], tr["peak_lr"], tr["min_lr"])
+                for pg in optimizer.param_groups: pg["lr"] = lr
+
+                elapsed = time.time() - t_start
+                pbar.set_postfix({"loss": f"{loss.item():.4f}", "lr": f"{lr:.2e}", "stp/s": f"{global_step/max(elapsed,1):.1f}"})
+                pbar.update(1)
+
+                if global_step % tr["val_interval"] == 0:
+                    vl = validate(model, val_dataset, val_criterion, device)
+                    pbar.write(f"[step {global_step:06d}] VAL loss={vl:.6f}")
+                    model.train()
+                    if vl < best_val:
+                        best_val = vl
+                        _save(model, optimizer, global_step, epoch, vl,
+                              ck["output_dir"], "best")
+                if global_step % tr["save_interval"] == 0:
+                    _save(model, optimizer, global_step, epoch, loss.item(),
+                          ck["output_dir"], f"step{global_step:06d}")
+                if global_step >= tr["max_steps"]: break
+
+    pbar.close()
+    _save(model, optimizer, global_step, epoch, loss.item(),
+          ck["output_dir"], "final")
+    print(f"\n[INFO] Done: {global_step} steps in {(time.time()-t_start)/60:.1f}min, best_val={best_val:.6f}")
+
+@torch.no_grad()
+def validate(model, dataset, criterion, device):
+    model.eval()
+    total, n = 0.0, 0
+    use_amp = (device == "cuda")
+    amp_dtype_val = getattr(torch, "float16")
+    _first_err = True
+    for idx in range(len(dataset)):
+        try:
+            invar, outvar, si, ti = dataset._get_raw_sample(idx)
+            invar_np = invar.numpy() if hasattr(invar, "numpy") else np.asarray(invar)
+            outvar_np = outvar.numpy() if hasattr(outvar, "numpy") else np.asarray(outvar)
+            inp, tgt = dataset.process_sample(invar_np, outvar_np, si, ti)
+            x = torch.from_numpy(inp).unsqueeze(0).unsqueeze(2).to(device)
+            y = torch.from_numpy(tgt).squeeze(1).unsqueeze(0).to(device)
+            with (torch.amp.autocast("cuda", dtype=amp_dtype_val) if use_amp
+                  else torch.no_grad()):
+                pred = model(x)                 # AIFS.forward already squeezes dim
+            total += criterion(pred, y).item(); n += 1
+        except Exception as e:
+            if _first_err:
+                print(f"[WARN] Validation error on sample {idx}: {e}")
+                _first_err = False
+    return total / max(n, 1)
+
+def _save(model, optimizer, step, epoch, loss, out_dir, tag):
+    """Save training checkpoint — 推理+续训合一，零外部依赖.
+
+    单个 .ckpt 文件包含:
+    - 完整 AnemoiModelEncProcDec (torch.load 直接得模型 → SimpleRunner 可用)
+    - ai-models.json + supporting arrays (inference metadata)
+    - optimizer.pkl (恢复训练用，推理时忽略)
+    """
+    import pickle, zipfile
+
+    path = os.path.join(out_dir, f"model_bak.ckpt")
+
+    # ---- 1.  构建 AnemoiModelInterface（含 predict_step）---------
+    from anemoi.models.interface import AnemoiModelInterface
+    import numpy as np
+
+    # 虚拟统计量（identity normalization：mean=0, stdev=1）
+    num_vars = len(model._interface_data_indices.data.input.name_to_index)
+    dummy_stats = {
+        "minimum": np.zeros(num_vars, dtype=np.float32),
+        "maximum": np.ones(num_vars, dtype=np.float32),
+        "mean": np.zeros(num_vars, dtype=np.float32),
+        "stdev": np.ones(num_vars, dtype=np.float32),
+    }
+
+    interface = AnemoiModelInterface(
+        config=model._interface_config,
+        graph_data=model._interface_graph_data,
+        statistics=dummy_stats,
+        data_indices=model._interface_data_indices,
+        metadata={},
+        truncation_data=getattr(model._model, '_truncation_data', {}),
+    )
+    interface.model.load_state_dict(model._model.state_dict())
+    interface.cpu()
+    torch.save(interface, path)
+    del interface
+
+    # ---- 2.  注入 metadata (从本地静态文件) ---------------------------
+    _inject_metadata(path)
+
+    # ---- 3.  追加 optimizer 到同一子目录 ----------------------------
+    with zipfile.ZipFile(path, "r") as zf:
+        for name in zf.namelist():
+            if name.endswith("/data.pkl") or name.endswith("/byteorder"):
+                base_dir = name.split("/")[0]
+                break
+        else:
+            base_dir = "checkpoint"
+
+    opt_data = {
+        "optimizer_state_dict": optimizer.state_dict(),
+        "step": step, "epoch": epoch, "loss": loss,
+    }
+    with zipfile.ZipFile(path, "a") as zf:
+        zf.writestr(f"{base_dir}/optimizer.pkl", pickle.dumps(opt_data))
+
+    print(f"[INFO] Saved: {path}")
+
+
+def _inject_metadata(dst_path: str):
+    """从项目静态文件生成完整 metadata 并注入 checkpoint ZIP。
+
+    自给自足——不依赖原始 .ckpt。metadata 写入 torch.save 使用的
+    同一个子目录，保持 PyTorch ZIP 结构兼容。
+    """
+    import json, zipfile
+    import numpy as np
+
+    config_path = str(ROOT / "model" / "aifs_config.json")
+    grid_path = str(ROOT / "model" / "grid-n320.npz")
+
+    with open(config_path) as f:
+        config_data = json.load(f)
+    grid = np.load(grid_path)
+
+    # ---- 1. 探测 torch.save 使用的子目录 -------------------------------
+    with zipfile.ZipFile(dst_path, "r") as zf:
+        for name in zf.namelist():
+            if name.endswith("/data.pkl") or name.endswith("/byteorder"):
+                base_dir = name.split("/")[0]
+                break
+        else:
+            base_dir = "checkpoint"
+        meta_dir = f"{base_dir}/anemoi-metadata"
+
+    # ---- 2. 构建完整 metadata -------------------------------------------
+    metadata = {
+        "config": config_data["model_config"],
+        "data_indices": config_data["data_indices"],
+        "dataset": config_data["dataset"],
+        "provenance_training": config_data.get("provenance_training", {}),
+        "training": config_data.get("training", {}),
+        "run_id": config_data.get("run_id", ""),
+        "seed": config_data.get("seed", 0),
+        "uuid": config_data.get("uuid", ""),
+        "timestamp": config_data.get("timestamp", ""),
+        "version": "1.0",
+        "supporting_arrays_paths": {
+            "latitudes": {
+                "path": f"{meta_dir}/latitudes.numpy",
+                "shape": [542080],
+                "dtype": "float64",
+            },
+            "longitudes": {
+                "path": f"{meta_dir}/longitudes.numpy",
+                "shape": [542080],
+                "dtype": "float64",
+            },
+        },
+    }
+
+    # ---- 3. 写入 checkpoint ZIP（与模型同一子目录）---------------------
+    meta_json = json.dumps(metadata).encode("utf-8")
+    with zipfile.ZipFile(dst_path, "a") as dst:
+        dst.writestr(f"{meta_dir}/ai-models.json", meta_json)
+        dst.writestr(
+            f"{meta_dir}/latitudes.numpy",
+            grid["latitudes"].astype(np.float64).tobytes(),
+        )
+        dst.writestr(
+            f"{meta_dir}/longitudes.numpy",
+            grid["longitudes"].astype(np.float64).tobytes(),
+        )
+
+    print(f"[INFO] Metadata injected into {base_dir}/")
+
+# ============================================================================
+# CLI
+# ============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="AIFS v1.1 Training")
+    parser.add_argument("--config", "-c", type=str,
+                        default=str(ROOT / "conf" / "config.yaml"))
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    hw = cfg["hardware"]
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(hw["device_ids"])
+    if hw["device"] == "dcu": os.environ["HIP_VISIBLE_DEVICES"] = str(hw["device_ids"])
+    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+    seed = cfg["training"]["seed"]
+    np.random.seed(seed); torch.manual_seed(seed)
+    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
+
+    print(f"\n{'='*60}\n  AIFS v1.1 — Pre-training\n{'='*60}")
+    print(f"  Device: {hw['device']}  Dataset: {cfg['data']['data_dir']}")
+    print(f"  LR: {cfg['training']['peak_lr']:.2e}  Steps: {cfg['training']['max_steps']}")
+    print(f"  Output: {cfg['checkpoint']['output_dir']}\n{'='*60}\n")
+
+    try: train(cfg)
+    except KeyboardInterrupt: print("\n[INFO] Interrupted.")
+    except Exception: traceback.print_exc(); sys.exit(1)
+
+if __name__ == "__main__": main()