Upload folder using huggingface_hub

.ms_upload_cache

@@ -0,0 +1 @@
+{"version": 3, "repo_id": "OneScience/GraphDOP", "files": {"README.md|1787118190.0|6673": {"hash": "014371f316b8c1a1c437e6c5fedd7a93415921b4871ad37485cf1e682df834e8", "size": 6673, "status": "c"}, "conf/config.yaml|1787189077.0|2517": {"hash": "04ee6d048c400bf199bea31516dfd7c8db381cdbfd34cd9ec37b7bc1e2b1e922", "size": 2517, "status": "c"}, "model/__pycache__/graphdop.cpython-311.pyc|1787129468.0|20086": {"hash": "2ca2bd82b23d7b8a3e58209ebd5ddf898689f9929949fde90ffc7474f13fa439", "size": 20086, "status": "c"}, "scripts/fake_data.py|1787118238.0|3109": {"hash": "9ecb538acc95858926ef9f0a1f3bcbe0a228cc6a4122c4872ec946179ac885ba", "size": 3109, "status": "c"}, "scripts/inference.py|1787060565.0|3460": {"hash": "c3623c7078ace0c3c91407a7c835131adb287015f7dd84a938d473058be3bef4", "size": 3460, "status": "c"}, "scripts/result.py|1787060586.0|8607": {"hash": "2b2b68ec178f50b303f66c4a3e89d40e133a3a1559da3aaf77ad9fbcffff74e2", "size": 8607, "status": "c"}, "configuration.json|1787060533.0|39": {"hash": "6b6f30f1e82ad5243c2ff583dcd33b77949442fbf942d80d9f236ffd7a9fe796", "size": 39, "status": "c"}, "model/graphdop.py|1787060478.0|11202": {"hash": "3ed51f897038bbad17a8fcd215dbee3e37a50f6d2e66353222175132b1fe608d", "size": 11202, "status": "c"}, "scripts/train.py|1787060554.0|9302": {"hash": "b17e5ae03dedb73ae34643054dc1b1e3c89f1468dac6a8804b4966e6e90c8b33", "size": 9302, "status": "c"}, "weight/.gitkeep|1787060604.0|1": {"hash": "01ba4719c80b6fe911b091a7c05124b64eeece964e09c058ef8f9805daca546b", "size": 1, "status": "c"}}}

README.md

@@ -0,0 +1,179 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth science
+- Weather forecasting
+- Medium-range weather forecasting
+- Observation-driven
+- Graph neural networks
+- ERA5
+frameworks: PyTorch
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">GraphDOP</span>
+  </strong>
+</p>
+
+
+# Model Introduction
+
+GraphDOP (Graph-based Direct Observation Prediction) was proposed by ECMWF. It is an end-to-end observation-driven (AI-DOP) weather forecasting model based on graph neural networks (GNNs). The model uses only Earth-system observations (polar-orbiting/geostationary satellite brightness temperatures, occultation bending angles, scatterometer backscatter, radar altimetry, radiosonde and conventional surface observations, and so on) as inputs and training targets. It does not use any physics-based reanalysis field and can produce medium-range forecasts beyond 5 days.
+
+Paper:GraphDOP: Towards skilful data-driven medium-range weather forecasts learnt and initialised directly from observations
+
+https://arxiv.org/abs/2412.15687
+
+# Model Description
+
+GraphDOP uses an encoder-processor-decoder architecture: the GNN encoder maps observations in the input window to an O96 (approximately 1°) latent grid according to spatial proximity, the Transformer processor advances the atmospheric state in latent space, and the GNN decoder maps the latent grid back to target observation locations and outputs forecasts channel by channel. The training objective is channel-weighted mean squared error (WMSE). 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 |
+| :---: | :--- |
+| Observation-driven medium-range weather forecasting research | Learn atmospheric-state representations directly from observations and forecast future windows. |
+| Graph + Transformer latent-space model research | Reproducible encoder-processor-decoder architecture and WMSE 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/GraphDOP --local-dir ./GraphDOP
+cd GraphDOP
+```
+
+### 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 the model input/output windows and `grid_shape`. Because ERA5Datapipe supports only regular grids, this project approximates the irregular Level-1 observations in the paper with six observation-grid channels.
+
+### 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 (1024 latent channels, an O96 latent grid, and 70k steps on 64×H100 GPUs) has no publicly released weights.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth`：
+
+```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 an ECMWF AI-DOP preprint, and no official implementation is publicly available (the source is based on PyTorch Geometric and builds dynamic observation graphs per batch). This repository's `model/graphdop.py` is a pure PyTorch minimal reproduction that preserves the paper's encoder-processor-decoder GNN structure and WMSE objective.
+- Differences from the paper (due to limitations of the OneScience gridded data pipeline): the paper consumes irregular raw Level-1 observations, with each observation connected to its nearest latent-grid neighbors using a dynamically constructed graph. This reproduction uses ERA5 gridded h5 channels as a placeholder observation field and fixes the latent grid to a regular 8-neighbor graph. The edge features (azimuth and Haversine distance) match the paper. Training augmentations such as random observation dropout (25% for satellites/50% for conventional observations) are not implemented.
+- `conf/config.yaml` uses a small configuration (32×32 grid, 8×8 latent grid, latent_dim=64) for connectivity validation by default; paper-level reproduction requires an O96 latent grid, latent_dim=1024, and a larger data scale.
+- The following details are not disclosed in the paper and are assumptions in this reproduction:per-channel weights `w_{c,i}` (currently all 1 by default), graph-construction details, and the specific implementation of windowed attention in the processor.
+
+# 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 GraphDOP reproduction (the model code is an original minimal implementation), with the architecture design based on the paper by Alexe et al. (2024).
+- Please cite:Alexe, M., E. Boucher, P. Lean, E. Pinnington, P. Laloyaux, A. McNally et al. GraphDOP: Towards skilful data-driven medium-range weather forecasts learnt and initialised directly from observations. arXiv:2412.15687, 2024.

conf/config.yaml

@@ -0,0 +1,66 @@
+# GraphDOP 训练配置示例
+# 论文配置：输入一个 12 小时观测窗口（O96 reduced Gaussian 网格约 1°、40320 潜节点），
+# 经 GNN 编码器映射到潜网格、Transformer 处理器推进时间、GNN 解码器预测下一窗口观测；
+# 潜空间通道 1024，WMSE 目标，18 年数据（2004-2021）训练，64×H100 70k 步。
+# 当前为连通性验证小配置：虚拟数据 32×32 网格、潜网格 8×8、latent_dim=64。
+model:
+  start_epoch: 0
+  max_epoch: 100
+  lr: 1E-3                 # 论文起始 lr=1e-3，cosine 退火到 3e-7（warmup 1000 步）
+  patience: 50
+  checkpoint_dir: "./data/checkpoints"
+
+  # GraphDOP 结构参数（论文值见注释）
+  in_channels: 6           # 观测通道数（论文为多仪器通道，如 ATMS/AMSU-A/IASI 等）
+  out_channels: 6          # 预报通道数
+  input_steps: 2           # 输入窗口帧数（论文为单个 12h 窗口；time_step=6h 故取 2 帧）
+  output_steps: 2          # 输出窗口帧数（论文为下一个 12h 窗口）
+  grid_shape: [32, 32]     # 观测网格尺寸（论文 O96 约 1°）
+  mesh_shape: [8, 8]       # 潜网格尺寸（论文 O96 reduced Gaussian 40320 节点）
+  latent_dim: 64           # 潜空间通道数（论文 1024）
+  num_encoder_layers: 2    # 编码器 GNN 消息传递层数
+  num_decoder_layers: 2    # 解码器 GNN 消息传递层数
+  num_processor_blocks: 1  # 处理器 Transformer 块数
+  n_heads: 4
+  hidden_dim: 64
+  channel_weights: [1, 1, 1, 1, 1, 1]  # WMSE 逐通道权重（论文 w_{c,i} 经验值）
+
+# 整个数据读取流程
+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
+
+    # 气象变量（论文观测类型占位：ATMS 亮温/掩星弯角/散射计后向散射/雷达高度计有效波高/常规观测）
+    channels: ['atms_brightness_temperature', 'gpsro_bending_angle', 'ascat_sigma0',
+               'significant_wave_height', '2m_temperature', '10m_wind_speed']
+
+  # DataLoader 配置
+  dataloader:
+    mask_dtype: "float32"
+    batch_size: 4
+    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,104 @@
+{
+  "model_name": "GraphDOP",
+  "model_type": "graphdop",
+  "architectures": [
+    "GraphDOP"
+  ],
+  "framework": "PyTorch",
+  "domain": "climate-and-atmosphere",
+  "task": "observation-driven-medium-range-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/graphdop.py",
+    "scope": "pure-PyTorch minimal reproduction using gridded ERA5 placeholders and fixed regular-mesh graphs instead of the paper's irregular Level-1 observations and dynamic graphs"
+  },
+  "architecture": {
+    "family": "GNN encoder-Transformer processor-GNN decoder",
+    "input_format": "B T C H W",
+    "encoder": "per-grid-cell MLP, adaptive pooling to the latent mesh, then residual mean-aggregation GNN layers",
+    "processor": "pre-normalized Transformer encoder over latent-mesh tokens with learned positional embeddings",
+    "decoder": "latent-mesh GNN, bilinear upsampling, and a per-grid-cell output MLP",
+    "edge_features": [
+      "forward bearing",
+      "Haversine distance"
+    ],
+    "activation": "GELU",
+    "normalization": "LayerNorm",
+    "loss": "channel-weighted mean squared error",
+    "repository_default_config": {
+      "purpose": "connectivity validation with synthetic gridded data",
+      "grid_shape": [
+        32,
+        32
+      ],
+      "mesh_shape": [
+        8,
+        8
+      ],
+      "in_channels": 6,
+      "out_channels": 6,
+      "input_steps": 2,
+      "output_steps": 2,
+      "latent_dim": 64,
+      "num_encoder_layers": 2,
+      "num_decoder_layers": 2,
+      "num_processor_blocks": 1,
+      "attention_heads": 4,
+      "hidden_dim": 64,
+      "channel_weights": [
+        1,
+        1,
+        1,
+        1,
+        1,
+        1
+      ]
+    },
+    "paper_reference_config": {
+      "latent_grid": "O96 reduced Gaussian grid with 40320 nodes",
+      "latent_dim": 1024,
+      "observation_graph": "dynamic graph over irregular Level-1 observations",
+      "training_steps": 70000,
+      "training_hardware": "64 H100 GPUs"
+    }
+  },
+  "data": {
+    "dataset": "ERA5",
+    "role": "regular-grid placeholder for the paper's multi-instrument observations",
+    "variables": [
+      "atms_brightness_temperature",
+      "gpsro_bending_angle",
+      "ascat_sigma0",
+      "significant_wave_height",
+      "2m_temperature",
+      "10m_wind_speed"
+    ],
+    "time_step_hours": 6,
+    "input_length": 2,
+    "output_length": 2,
+    "channels": 6,
+    "spatial_size": [
+      32,
+      32
+    ],
+    "storage_format": "HDF5 fields with T C H W layout",
+    "train_years": [
+      1951,
+      1952
+    ],
+    "validation_years": [
+      1953
+    ],
+    "test_years": [
+      1954
+    ]
+  },
+  "configuration_sources": [
+    "README.md",
+    "conf/config.yaml",
+    "model/graphdop.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/fake_data.py",
+    "configuration.json"
+  ]
+}

configuration.json

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

model/graphdop.py

@@ -0,0 +1,275 @@
+# coding=utf-8
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# Minimal reproduction of GraphDOP (ECMWF, "Towards skilful medium-range
+# forecasts learnt directly from observations", 2025 preprint) following the
+# encoder -- processor -- decoder design:
+#
+#   * Encoder: a GNN that projects gridded "observations" inside the input
+#     window onto a latent mesh (a coarse regular lat/lon grid), using graph
+#     edges with (forward bearing, haversine distance) features.
+#   * Processor: a transformer that advances the latent atmospheric state
+#     forward in time, once per output frame (latent-space rollout).
+#   * Decoder: a GNN that maps the latent mesh back onto the target grid and
+#     predicts per-channel observations with instrument-like output MLPs.
+#
+# Differences from the paper (documented in README.md): the paper consumes
+# irregular, instrument-specific Level-1 observations with dynamic graphs built
+# per batch (PyTorch Geometric); here the OneScience ERA5-h5 gridded pipeline is
+# used as the observation placeholder, and the graphs are fixed regular-grid
+# meshes. The weighted MSE objective is kept (per-channel weights).
+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 between points given 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 GNNLayer(nn.Module):
+    """Message-passing layer with edge features (mean-aggregate, residual)."""
+
+    def __init__(self, dim, edge_dim=2, hidden_dim=64):
+        super().__init__()
+        self.edge_mlp = _mlp(2 * dim + edge_dim, dim, hidden_dim)
+        self.node_mlp = _mlp(dim, dim, hidden_dim)
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x, edge_index, edge_attr):
+        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)
+        msg = self.edge_mlp(torch.cat([xb[src_b], xb[dst_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))
+
+
+class ObsEncoder(nn.Module):
+    """
+    Maps the observation grid onto the latent mesh with a per-cell input MLP,
+    an adaptive pooling to the mesh resolution, and graph message passing.
+    """
+
+    def __init__(self, in_channels, latent_dim, mesh_shape, num_layers=2, hidden_dim=64):
+        super().__init__()
+        self.in_channels = in_channels
+        self.input_mlp = _mlp(in_channels, latent_dim, hidden_dim)
+        self.gnn = nn.ModuleList([GNNLayer(latent_dim, hidden_dim=hidden_dim) 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):
+        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)
+        mesh = mesh.permute(0, 2, 3, 1).reshape(B, -1, mesh.size(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)
+        return mesh
+
+
+class LatentProcessor(nn.Module):
+    """Transformer over latent mesh tokens that advances the state in time."""
+
+    def __init__(self, latent_dim, mesh_shape, num_blocks=1, n_heads=4, hidden_dim=128):
+        super().__init__()
+        n_nodes = mesh_shape[0] * mesh_shape[1]
+        self.pos_emb = nn.Parameter(torch.zeros(1, n_nodes, latent_dim))
+        nn.init.trunc_normal_(self.pos_emb, std=0.02)
+        block = nn.TransformerEncoderLayer(
+            d_model=latent_dim, nhead=n_heads, dim_feedforward=hidden_dim,
+            dropout=0.0, activation="gelu", batch_first=True, norm_first=True,
+        )
+        self.blocks = nn.ModuleList([block for _ in range(num_blocks)])
+
+    def forward(self, mesh):
+        tokens = mesh + self.pos_emb
+        for block in self.blocks:
+            tokens = block(tokens)
+        return tokens
+
+
+class ObsDecoder(nn.Module):
+    """
+    Maps the latent mesh back onto the target grid (bilinear upsample) and
+    predicts per-channel observations with an output MLP.
+    """
+
+    def __init__(self, latent_dim, out_channels, grid_shape, mesh_shape, num_layers=2, hidden_dim=64):
+        super().__init__()
+        self.gnn = nn.ModuleList([GNNLayer(latent_dim, hidden_dim=hidden_dim) 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):
+        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)
+        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 GraphDOP(nn.Module):
+    """
+    Config-driven GraphDOP wrapper.
+
+    Args:
+        in_channels: Number of observation channels per frame.
+        out_channels: Number of forecast channels per frame.
+        input_steps: Number of input (observation window) frames.
+        output_steps: Number of forecast frames.
+        grid_shape: Spatial shape of the (gridded) observation field.
+        mesh_shape: Latent mesh resolution (each dimension, powers of two fine).
+        latent_dim: Feature dimension of latent mesh tokens.
+        num_encoder_layers / num_decoder_layers: GNN message-passing layers.
+        num_processor_blocks: Transformer blocks in the processor.
+        n_heads: Attention heads of the processor.
+        channel_weights: Per-channel weights for the weighted MSE objective.
+    """
+
+    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,
+        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.encoder = ObsEncoder(
+            self.in_channels, int(latent_dim), self.mesh_shape, num_layers=int(num_encoder_layers), hidden_dim=int(hidden_dim)
+        )
+        self.processor = LatentProcessor(
+            int(latent_dim), self.mesh_shape, num_blocks=int(num_processor_blocks), n_heads=int(n_heads), hidden_dim=int(hidden_dim)
+        )
+        self.decoder = ObsDecoder(
+            int(latent_dim), self.out_channels, self.grid_shape, self.mesh_shape,
+            num_layers=int(num_decoder_layers), hidden_dim=int(hidden_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 forward(self, x):
+        """
+        Args:
+            x: Observation frames, shape [batch, input_steps, C, H, W].
+        Returns:
+            Forecast frames, shape [batch, output_steps, C, H, W].
+        """
+        latents = torch.stack([self.encoder(x[:, t]) for t in range(self.input_steps)], dim=0)
+        latent = latents.mean(dim=0)
+        outs = []
+        for _ in range(self.output_steps):
+            latent = self.processor(latent)
+            outs.append(self.decoder(latent))
+        return torch.stack(outs, dim=1)
+
+    def wmse_loss(self, pred, target):
+        """Weighted mean squared error objective (Eq. 1 of the paper)."""
+        diff = (pred - target) ** 2
+        w = self.channel_weights.view(1, 1, self.out_channels, 1, 1)
+        return (diff * w).mean()

scripts/fake_data.py

@@ -0,0 +1,74 @@
+import os
+import h5py
+import numpy as np
+from onescience.utils.YParams import YParams
+
+
+# GraphDOP 的当前 OneScience 适配层把不规则观测投影为规则网格通道。
+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,90 @@
+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.graphdop import GraphDOP
+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 = GraphDOP(
+        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,
+        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/'")
+    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]
+            pred = model(invar).cpu().numpy()                    # [1, output_steps, C, H, W]
+            for t in range(pred.shape[1]):
+                fname = data[4][cfg.input_steps + t][0]          # 该预测帧对应的时刻
+                pred_var = pred[0, 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, cfg in zip(axes, plot_configs):
+        im = ax.imshow(cfg['data'], cmap=cfg['cmap'], vmin=cfg['vmin'], vmax=cfg['vmax'])
+        ax.set_title(cfg['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,222 @@
+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.graphdop import GraphDOP
+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 = GraphDOP(
+        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,
+        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)                           # [B, output_steps, C, H, W]
+            loss = model.wmse_loss(outvar_pred, outvar)          # 论文式(1) WMSE
+            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)
+                loss = model.wmse_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()

weight/.gitkeep

@@ -0,0 +1 @@
+