Upload folder using huggingface_hub

.ms_upload_cache

@@ -0,0 +1 @@
+{"version": 3, "repo_id": "OneScience/WeatherNext2", "files": {"model/__pycache__/fgn.cpython-311.pyc|1787129521.0|26822": {"hash": "194470d61ce734af9a3531410be74b5f6344e55028cf50334b5ffb63e9359484", "size": 26822, "status": "c"}, "conf/config.yaml|1787189086.0|3101": {"hash": "837175aa730e9754739f9eb85b5ccaaf198a83e26bdc1485410e4a5d8a39330d", "size": 3101, "status": "c"}, "README.md|1787118190.0|7664": {"hash": "4058aaeb1f88068ac319b037830b1e6b0667d5a58f9f3cbed9b62ed95cb5face", "size": 7664, "status": "c"}, "scripts/fake_data.py|1787118238.0|3101": {"hash": "478615dcf3e95eaae7a8c24f913c2c938f5aa65b18292ca7e0db3b224b2ecfa9", "size": 3101, "status": "c"}, "scripts/result.py|1787110546.0|8612": {"hash": "7bb894cb829949e18ae2fa3702f26eec5b97cdfcbd97d515965e13ed811a5932", "size": 8612, "status": "c"}, "scripts/inference.py|1787110490.0|3682": {"hash": "158ba56d5e71f27cff76254205895505a39d2ad4a31421b697c1c3df1542f79c", "size": 3682, "status": "c"}, "model/fgn.py|1787110319.0|17001": {"hash": "21c87f83975836c72f29da3c7dea2cbebe777bd2b94cd8962815546ef245d12d", "size": 17001, "status": "c"}, "configuration.json|1787110570.0|38": {"hash": "97d4072fc4a7a3b71e3184609fc491e77a324a982d9df860d9379c009dff84de", "size": 38, "status": "c"}, "scripts/train.py|1787110476.0|9363": {"hash": "c2abc2179712f6366dcb7ef6fa074ec9b62bd7e0a4c466c56983934233936cce", "size": 9363, "status": "c"}, "weight/.gitkeep|1787110570.0|0": {"hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "size": 0, "status": "c"}}}

README.md

@@ -0,0 +1,182 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth science
+- Weather forecasting
+- Probabilistic weather forecasting
+- Ensemble forecasting
+- Graph neural networks
+- Conditional normalization
+- CRPS
+- ERA5
+frameworks: PyTorch
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">WeatherNext2 · FGN</span>
+  </strong>
+</p>
+
+
+# Model Introduction
+
+FGN (**F**unctional **G**enerative **N**etworks) was proposed by Google DeepMind and is the probabilistic global weather forecasting model in the WeatherNext series (WeatherNext2). FGN uses the fair continuous ranked probability score (fair CRPS) of the marginal distribution as its training objective. It models aleatoric uncertainty by injecting a global noise vector into conditional normalization layers and epistemic uncertainty through an ensemble of independently trained models. In this way, it captures the joint spatial structure of ensemble forecasts while optimizing only pointwise marginal objectives, and surpasses GenCast and ECMWF ENS across 15-day medium-range forecasts.
+
+Paper:Skilful joint probabilistic weather forecasting from marginals
+
+https://arxiv.org/abs/2506.14285
+
+# Model Description
+
+FGN uses a GNN encoder-processor-decoder architecture consistent with the GenCast denoiser: a sparse GNN encoder maps latitude-longitude grid inputs to a latent space on a six-times-subdivided icosahedral grid, a graph-transformer processor advances the atmospheric state on that grid, and a GNN decoder maps the latent grid back to the output grid. Each forecast samples a 32-dimensional global noise vector, embeds it with a single matrix multiplication, and injects it into all conditional LayerNorm layers (equivalent to applying a learned functional perturbation to the network parameters) as the source of ensemble spread. Under a second-order Markov assumption (using the two previous frames), the model generates forecasts autoregressively at 6-hour intervals. This repository is a minimal reproduction based on the paper and is integrated with the OneScience data loading and training workflow.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Probabilistic/ensemble medium-range weather forecasting research | Learn pointwise marginal distributions with CRPS and generate joint ensemble forecasts. |
+| Uncertainty modeling research | Reproducible parameter-space noise injection (conditional normalization) and deep ensemble mechanisms. |
+| Graph + Transformer latent-space model research | Encoder-processor-decoder architecture and fair CRPS objective. |
+| Local quick validation | Use synthetic data to check data loading, training, inference, and result scripts. |
+| ModelScope/OneCode execution | Download the model package, install dependencies, and run the scripts directly. |
+| Multi-card training | Launch multi-process training with `torchrun`. |
+
+
+# Usage
+
+## 1. OneCode Usage
+
+Use the OneCode online environment for intelligent one-click AI4S programming:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- GPU or DCU is recommended.
+- CPU can be used for imports and small-configuration connectivity validation, but full training and inference are slower.
+- DCU users must install DTK beforehand. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/WeatherNext2 --local-dir ./WeatherNext2
+cd WeatherNext2
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+```bash
+# 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
+
+The OneScience community provides ERA5 data for training (the current repository contains complete data slices subject to data-file size limits). Download it with the command below and confirm that the data path in `conf/config.yaml` is correct:
+
+```bash
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+For a quick workflow validation, run the synthetic data script first:
+
+```bash
+python scripts/fake_data.py
+```
+
+> Note: `scripts/fake_data.py` generates `[T, C, H, W]` data from second-order Markov inputs, forecast steps, batch size, and `grid_shape`. The current small configuration uses 6 channels and a 32×32 grid.
+
+### Training
+
+Single card：
+
+```bash
+python scripts/train.py
+```
+
+Multiple cards：
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training outputs:
+
+```text
+data/checkpoints/model_bak.pth
+data/checkpoints/trloss.npy
+data/checkpoints/valoss.npy
+```
+
+### Training Weights
+The `weight/` folder is reserved for model weights. Pretrained weights are not provided by default; users may train the model using the paper configuration. The paper model (per-seed latent dimension 768, 24 processor layers, an ensemble of 4 model seeds, and approximately 490 TPU-days of total computation) has no publicly released weights.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth` by default, generates `num_members` ensemble members for each initialization time (independently sampling global noise for each), and uses the member mean as the deterministic forecast output:
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are written frame by frame to:
+
+```text
+result/output/
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Outputs include：
+
+- `result/rmse.npy`
+- `result/acc.npy`
+- `result/loss.png`
+- Forecast comparison plots for the specified date and variables
+
+
+# Official Source and Reproduction Notes
+
+- The paper is a Google DeepMind preprint (© 2025 Google DeepMind. All rights reserved.), and no official implementation or weights are publicly available. This repository's `model/fgn.py` is a pure PyTorch minimal reproduction that preserves the paper's encoder-processor-decoder GNN structure, conditional LayerNorm global-noise injection, and fair CRPS objective (with N=2 ensemble samples during training).
+- Differences from the paper (due to the OneScience gridded data pipeline and connectivity-validation scale): the paper uses a six-times-subdivided icosahedral latent grid (approximately 40k nodes) and a 0.25° (1440×721) output grid, with approximately 180M parameters per seed. This reproduction uses a fixed regular 8-neighbor latent grid (`mesh_shape`), ERA5 gridded h5 channels as input placeholders, and a default latent dimension of 64. The number of latent-grid nodes and edge features are determined by the `mesh_shape` and `channel_weights` configurations.
+- The following details are not disclosed in the paper and are assumptions in this reproduction:the specific projection from noise to each layer's scale/shift in conditional LayerNorm (currently implemented with per-layer linear projections initialized to zero so the initial behavior is standard LayerNorm); per-channel weights for the multi-task loss (currently all 1 by default); and the evaluation convention of using the ensemble mean as the deterministic forecast for probabilistic output.
+- Paper-level reproduction requires the four-stage training procedure in the paper (ERA5 1°12h → 1°6h → 0.25°6h → HRES-fc0 0.25° AR fine-tuning); `conf/config.yaml` uses a small configuration for connectivity validation by default.
+
+# Official OneScience 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 and License
+
+- This repository is an independent FGN (WeatherNext2) reproduction (the model code is an original minimal implementation), with the architecture design based on the paper by Alet et al. (2025).
+- Please cite:Alet, F., Price, I., El-Kadi, A., Masters, D., Markou, S., Andersson, T. R., Stott, J., Lam, R., Willson, M., Sanchez-Gonzalez, A. and Battaglia, P. Skilful joint probabilistic weather forecasting from marginals. arXiv:2506.14285, 2025.

conf/config.yaml

@@ -0,0 +1,72 @@
+# FGN (Functional Generative Networks) 训练配置示例
+# 论文配置：第二阶 Markov（输入两个先验状态帧 x_{t-2},x_{t-1}）预测下一状态 x_t；
+# 网格 0.25°（1440×721），潜空间为 6 次细分球面 icosahedral 网格，潜维度 768、
+# 处理器 24 层、6 attention heads（每模型种子约 180M 参数）；训练 4 个模型种子，
+# 评估时每个种子生成 14 个集合成员（共 56）；目标为公平 CRPS（Eq.4），训练时 N=2。
+# 时间步长 6h，15 天预报共 60 帧。
+# 当前为连通性验证小配置：虚拟数据 32×32 网格、潜网格 8×8、latent_dim=64。
+model:
+  start_epoch: 0
+  max_epoch: 100
+  lr: 1E-3                 # 论文 peak lr=8e-4，cosine 退火（warmup 1000 步）
+  patience: 50
+  checkpoint_dir: "./data/checkpoints"
+
+  # FGN 结构参数（论文值见注释）
+  in_channels: 6           # 通道数（论文 6 个大气变量×13 层 + 6 个表面变量 = 84 通道）
+  out_channels: 6
+  input_steps: 2           # 第二阶 Markov，输入先验两帧（论文 input_steps=2）
+  output_steps: 2          # 自回归预报帧数（论文 6h 步长 15 天 = 60 帧）
+  grid_shape: [32, 32]     # 观测网格尺寸（论文 0.25° 1440×721）
+  mesh_shape: [8, 8]       # 潜网格尺寸（论文 6 次细分 icosahedral 网格约 40k 节点）
+  latent_dim: 64           # 潜空间维度（论文 768）
+  num_encoder_layers: 2    # 编码器 GNN 消息传递层数
+  num_decoder_layers: 2    # 解码器 GNN 消息传递层数
+  num_processor_blocks: 2  # 处理器 graph-transformer 块数（论文 24 层）
+  n_heads: 4               # 处理器 attention 头数（论文 6）
+  hidden_dim: 64           # 处理器 FFN / MLP 隐藏维度
+  noise_dim: 32            # 全局噪声向量维度（论文 32，经条件 LayerNorm 注入）
+  num_ensemble_models: 1   # 深度集成的模型种子数（论文 4）
+  num_members: 2           # 推理时每个模型种子生成的集合成员数（论文每种子 14，共 56）
+  channel_weights: [1, 1, 1, 1, 1, 1]  # CRPS 逐通道权重（论文取 GraphCast/GenCast 权重）
+
+# 整个数据读取流程
+datapipe:
+  name: "ERA5"
+  task: "weather_forecasting"
+
+  dataset:
+    type: "hdf5"
+    data_dir: './data/'
+    train_time: [1951, 1952]
+    val_time: [1953]
+    test_time: [1954]
+    img_size: [32, 32]
+    verbose: true
+    cache: false
+
+    # 气象变量（论文 t/z/q/u/v/w 大气变量 + 2t/10u/10v/msl/sst/tp 表面变量；
+    # 此处使用 6 个表面变量占位）
+    channels: ['2m_temperature', '10m_u_component_of_wind', '10m_v_component_of_wind',
+               'mean_sea_level_pressure', 'sea_surface_temperature', 'total_precipitation']
+
+  # DataLoader 配置
+  dataloader:
+    mask_dtype: "float32"
+    batch_size: 2
+    num_workers: 1
+    pin_memory: true
+    drop_last: true
+    shuffle: false
+    prefetch_factor: 2
+    persistent_workers: true
+
+  # 分布式配置
+  distributed:
+    enabled: true
+    sampler: "DistributedSampler"
+    rank: 0
+    world_size: 2
+    shuffle: true
+    seed: 42
+    drop_last: true

config.json

@@ -0,0 +1,107 @@
+{
+  "model_name": "WeatherNext2",
+  "model_type": "fgn",
+  "architectures": [
+    "FGN"
+  ],
+  "framework": "PyTorch",
+  "domain": "climate-and-atmosphere",
+  "task": "probabilistic-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/fgn.py",
+    "scope": "minimal PyTorch FGN reproduction with regular-grid graph operations, conditional LayerNorm noise injection, autoregressive ensembles, and fair CRPS"
+  },
+  "architecture": {
+    "family": "probabilistic GNN encoder, graph-Transformer processor, and GNN decoder",
+    "attention_mechanism": "multi-head self-attention over a regular latent mesh with fixed wrap-around 8-neighbor graph message passing",
+    "input_format": "BSTCHW",
+    "output_format": "BMTCHW",
+    "encoder": "per-cell MLP, adaptive pooling from the observation grid to the latent mesh, and GNN message passing",
+    "processor": "graph-Transformer blocks conditioned by a sampled global noise vector",
+    "decoder": "latent-mesh GNN, bilinear upsampling, and per-cell output MLP",
+    "activation": "GELU",
+    "normalization": "ConditionalLayerNorm",
+    "repository_default_config": {
+      "purpose": "small connectivity-validation configuration",
+      "in_channels": 6,
+      "out_channels": 6,
+      "input_steps": 2,
+      "output_steps": 2,
+      "grid_shape": [
+        32,
+        32
+      ],
+      "mesh_shape": [
+        8,
+        8
+      ],
+      "latent_dim": 64,
+      "num_encoder_layers": 2,
+      "num_decoder_layers": 2,
+      "num_processor_blocks": 2,
+      "n_heads": 4,
+      "hidden_dim": 64,
+      "noise_dim": 32,
+      "num_ensemble_models": 1,
+      "num_members": 2,
+      "channel_weights": [
+        1,
+        1,
+        1,
+        1,
+        1,
+        1
+      ]
+    },
+    "paper_configuration": {
+      "in_channels": 84,
+      "out_channels": 84,
+      "input_steps": 2,
+      "output_steps": 60,
+      "grid_shape": [
+        721,
+        1440
+      ],
+      "latent_mesh": "six-times-subdivided icosahedral grid with approximately 40,000 nodes",
+      "latent_dim": 768,
+      "num_processor_blocks": 24,
+      "n_heads": 6,
+      "noise_dim": 32,
+      "num_ensemble_models": 4,
+      "num_members_per_model": 14
+    }
+  },
+  "data": {
+    "dataset": "ERA5-format HDF5",
+    "variables": [
+      "2m_temperature",
+      "10m_u_component_of_wind",
+      "10m_v_component_of_wind",
+      "mean_sea_level_pressure",
+      "sea_surface_temperature",
+      "total_precipitation"
+    ],
+    "frame_interval_hours": 6,
+    "input_length": 2,
+    "output_length": 2,
+    "channels": 6,
+    "default_smoke_spatial_size": [
+      32,
+      32
+    ],
+    "paper_spatial_size": [
+      721,
+      1440
+    ],
+    "paper_output_length": 60,
+    "normalization": "per-channel means and standard deviations stored in HDF5"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/fgn.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/fake_data.py",
+    "README.md"
+  ]
+}

configuration.json

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

model/fgn.py

@@ -0,0 +1,397 @@
+# coding=utf-8
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Minimal reproduction of FGN (Functional Generative Networks, Google DeepMind
+# "WeatherNext2", arXiv 2506.14285, 2025) following the paper's architecture:
+#
+#   * Grid encoder: a GNN that maps the two-prior-state gridded input onto a
+#     latent mesh (a coarse regular lat/lon grid).
+#   * Processor: a graph-transformer that operates on the latent mesh nodes
+#     with conditional layer-norm layers.
+#   * Grid decoder: a GNN that maps the latent mesh back onto the target grid.
+#
+# The probabilistic core of FGN is preserved:
+#   * a global noise vector n ~ N(0, I)^32 is sampled per ensemble member and
+#     per autoregressive step, embedded by a single matrix multiplication and
+#     passed into *all* conditional layer-norm layers (learned functional
+#     perturbations). This models aleatoric uncertainty.
+#   * epistemic uncertainty is modelled by an ensemble of independently
+#     trained models (deep ensembles); the mini constant model here uses one
+#     seed by default (see conf/config.yaml).
+#   * training objective is the fair CRPS estimator (Eq. 4) with N=2 samples.
+#
+# Differences from the paper (documented in README.md): the paper uses a
+# spherical 6-times-refined icosahedral mesh and a full 768-latent / 24-layer
+# / 6-head processor (about 180M params per seed); here the latent mesh is a
+# fixed regular grid with a wrap-around neighbor graph, and the hyper
+# parameters are reduced to CPU-friendly sizes for connectivity validation.
+# sampled-per-step noise inside a single model plus the AR rollout are kept.
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _latlon_grid(shape):
+    """Regular lat/lon coordinates for a (H, W) grid, North-to-South rows."""
+    H, W = shape
+    lat = torch.linspace(90.0, -90.0, H)
+    lon = torch.linspace(0.0, 360.0 - 360.0 / W, W)
+    return lat, lon
+
+
+def _haversine(lat1, lon1, lat2, lon2):
+    """Haversine distance in metres given points in degrees."""
+    R = 6371000.0
+    p1 = torch.deg2rad(lat1)
+    p2 = torch.deg2rad(lat2)
+    dp = torch.deg2rad(lat2 - lat1)
+    dl = torch.deg2rad(lon2 - lon1)
+    a = torch.sin(dp / 2) ** 2 + torch.cos(p1) * torch.cos(p2) * torch.sin(dl / 2) ** 2
+    return 2 * R * torch.asin(torch.sqrt(a.clamp(0, 1)))
+
+
+def _bearing(lat1, lon1, lat2, lon2):
+    """Initial forward bearing in radians from point 1 to point 2."""
+    p1 = torch.deg2rad(lat1)
+    p2 = torch.deg2rad(lat2)
+    dl = torch.deg2rad(lon2 - lon1)
+    y = torch.sin(dl) * torch.cos(p2)
+    x = torch.cos(p1) * torch.sin(p2) - torch.sin(p1) * torch.cos(p2) * torch.cos(dl)
+    return torch.atan2(y, x)
+
+
+def build_mesh_graph(mesh_shape):
+    """
+    Build a fixed 8-neighbourhood graph over a regular latent mesh.
+    Longitude wraps around; edge features are (forward bearing [rad],
+    haversine distance [km]).
+    """
+    H, W = mesh_shape
+    lat, lon = _latlon_grid(mesh_shape)
+    lat = lat.view(-1, 1).expand(H, W)
+    lon = lon.view(1, -1).expand(H, W)
+
+    src_list, dst_list, feat_list = [], [], []
+    for i in range(H):
+        for j in range(W):
+            for di, dj in ((-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)):
+                ni, nj = i + di, (j + dj) % W
+                if not (0 <= ni < H):
+                    continue
+                s = i * W + j
+                d = ni * W + nj
+                dist_km = _haversine(lat[i, j], lon[i, j], lat[ni, nj], lon[ni, nj]) / 1000.0
+                bear = _bearing(lat[i, j], lon[i, j], lat[ni, nj], lon[ni, nj])
+                src_list.append(s)
+                dst_list.append(d)
+                feat_list.append(torch.stack([bear / math.pi, dist_km / 1000.0]))
+    edge_index = torch.stack([torch.as_tensor(src_list), torch.as_tensor(dst_list)], dim=0)
+    edge_attr = torch.stack(feat_list)
+    return edge_index, edge_attr
+
+
+def _mlp(in_dim, out_dim, hidden_dim, n_layers=2):
+    dims = [in_dim] + [hidden_dim] * (n_layers - 1) + [out_dim]
+    layers = []
+    for i in range(len(dims) - 1):
+        layers.append(nn.Linear(dims[i], dims[i + 1]))
+        if i < len(dims) - 2:
+            layers.append(nn.GELU())
+    return nn.Sequential(*layers)
+
+
+class ConditionalLayerNorm(nn.Module):
+    """
+    Conditional layer-norm as used by FGN: a global noise vector is embedded
+    (single matrix multiplication) and injected, via learned scale/shift, into
+    every normalised module of the network. Sampling different noise vectors
+    n for each ensemble member / timestep is what generates the variance
+    across the ensemble (learned functional perturbations in weight space).
+    """
+
+    def __init__(self, dim, noise_dim=32):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.gamma = nn.Linear(noise_dim, dim)
+        self.beta = nn.Linear(noise_dim, dim)
+        nn.init.zeros_(self.gamma.weight)
+        nn.init.zeros_(self.beta.weight)
+        nn.init.zeros_(self.gamma.bias)
+        nn.init.zeros_(self.beta.bias)
+
+    def forward(self, x, noise_emb):
+        # x: [B, N, D]; noise_emb: [B, noise_dim]
+        scale = self.gamma(noise_emb).unsqueeze(1)          # [B, 1, D]
+        shift = self.beta(noise_emb).unsqueeze(1)           # [B, 1, D]
+        return self.norm(x) * (1.0 + scale) + shift
+
+
+class GNNLayer(nn.Module):
+    """
+    Message-passing layer with edge features (mean-aggregate, residual).
+    For the grid->mesh encoder the message function knly conditions on the
+    sender node features and the edge features, mirroring FGN's removal of
+    the receiver-mesh-node conditioning in the encoder message function.
+    """
+
+    def __init__(self, dim, edge_dim=2, hidden_dim=64, noise_dim=32, receiver_cond=True):
+        super().__init__()
+        msg_in = 2 * dim + edge_dim if receiver_cond else dim + edge_dim
+        self.edge_mlp = _mlp(msg_in, dim, hidden_dim)
+        self.node_mlp = _mlp(dim, dim, hidden_dim)
+        self.norm = ConditionalLayerNorm(dim, noise_dim)
+        self.receiver_cond = receiver_cond
+
+    def forward(self, x, edge_index, edge_attr, noise_emb):
+        B, N, D = x.shape
+        src, dst = edge_index
+        offsets = torch.arange(B, device=x.device) * N
+        src_b = (src.unsqueeze(0) + offsets.view(B, 1)).reshape(-1)
+        dst_b = (dst.unsqueeze(0) + offsets.view(B, 1)).reshape(-1)
+        edge_attr_b = edge_attr.unsqueeze(0).expand(B, -1, -1).reshape(-1, edge_attr.size(1))
+        xb = x.reshape(B * N, D)
+        if self.receiver_cond:
+            msg = self.edge_mlp(torch.cat([xb[src_b], xb[dst_b], edge_attr_b], dim=1))
+        else:
+            msg = self.edge_mlp(torch.cat([xb[src_b], edge_attr_b], dim=1))
+        agg = torch.zeros_like(xb)
+        agg.index_add_(0, dst_b, msg)
+        cnt = torch.bincount(dst_b, minlength=B * N).clamp(min=1).unsqueeze(1)
+        agg = agg / cnt
+        agg = agg.reshape(B, N, D)
+        return self.norm(x + self.node_mlp(agg), noise_emb)
+
+
+class GridEncoder(nn.Module):
+    """
+    Maps the gridded input states onto the latent mesh with a per-cell input
+    MLP, an adaptive pooling to the mesh resolution, and graph message passing
+    on the mesh graph (encoder message function without receiver conditioning).
+    """
+
+    def __init__(self, in_channels, latent_dim, mesh_shape, num_layers=2, hidden_dim=64,
+                 edge_dim=2, noise_dim=32):
+        super().__init__()
+        self.input_mlp = _mlp(in_channels, latent_dim, hidden_dim)
+        self.gnn = nn.ModuleList([
+            GNNLayer(latent_dim, edge_dim=edge_dim, hidden_dim=hidden_dim,
+                     noise_dim=noise_dim, receiver_cond=False)
+            for _ in range(num_layers)
+        ])
+        self.mesh_shape = mesh_shape
+        self.edge_index, self.edge_attr = build_mesh_graph(mesh_shape)
+
+    def forward(self, x, noise_emb):
+        B, C, H, W = x.shape
+        feat = x.permute(0, 2, 3, 1).reshape(-1, C)
+        feat = self.input_mlp(feat).reshape(B, H, W, -1).permute(0, 3, 1, 2)
+        mesh = F.adaptive_avg_pool2d(feat, self.mesh_shape)     # [B, D, Hm, Wm]
+        mesh = mesh.permute(0, 2, 3, 1).reshape(B, self.mesh_shape[0] * self.mesh_shape[1], -1)
+        edge_index, edge_attr = self.edge_index.to(x.device), self.edge_attr.to(x.device)
+        for layer in self.gnn:
+            mesh = layer(mesh, edge_index, edge_attr, noise_emb)
+        return mesh
+
+
+class GraphTransformerBlock(nn.Module):
+    """
+    One graph-transformer block of the processor: multi-head self-attention
+    over mesh tokens plus edge message passing, each with residual connection
+    and conditioned layer-norm.
+    """
+
+    def __init__(self, latent_dim, n_heads, hidden_dim, edge_dim=2, noise_dim=32):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(
+            latent_dim, n_heads, batch_first=True, dropout=0.0
+        )
+        self.attn_norm = ConditionalLayerNorm(latent_dim, noise_dim)
+        self.gnn = GNNLayer(latent_dim, edge_dim=edge_dim, hidden_dim=hidden_dim,
+                            noise_dim=noise_dim, receiver_cond=True)
+        self.gnn_norm = ConditionalLayerNorm(latent_dim, noise_dim)
+        self.ff = nn.Sequential(
+            nn.Linear(latent_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, latent_dim),
+        )
+        self.ff_norm = ConditionalLayerNorm(latent_dim, noise_dim)
+
+    def forward(self, x, edge_index, edge_attr, noise_emb):
+        # self-attention over mesh tokens
+        h = self.attn_norm(x, noise_emb)
+        h, _ = self.attn(h, h, h)
+        x = x + self.gnn_norm(h, noise_emb)
+        # edge message passing on the mesh graph
+        x = self.gnn(x, edge_index, edge_attr, noise_emb)
+        # feed-forward
+        h = self.ff_norm(x, noise_emb)
+        x = x + self.ff(h)
+        return x
+
+
+class GridDecoder(nn.Module):
+    """
+    Maps the latent mesh back onto the target grid (bilinear upsample) and
+    predicts per-channel fields with an output MLP.
+    """
+
+    def __init__(self, latent_dim, out_channels, grid_shape, mesh_shape, num_layers=2,
+                 hidden_dim=64, edge_dim=2, noise_dim=32):
+        super().__init__()
+        self.gnn = nn.ModuleList([
+            GNNLayer(latent_dim, edge_dim=edge_dim, hidden_dim=hidden_dim,
+                     noise_dim=noise_dim, receiver_cond=True)
+            for _ in range(num_layers)
+        ])
+        self.grid_shape = grid_shape
+        self.mesh_shape = mesh_shape
+        self.edge_index, self.edge_attr = build_mesh_graph(mesh_shape)
+        self.output_mlp = _mlp(latent_dim, out_channels, hidden_dim)
+
+    def forward(self, mesh, noise_emb):
+        B, N, D = mesh.shape
+        edge_index = self.edge_index.to(mesh.device)
+        edge_attr = self.edge_attr.to(mesh.device)
+        for layer in self.gnn:
+            mesh = layer(mesh, edge_index, edge_attr, noise_emb)
+        H, W = self.grid_shape
+        Hm, Wm = self.mesh_shape
+        mesh = mesh.transpose(1, 2).reshape(B, D, Hm, Wm)
+        grid = F.interpolate(mesh, size=self.grid_shape, mode="bilinear", align_corners=False)
+        grid = grid.permute(0, 2, 3, 1).reshape(B, H * W, D)
+        return self.output_mlp(grid).reshape(B, H, W, -1).permute(0, 3, 1, 2)
+
+
+class FGN(nn.Module):
+    """
+    Config-driven FGN (Functional Generative Networks) wrapper.
+
+    Args:
+        in_channels: Number of state channels per frame (concatenated prior
+            states fed to the grid encoder).
+        out_channels: Number of forecast channels per frame.
+        input_steps: Number of input (prior weather state) frames. FGN uses a
+            second-order Markov assumption, input_steps=2.
+        output_steps: Number of autoregressive forecast frames.
+        grid_shape: Spatial shape of the (gridded) input state.
+        mesh_shape: Latent mesh resolution (each dimension).
+        latent_dim: Feature dimension of latent mesh tokens.
+        num_encoder_layers / num_decoder_layers: GNN message-passing layers.
+        num_processor_blocks: Graph-transformer blocks in the processor.
+        n_heads: Attention heads of the processor.
+        hidden_dim: Feed-forward / MLP hidden size.
+        noise_dim: Dimension of the global noise vector injected through the
+            conditional layer-norm layers (paper: 32).
+        channel_weights: Per-channel weights for the fair-CRPS objective
+            (taken from the GenCast/GraphCast loss weighting by default).
+    """
+
+    def __init__(
+        self,
+        in_channels=6,
+        out_channels=6,
+        input_steps=2,
+        output_steps=2,
+        grid_shape=(32, 32),
+        mesh_shape=(8, 8),
+        latent_dim=64,
+        num_encoder_layers=2,
+        num_decoder_layers=2,
+        num_processor_blocks=1,
+        n_heads=4,
+        hidden_dim=64,
+        noise_dim=32,
+        channel_weights=None,
+    ):
+        super().__init__()
+        self.in_channels = int(in_channels)
+        self.out_channels = int(out_channels)
+        self.input_steps = int(input_steps)
+        self.output_steps = int(output_steps)
+        self.grid_shape = (int(grid_shape[0]), int(grid_shape[1]))
+        self.mesh_shape = (int(mesh_shape[0]), int(mesh_shape[1]))
+        self.noise_dim = int(noise_dim)
+
+        self.noise_embed = nn.Linear(self.noise_dim, self.noise_dim)
+
+        self.encoder = GridEncoder(
+            self.in_channels * self.input_steps, int(latent_dim), self.mesh_shape,
+            num_layers=int(num_encoder_layers), hidden_dim=int(hidden_dim),
+            noise_dim=self.noise_dim,
+        )
+        self.processor = nn.ModuleList([
+            GraphTransformerBlock(
+                int(latent_dim), int(n_heads), int(hidden_dim), noise_dim=self.noise_dim
+            )
+            for _ in range(int(num_processor_blocks))
+        ])
+        self.decoder = GridDecoder(
+            int(latent_dim), self.out_channels, self.grid_shape, self.mesh_shape,
+            num_layers=int(num_decoder_layers), hidden_dim=int(hidden_dim),
+            noise_dim=self.noise_dim,
+        )
+
+        if channel_weights is None:
+            channel_weights = torch.ones(self.out_channels)
+        self.register_buffer("channel_weights", torch.as_tensor(channel_weights, dtype=torch.float32))
+
+    def _rollout(self, x, noise):
+        """
+        Autoregressive rollout: at each output step sample the next state
+        conditional on the last `input_steps` prior states x_{t-2}, x_{t-1}
+        (second-order Markov), sampling a fresh global noise vector per step.
+        """
+        B, S, C, H, W = x.shape
+        state = list(torch.unbind(x, dim=1))
+        outs = []
+        for t in range(self.output_steps):
+            embrace = self.noise_embed(noise[t])              # [B, noise_dim]
+            inp = torch.cat(state, dim=1)                     # [B, S*C, H, W]
+            latent = self.encoder(inp, embrace)
+            edge_index, edge_attr = self.encoder.edge_index, self.encoder.edge_attr
+            edge_index = edge_index.to(x.device)
+            edge_attr = edge_attr.to(x.device)
+            for block in self.processor:
+                latent = block(latent, edge_index, edge_attr, embrace)
+            frame = self.decoder(latent, embrace)             # [B, C, H, W]
+            outs.append(frame)
+            state.append(frame)
+            state = state[-self.input_steps:]
+        return torch.stack(outs, dim=1)                       # [B, output_steps, C, H, W]
+
+    def forward(self, x, num_members=1):
+        """
+        Args:
+            x: Input state frames, shape [batch, input_steps, C, H, W].
+            num_members: Number of independent ensemble members to generate
+                (each member samples independent global noise per step).
+        Returns:
+            Forecast frames, shape [batch, num_members, output_steps, C, H, W].
+        """
+        device = x.device
+        members = []
+        for _ in range(int(num_members)):
+            noise = torch.randn(self.output_steps, x.size(0), self.noise_dim, device=device)
+            members.append(self._rollout(x, noise))
+        return torch.stack(members, dim=1)
+
+    def crps_loss(self, pred, target):
+        """
+        Fair CRPS objective (Eq. 4 of the paper) with an N-member ensemble,
+        averaged over all locations, variables, levels and output steps:
+
+            fCRPS(F_1:N, y) = 1/N sum_i |F_i - y|
+                              - 1/(2 N (N-1)) sum_{i != i'} |F_i - F_i'|
+
+        With N=2 this reduces to 0.5(|F1-y|+|F2-y|) - 0.5|F1-F2|. The loss is
+        weighted per channel to match the GenCast/GraphCast loss weighting.
+        """
+        N = pred.size(1)
+        mae = torch.abs(pred - target.unsqueeze(1)).mean(dim=1)      # (1/N) sum_i |F_i - y|
+        per = torch.abs(pred.unsqueeze(2) - pred.unsqueeze(1)).sum(dim=(1, 2)) / (N * (N - 1))
+        crps = mae - 0.5 * per                                        # [B, T, C, H, W]
+        w = self.channel_weights.view(1, 1, self.out_channels, 1, 1)
+        return (crps * w).mean()

scripts/fake_data.py

@@ -0,0 +1,74 @@
+import os
+import h5py
+import numpy as np
+from onescience.utils.YParams import YParams
+
+
+# FGN 使用两个先验状态，自回归预测配置指定的后续状态数。
+def get_dims(cfg_model, cfg_data):
+    H, W = map(int, cfg_model.grid_shape)
+    if tuple(map(int, cfg_data.dataset.img_size)) != (H, W):
+        raise ValueError("model.grid_shape and datapipe.dataset.img_size must match")
+    input_steps = int(cfg_model.input_steps)
+    output_steps = int(cfg_model.output_steps)
+    samples = int(cfg_data.dataloader.batch_size)
+    T = input_steps + output_steps + samples - 1
+    return {
+        "T": T, "H": H, "W": W, "time_step": 6,
+        "input_steps": input_steps, "output_steps": output_steps,
+    }
+
+
+def generate_fake_h5(data_dir, var_names, years, dims):
+    """
+    为每个年份生成一个空 h5 文件。
+    利用 HDF5 chunked 数据集未写入 chunk 即返回 fill_value=0 的特性，
+    文件实际只含元数据，极小，但 shape 与真实数据完全一致。
+    均值/标准差也作为数据集内嵌进每年的 h5，与 era5.py 新版读取方式对应。
+
+    注意：ERA5Datapipe 要求 samples_per_year = T - input_steps - output_steps + 1 >= 1，
+        T 由 input_steps、output_steps 与 batch_size 自动计算。
+    """
+    os.makedirs(os.path.join(data_dir, "data"), exist_ok=True)
+    T, C = dims["T"], len(var_names)
+    H, W = dims["H"], dims["W"]
+
+    means = np.zeros((1, C, 1, 1), dtype=np.float32)
+    stds = np.ones((1, C, 1, 1), dtype=np.float32)
+
+    for year in years:
+        path = os.path.join(data_dir, "data", 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_kb = os.path.getsize(path) / 1024
+        print(f"  {year}.h5  shape=({T},{C},{H},{W})  "
+              f"logical={T*C*H*W*4/1024**3:.1f}GB  actual={size_kb:.1f}KB")
+
+
+if __name__ == "__main__":
+    cfg_model = YParams("conf/config.yaml", "model")
+    cfg_datapipe = YParams("conf/config.yaml", "datapipe")
+
+    if cfg_datapipe.dataset.data_dir.startswith("/public/") or cfg_datapipe.dataset.data_dir.startswith("/work2/"):
+        print("请检查 config，确保各 *_dir 指向本地测试路径而非生产路径。")
+        exit()
+
+    years = cfg_datapipe.dataset.train_time + cfg_datapipe.dataset.val_time + cfg_datapipe.dataset.test_time
+    atm_vars = cfg_datapipe.dataset.channels
+    if len(atm_vars) != int(cfg_model.in_channels) or len(atm_vars) != int(cfg_model.out_channels):
+        raise ValueError("channel count must match model input/output channels")
+
+    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, get_dims(cfg_model, cfg_datapipe))
+
+    print("\n✅ Fake datasets generated.")

scripts/inference.py

@@ -0,0 +1,93 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（inference.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import torch
+import os
+import glob
+import numpy as np
+import h5py
+from tqdm import tqdm
+from model.fgn import FGN
+from onescience.utils.YParams import YParams
+from onescience.datapipes.climate import ERA5Datapipe
+
+
+def get_stats(data_dir, channels):
+    """从新版 h5 中读取变量列表与归一化参数（均值/标准差）"""
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        ds = f["fields"]
+        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+        mu = f["global_means"][:]   # [1, C, 1, 1]
+        std = f["global_stds"][:]
+
+    channel_indices = [all_variables.index(v) for v in channels]
+    means = mu[:, channel_indices, :, :]
+    stds = std[:, channel_indices, :, :]
+    return means, stds
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+
+    ## Model config init
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    means, stds = get_stats(cfg_data.dataset.data_dir, cfg_data.dataset.channels)
+
+    cfg['N_in_channels'] = len(cfg_data.dataset.channels)
+    cfg['N_out_channels'] = len(cfg_data.dataset.channels)
+
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.test_time,
+        distributed=False,
+        input_steps=cfg.input_steps,
+        output_steps=cfg.output_steps,
+        batch_size=1,
+        num_workers=4,
+    )
+    test_dataloader, _ = datapipe.get_dataloader("test")
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=device, weights_only=False)
+    model = FGN(
+        in_channels=cfg['N_in_channels'],
+        out_channels=cfg['N_out_channels'],
+        input_steps=cfg.input_steps,
+        output_steps=cfg.output_steps,
+        grid_shape=cfg.grid_shape,
+        mesh_shape=cfg.mesh_shape,
+        latent_dim=cfg.latent_dim,
+        num_encoder_layers=cfg.num_encoder_layers,
+        num_decoder_layers=cfg.num_decoder_layers,
+        num_processor_blocks=cfg.num_processor_blocks,
+        n_heads=cfg.n_heads,
+        hidden_dim=cfg.hidden_dim,
+        noise_dim=cfg.noise_dim,
+        channel_weights=cfg.channel_weights,
+    ).to(device)
+    model.load_state_dict(ckpt["model_state_dict"])
+
+    model.eval()
+    os.makedirs('result/output/', exist_ok=True)
+    print(f"📂 infer results will be generated to './result/output/'")
+    print(f"📂 generating {cfg.num_members} ensemble members per init, saving their mean per frame")
+    with torch.no_grad():
+        for data in tqdm(test_dataloader, desc="Inferring testset", unit="batch"):
+            invar = data[0].to(device, dtype=torch.float32)      # [1, input_steps, C, H, W]
+            members = model(invar, num_members=cfg.num_members)  # [1, M, output_steps, C, H, W]
+            pred = members[0].mean(dim=0).cpu().numpy()          # 集合成员均值 [T, C, H, W]
+            for t in range(pred.shape[0]):
+                fname = data[4][cfg.input_steps + t][0]          # 该预测帧对应的时刻
+                pred_var = pred[t]                               # [C, H, W]
+                pred_var = pred_var * stds + means
+                np.save(f"result/output/{fname}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,203 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+import sys
+import glob
+import h5py
+from datetime import datetime
+from tqdm import tqdm
+from onescience.utils.fcn.YParams import YParams
+from matplotlib import rcParams
+
+# rcParams['font.family'] = 'serif'
+# rcParams['font.serif'] = ['DejaVu Serif']
+rcParams['mathtext.fontset'] = 'stix'
+rcParams['axes.linewidth'] = 0.9
+rcParams['xtick.major.width'] = 0.9
+rcParams['ytick.major.width'] = 0.9
+
+
+def get_metadata(data_dir, channels):
+    """从新版 h5 attrs 中读取变量列表和 time_step"""
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        ds = f["fields"]
+        all_variables = [v.decode() if isinstance(v, bytes) else v for v in ds.attrs["variables"]]
+        time_step = int(ds.attrs["time_step"])
+
+    channel_indices = [all_variables.index(v) for v in channels]
+
+    total_files = [f for f in os.listdir('./result/output/') if f.endswith('.npy')]
+    total_files.sort()
+    return total_files, channel_indices, time_step
+
+
+def filename_to_index(filename, time_step):
+    """将 YYYYMMDDHH 格式的文件名转换为年度 h5 文件中的时间步索引"""
+    dt = datetime.strptime(filename, "%Y%m%d%H")
+    year_start = datetime(dt.year, 1, 1)
+    hours = (dt - year_start).total_seconds() / 3600
+    return int(hours / time_step)
+
+
+def get_result(total_files, channel_indices, time_step, data_dir, clim_mean):
+    channel_rmse = np.zeros(len(channel_indices))
+    channel_acc = np.zeros(len(channel_indices))
+    clim_mean = clim_mean[0, :, :, :]
+    if not os.path.exists('./result/rmse.npy') or not os.path.exists('result/acc.npy'):
+        numerator = np.zeros(len(channel_indices))
+        pred_sq_sum = np.zeros(len(channel_indices))
+        label_sq_sum = np.zeros(len(channel_indices))
+        for file in tqdm(total_files, unit="files"):
+            fname = file[:-4]  # 去掉 .npy
+            year = fname[:4]
+            t_idx = filename_to_index(fname, time_step)
+            with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+                label = f["fields"][t_idx]  # [C, H, W]
+                label = label[channel_indices]
+            pred = np.load(f'result/output/{file}').squeeze()
+            if pred.ndim == 2:
+                pred = pred[np.newaxis]  # 单通道时 squeeze 会压缩掉通道维，恢复为 [C, H, W]
+
+            label_anom = label - clim_mean
+            pred_anom = pred - clim_mean
+            # 累加
+            numerator += np.sum(pred_anom * label_anom, axis=(1, 2))
+            pred_sq_sum += np.sum(pred_anom ** 2, axis=(1, 2))
+            label_sq_sum += np.sum(label_anom ** 2, axis=(1, 2))
+
+            channel_rmse += np.sqrt(np.mean((label - pred) ** 2, axis=(1, 2)))
+        channel_rmse /= len(total_files)
+        channel_acc = numerator / (np.sqrt(pred_sq_sum * label_sq_sum) + 1e-8)
+        np.save('./result/acc.npy', channel_acc)
+        np.save('./result/rmse.npy', channel_rmse)
+
+
+def show_result():
+    channel_rmse = np.load('./result/rmse.npy')
+    channel_acc = np.load('./result/acc.npy')
+
+    channels = [cfg_data.dataset.channels[i] for i in range(len(channel_indices))]
+    w = 24  # 最长 channel 名宽度
+
+    # 表头
+    print(f"┌{'─' * (w + 2)}┬{'─' * 14}┬{'─' * 14}┐")
+    print(f"│ {'Channel':<{w}} │ {'RMSE':>12} │ {'ACC':>12} │")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    # 数据行
+    for i, ch in enumerate(channels):
+        print(f"│ {ch:<{w}} │ {channel_rmse[i]:>12.4f} | {channel_acc[i]:>12.4f} |")
+    print(f"├{'─' * (w + 2)}┼{'─' * 14}┼{'─' * 14}┤")
+    print(f"│ {'Average':<{w}} │ {np.mean(channel_rmse):>12.4f} │ {np.mean(channel_acc):>12.4f} │")
+    print(f"└{'─' * (w + 2)}┴{'─' * 14}┴{'─' * 14}┘")
+
+
+def plot(label, pred, var, filename):
+    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+
+    xtick_labels = ['180°W', '90°W', '0°', '90°E', '180°E']
+    ytick_labels = ['90°S', '45°S', '0°', '45°N', '90°N']
+    xticks = np.linspace(0, label.shape[-1] - 1, 5)
+    yticks = np.linspace(0, label.shape[-2] - 1, 5)
+
+    vmin = min(label.min(), pred.min())
+    vmax = max(label.max(), pred.max())
+
+    diff = label - pred
+    rmse = np.sqrt(np.mean(diff ** 2))
+    diff_abs_max = np.abs(diff).max()
+
+    plot_configs = [
+        {'data': label, 'title': 'Truth', 'cmap': 'viridis', 'vmin': vmin, 'vmax': vmax},
+        {'data': pred,  'title': 'Prediction', 'cmap': 'viridis', 'vmin': vmin, 'vmax': vmax},
+        {'data': diff,  'title': f'Difference (RMSE={rmse:.2f})', 'cmap': 'RdBu_r', 'vmin': -diff_abs_max, 'vmax': diff_abs_max},
+    ]
+
+    for ax, cfg2 in zip(axes, plot_configs):
+        im = ax.imshow(cfg2['data'], cmap=cfg2['cmap'], vmin=cfg2['vmin'], vmax=cfg2['vmax'])
+        ax.set_title(cfg2['title'], fontsize=12, pad=4)
+        ax.set_xlabel('Longitude')
+        ax.set_ylabel('Latitude')
+        ax.set_xticks(xticks)
+        ax.set_xticklabels(xtick_labels)
+        ax.set_yticks(yticks)
+        ax.set_yticklabels(ytick_labels)
+        plt.colorbar(im, ax=ax, orientation='horizontal')
+
+    fig.suptitle(var, fontsize=14, fontweight='bold', y=0.98)
+    plt.savefig(filename, dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+def plot_loss(train_loss, valid_loss):
+    mask = ~(np.isnan(train_loss) | np.isnan(valid_loss))
+    train_loss = train_loss[mask]
+    valid_loss = valid_loss[mask]
+
+    fig, ax = plt.subplots(figsize=(5, 3.5))
+    colors = {'train': '#2563EB', 'valid': '#EA580C'}
+    epochs = np.arange(1, len(train_loss) + 1)
+
+    ax.plot(epochs, train_loss, color=colors['train'], linewidth=1.5, label='Train')
+    ax.plot(epochs, valid_loss, color=colors['valid'], linewidth=1.5, label='Valid', linestyle='--')
+    min_idx = np.argmin(valid_loss)
+    ax.scatter(epochs[min_idx], valid_loss[min_idx],
+               color=colors['valid'], s=40, zorder=5, edgecolors='white')
+    ax.annotate(f'Best: {valid_loss[min_idx]:.3f}',
+                xy=(epochs[min_idx], valid_loss[min_idx]),
+                xytext=(10, 10), textcoords='offset points', fontsize=8, color=colors['valid'],
+                arrowprops=dict(arrowstyle='-', color=colors['valid'], lw=0.5))
+
+    ax.set(xlabel='Epoch', ylabel='Loss', xlim=(0, len(train_loss) + 1))
+    ax.legend(frameon=False, loc='upper right')
+    ax.grid(True, linestyle='--', alpha=0.3)
+    ax.spines[['top', 'right']].set_visible(False)
+
+    plt.tight_layout()
+    plt.savefig('./result/loss.png', dpi=300, bbox_inches='tight')
+    plt.close()
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    config_file_path = os.path.join(current_path, 'conf/config.yaml')
+    cfg = YParams(config_file_path, 'model')
+    cfg_data = YParams(config_file_path, "datapipe")
+
+    train_loss = np.load('./data/checkpoints/trloss.npy')
+    valid_loss = np.load('./data/checkpoints/valoss.npy')
+    plot_loss(train_loss, valid_loss)
+
+    data_dir = cfg_data.dataset.data_dir
+    total_files, channel_indices, time_step = get_metadata(data_dir, cfg_data.dataset.channels)
+
+    # Load data & Compute RMSE/ACC per channel
+    h5_files = sorted(glob.glob(os.path.join(data_dir, "data", "*.h5")))
+    with h5py.File(h5_files[0], "r") as f:
+        mu = f["global_means"][:]
+    clim_mean = mu[:, channel_indices, :, :]
+    get_result(total_files, channel_indices, time_step, data_dir, clim_mean)
+    show_result()
+
+    ##### 默认绘制第一个预测输出的时刻与全部通道，用户可自行指定 #####
+    eg_files = [total_files[0][:-4]]
+    # 最多绘制 3 个通道的对比图
+    channel_index = list(range(min(3, len(cfg_data.dataset.channels))))
+
+    selected_var = [cfg_data.dataset.channels[int(i)] for i in channel_index]
+    print(f"seleted date: {eg_files}")
+    print(f"selected channels: {selected_var}")
+    for file in eg_files:
+        year = file[:4]
+        t_idx = filename_to_index(file, time_step)
+        with h5py.File(os.path.join(data_dir, 'data', f'{year}.h5'), "r") as f:
+            label = f["fields"][t_idx]  # [C, H, W]
+            label = label[channel_indices]
+        pred = np.load(f'result/output/{file}.npy').squeeze()
+        if pred.ndim == 2:
+            pred = pred[np.newaxis]  # 单通道时 squeeze 会压缩掉通道维，恢复为 [C, H, W]
+        for i in range(len(selected_var)):
+            filename = f'./result/{file}_{selected_var[i]}.png'
+            plot(label[channel_index[i]], pred[channel_index[i]], selected_var[i], filename)
+            print(f'✅plot {filename}')

scripts/train.py

@@ -0,0 +1,223 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+import torch
+import os
+import numpy as np
+import torch.distributed as dist
+import logging
+import time
+
+from model.fgn import FGN
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+
+try:
+    from apex import optimizers
+    _FUSED_ADAM = True
+except Exception:
+    _FUSED_ADAM = False
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    logger = logging.getLogger()
+
+    ## Model config init
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    ## Distributed config init
+    cfg.world_size = 1
+    if "WORLD_SIZE" in os.environ:
+        cfg.world_size = int(os.environ["WORLD_SIZE"])
+    world_rank = 0
+    local_rank = 0
+    if cfg.world_size > 1 and torch.cuda.is_available():
+        dist.init_process_group(backend="nccl", init_method="env://")
+        local_rank = int(os.environ["LOCAL_RANK"])
+        world_rank = dist.get_rank()
+    device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
+
+    ## DataLoader init
+    cfg_data = YParams(config_file_path, "datapipe")
+    cfg['N_in_channels'] = len(cfg_data.dataset.channels)
+    cfg['N_out_channels'] = len(cfg_data.dataset.channels)
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.train_time,
+        distributed=dist.is_initialized(),
+        input_steps=cfg.input_steps,
+        output_steps=cfg.output_steps,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers,
+    )
+    train_dataloader, train_sampler = datapipe.get_dataloader("train")
+    datapipe = ERA5Datapipe(
+        dataset_dir=cfg_data.dataset.data_dir,
+        used_variables=cfg_data.dataset.channels,
+        used_years=cfg_data.dataset.val_time,
+        distributed=dist.is_initialized(),
+        input_steps=cfg.input_steps,
+        output_steps=cfg.output_steps,
+        batch_size=cfg_data.dataloader.batch_size,
+        num_workers=cfg_data.dataloader.num_workers,
+    )
+    val_dataloader, val_sampler = datapipe.get_dataloader("valid")
+
+    # Model init
+    model = FGN(
+        in_channels=cfg['N_in_channels'],
+        out_channels=cfg['N_out_channels'],
+        input_steps=cfg.input_steps,
+        output_steps=cfg.output_steps,
+        grid_shape=cfg.grid_shape,
+        mesh_shape=cfg.mesh_shape,
+        latent_dim=cfg.latent_dim,
+        num_encoder_layers=cfg.num_encoder_layers,
+        num_decoder_layers=cfg.num_decoder_layers,
+        num_processor_blocks=cfg.num_processor_blocks,
+        n_heads=cfg.n_heads,
+        hidden_dim=cfg.hidden_dim,
+        noise_dim=cfg.noise_dim,
+        channel_weights=cfg.channel_weights,
+    ).to(device)
+
+    if _FUSED_ADAM:
+        optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.lr)
+    else:
+        optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=5, mode='min')
+
+    ## Train process init
+    os.makedirs(cfg.checkpoint_dir, exist_ok=True)
+    train_loss_file = f"{cfg.checkpoint_dir}/trloss.npy"
+    valid_loss_file = f"{cfg.checkpoint_dir}/valoss.npy"
+    best_valid_loss = float("inf")
+    best_loss_epoch = 0
+    train_losses = np.empty((0,), dtype=np.float32)
+    valid_losses = np.empty((0,), dtype=np.float32)
+
+    ## Get model params count
+    if cfg.world_size == 1:
+        total_params = sum(p.numel() for p in model.parameters())
+        print("\n\n")
+        print("-" * 50)
+        print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B")
+        print("-" * 50, "\n")
+
+    ## Load model weight if there exist well-trained model
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"):
+        if world_rank == 0:
+            print("\n\n")
+            print("-" * 50)
+            print(f"✅ There has a model weight, load and continue training...")
+            print(f'If you want to train a new model, ensure there is no *.pth file in {cfg.checkpoint_dir}')
+            print("-" * 50, "\n")
+        ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location=device, weights_only=False)
+        model.load_state_dict(ckpt["model_state_dict"])
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
+        best_valid_loss = ckpt["best_valid_loss"]
+        best_loss_epoch = ckpt["best_loss_epoch"]
+        train_losses = np.load(train_loss_file)
+        valid_losses = np.load(valid_loss_file)
+
+    ## Distributed model
+    if dist.is_initialized():
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+    world_rank == 0 and logger.info(f"start training ...")
+
+    for epoch in range(cfg.max_epoch):
+        if dist.is_initialized():
+            train_sampler.set_epoch(epoch)
+            val_sampler.set_epoch(epoch)
+        model.train()
+        train_loss = 0
+        start_time = time.time()
+        for j, data in enumerate(train_dataloader):
+            invar = data[0].to(device, dtype=torch.float32)      # [B, input_steps, C, H, W]
+            outvar = data[1].to(device, dtype=torch.float32)     # [B, output_steps, C, H, W]
+            outvar_pred = model(invar, num_members=cfg.num_members)   # [B, M, output_steps, C, H, W]
+            loss = model.crps_loss(outvar_pred, outvar)          # 论文式(4) 公平 CRPS
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item()
+            if world_rank == 0:
+                logger.info(f'Train: Epoch {epoch}-{j+1}/{len(train_dataloader)} '
+                            f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] '
+                            f'[{(time.time()-start_time)/(j+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                            f'loss:{train_loss / (j+1): .04f}')
+
+        train_loss /= len(train_dataloader)
+
+        model.eval()
+        valid_loss = 0
+        with torch.no_grad():
+            start_time = time.time()
+            for j, data in enumerate(val_dataloader):
+                invar = data[0].to(device, dtype=torch.float32)
+                outvar = data[1].to(device, dtype=torch.float32)
+                outvar_pred = model(invar, num_members=cfg.num_members)
+                loss = model.crps_loss(outvar_pred, outvar)
+
+                if dist.is_initialized():
+                    loss_tensor = loss.detach().to(device)
+                    dist.all_reduce(loss_tensor)
+                    loss = loss_tensor.item() / cfg.world_size
+                    valid_loss += loss
+                else:
+                    valid_loss += loss.item()
+                if world_rank == 0:
+                    logger.info(f'Valid: Epoch {epoch}-{j+1}/{len(val_dataloader)} '
+                            f'[cost {int((time.time()-start_time) // 60):02}:{int((time.time()-start_time) % 60):02}] '
+                            f'loss:{valid_loss / (j+1): .04f}')
+
+        valid_loss /= len(val_dataloader)
+        is_save_ckp = False
+        if valid_loss < best_valid_loss:
+            best_valid_loss = valid_loss
+            best_loss_epoch = epoch
+            world_rank == 0 and save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, cfg.checkpoint_dir)
+            is_save_ckp = True
+        scheduler.step(valid_loss)
+
+        if world_rank == 0:
+            logger.info(f"Epoch [{epoch + 1}/{cfg.max_epoch}], "
+                        f"Train Loss: {train_loss:.4f}, "
+                        f"Valid Loss: {valid_loss:.4f}, "
+                        f"Best loss at Epoch: {best_loss_epoch + 1}"
+                        + (", saving checkpoint" if is_save_ckp else "")
+                        )
+            train_losses = np.append(train_losses, train_loss)
+            valid_losses = np.append(valid_losses, valid_loss)
+            np.save(train_loss_file, train_losses)
+            np.save(valid_loss_file, valid_losses)
+
+        if epoch - best_loss_epoch > cfg.patience:
+            print(f"Loss has not decrease in {cfg.patience} epochs, stopping training...")
+            exit()
+
+
+def save_checkpoint(model, optimizer, scheduler, best_valid_loss, best_loss_epoch, model_path):
+    model_to_save = model.module if hasattr(model, "module") else model
+    state = {"model_state_dict": model_to_save.state_dict(),
+             "optimizer_state_dict": optimizer.state_dict(),
+             "scheduler_state_dict": scheduler.state_dict(),
+             "best_valid_loss": best_valid_loss,
+             "best_loss_epoch": best_loss_epoch,
+            }
+    torch.save(state, f"{model_path}/model.pth")
+    ### the weight file saving may interrupted due to DCU queue limit, get a backup to ensure there at least has one model
+    os.system(f"mv {model_path}/model.pth {model_path}/model_bak.pth")
+
+
+if __name__ == "__main__":
+    current_path = os.getcwd()
+    sys.path.append(current_path)
+    main()