Upload folder using huggingface_hub

.ms_upload_cache

@@ -0,0 +1 @@
+{"version": 3, "repo_id": "OneScience/Spherical-Fourier-Neural-Operators", "files": {"README.md|1787118190.0|5887": {"hash": "bec28fe4ac1b73be089ebdb23b26f499685ec2fc1e8929dd5d9468381fce8149", "size": 5887, "status": "c"}, "conf/config.yaml|1787189043.0|1979": {"hash": "f6a18e6a5e34e1144994e28c17753ea7e4ef4e2c50377bf500f94f20059b482e", "size": 1979, "status": "c"}, "model/__pycache__/sfno.cpython-311.pyc|1787193344.0|3023": {"hash": "866a2151bf8c896e448c3fcedcd6d5af2ff0e61f84c5a99f526b5f71a056a3ad", "size": 3023, "status": "c"}, "scripts/fake_data.py|1787118238.0|4258": {"hash": "6c7667f81925261a97bd297eade50bfd7afad4c7311e7f48535cbae7975cedb9", "size": 4258, "status": "c"}, "model/sfno.py|1787189031.0|4142": {"hash": "b105e978205217f5b5535fd832c802f1ac12e1f1bbde4da16448a0e0d9021252", "size": 4142, "status": "c"}, "scripts/inference.py|1787037053.0|3214": {"hash": "ee20e9825d9533b521c1e5b955abc29ebc8d636f9339d28dfbbcc972411e809e", "size": 3214, "status": "c"}, "scripts/result.py|1787129353.0|8391": {"hash": "9a5b66ed2b41bbb766b3ce26a286be904d6d2483d8fde2ada6ea6c6044e6aed0", "size": 8391, "status": "c"}, "configuration.json|1787036618.0|38": {"hash": "97d4072fc4a7a3b71e3184609fc491e77a324a982d9df860d9379c009dff84de", "size": 38, "status": "c"}, "scripts/train.py|1787036703.0|9260": {"hash": "ef4f334a62fc97e80b2333e4c0e3664298778275ed3c141c161106133a8646ed", "size": 9260, "status": "c"}, "weight/.gitkeep|1787036758.0|0": {"hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "size": 0, "status": "c"}}}

README.md

@@ -0,0 +1,179 @@
+---
+license: bsd-3-clause
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth science
+- Weather forecasting
+- Short- to medium-range weather forecasting
+- Spherical Fourier neural operators
+- ERA5
+frameworks: PyTorch
+datasets:
+  - OneScience/ERA5
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Spherical Fourier Neural Operators</span>
+  </strong>
+</p>
+
+
+# Model Introduction
+
+SFNO（Spherical Fourier Neural Operator）was proposed jointly by NVIDIA, Caltech, and other institutions，replaces the planar FFT in standard FNOs with spherical harmonic transforms (SHT), making spectral convolution respect spherical geometry and mitigating polar artifacts, spectral artifacts, and instability during long autoregressive rollouts.
+
+Paper:Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere
+
+https://arxiv.org/abs/2306.03838
+
+# Model Description
+
+SFNO is a deterministic, single-state global weather dynamics model: it takes the atmospheric state at one 6-hour time step and outputs the same set of variables at the next 6-hour time step. Longer lead times are obtained through autoregressive rollout.This repository is organized from the official reference implementation in `NVIDIA/torch-harmonics` and integrated with the OneScience data loading and training workflow.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global weather dynamics research | Train a spherical-equivariant neural operator forecasting model on ERA5 data. |
+| Long-horizon autoregressive stability research | Evaluate polar artifacts and dissipation during multi-step rollout. |
+| 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.
+- The model depends on `torch-harmonics`(spherical harmonic transforms)，Confirm that its version matches PyTorch/CUDA before installation.
+
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/SphericalFourierNeuralOperators --local-dir ./SphericalFourierNeuralOperators
+cd SphericalFourierNeuralOperators
+```
+
+### 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] torch-harmonics -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] torch-harmonics -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]` HDF5 data from the model configuration. The current small configuration uses 6 channels and a 32×64 grid, and automatically calculates a time length sufficient for the batch.
+
+### 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. The paper does not clearly release weights for the 26/73-channel weather model; this repository does not currently provide official weights, and users may train the model using the paper configuration.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth`：
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are written 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 model implementation comes from the SFNO reference implementation in the official `NVIDIA/torch-harmonics` (BSD-3-Clause).
+- Commit fetched for the current case directory: `49bac755cd8306fbd27a3604acafa65adf7ca202` (2026-08-14).
+- `conf/config.yaml` uses a small configuration (`img_size=[32, 64]`, `embed_dim=16`, `num_layers=2`) for connectivity validation by default; paper-level reproduction requires a 0.25° (721×1440) grid, 26/73 channels, and a larger network as specified in the paper.
+- The following details are not disclosed in the paper and are assumptions in this reproduction:the internal spectral downsampling factor of the weather model, the positional embedding form, per-variable normalization statistics, and training batch size.
+
+# 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 organization and adaptation of SFNO. The model source is based on the official `torch-harmonics` implementation by Bonev et al. (2023) and follows BSD-3-Clause.
+- Please cite:Bonev, Kurth, Hundt, Pathak, Baust, Kashinath, Anandkumar. Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere. ICML 2023.

conf/config.yaml

@@ -0,0 +1,65 @@
+# SFNO 训练配置示例
+# 论文默认配置为 0.25° 全球网格（721×1440）与 26/73 通道；
+# 当前 img_size/embed_dim 等为连通性验证用的小配置，完整论文复现需按论文调整。
+model:
+  start_epoch: 0
+  max_epoch: 100
+  lr: 1E-3
+  patience: 50
+  checkpoint_dir: "./data/checkpoints"
+
+  # SFNO 网络参数（模型名对应 Bonev et al. 2023 Spherical Fourier Neural Operator）
+  img_size: [32, 64]          # 论文为 [721, 1440]（0.25° ERA5）
+  scale_factor: 2             # 潜在空间频谱降采样倍数（论文天气模型未公开精确值）
+  embed_dim: 16               # 嵌入维度（论文为 256，需大显存）
+  num_layers: 2               # SFNO block 层数（论文约 8 层）
+  activation_function: "gelu"
+  use_mlp: true
+  mlp_ratio: 2.0
+  drop_rate: 0.0
+  drop_path_rate: 0.0
+  normalization_layer: "instance_norm"
+  hard_thresholding_fraction: 1.0
+  residual_prediction: false
+  pos_embed: "none"           # 可选: none / sequence / spectral / learnable lat / learnable latlon
+  bias: false
+
+# 整个数据读取流程
+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, 64]
+    verbose: true
+    cache: false
+
+    # 气象变量（此处为论文 26/73 变量配置的一个小子集，用于连通性验证）
+    channels: ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature',
+               'mean_sea_level_pressure', 'geopotential_500', 'temperature_850']
+
+  # 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,92 @@
+{
+  "model_name": "Spherical Fourier Neural Operators",
+  "model_type": "sfno",
+  "architectures": [
+    "SFNO"
+  ],
+  "framework": "PyTorch",
+  "domain": "climate-and-atmosphere",
+  "task": "global-weather-forecasting",
+  "implementation": {
+    "entry_point": "model/sfno.py",
+    "scope": "configurable wrapper around torch_harmonics.examples.models.sfno.SphericalFourierNeuralOperator for deterministic single-state 6-hour forecasting"
+  },
+  "architecture": {
+    "family": "Spherical Fourier Neural Operator",
+    "spectral_operator": "spherical harmonic transform",
+    "input_format": "B C H W",
+    "prediction": "one atmospheric state to the next 6-hour state; longer forecasts use autoregressive rollout",
+    "activation": "GELU",
+    "normalization": "InstanceNorm",
+    "repository_default_config": {
+      "purpose": "connectivity validation with synthetic ERA5 data",
+      "img_size": [
+        32,
+        64
+      ],
+      "scale_factor": 2,
+      "in_channels": 6,
+      "out_channels": 6,
+      "embed_dim": 16,
+      "num_layers": 2,
+      "use_mlp": true,
+      "mlp_ratio": 2.0,
+      "drop_rate": 0.0,
+      "drop_path_rate": 0.0,
+      "hard_thresholding_fraction": 1.0,
+      "residual_prediction": false,
+      "positional_embedding": "none",
+      "bias": false
+    },
+    "paper_reference_config": {
+      "grid_resolution_degrees": 0.25,
+      "grid_size": [
+        721,
+        1440
+      ],
+      "channels": "26 or 73",
+      "embed_dim": 256,
+      "num_layers": "approximately 8",
+      "note": "the weather model's exact internal spectral downsampling factor is not disclosed"
+    }
+  },
+  "data": {
+    "dataset": "ERA5",
+    "variables": [
+      "10m_u_component_of_wind",
+      "10m_v_component_of_wind",
+      "2m_temperature",
+      "mean_sea_level_pressure",
+      "geopotential_500",
+      "temperature_850"
+    ],
+    "time_step_hours": 6,
+    "input_length": 1,
+    "output_length": 1,
+    "channels": 6,
+    "spatial_size": [
+      32,
+      64
+    ],
+    "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/sfno.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/fake_data.py",
+    "configuration.json"
+  ]
+}

configuration.json

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

model/sfno.py

@@ -0,0 +1,100 @@
+# coding=utf-8
+#
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This module is adapted from the official Spherical Fourier Neural Operator
+# reference implementation of Boris Bonev et al. (ICML 2023), published in the
+# NVIDIA/torch-harmonics repository (BSD-3-Clause). Only a thin configurable
+# wrapper is added so that a single YAML config can drive the model.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice,
+#    this list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its contributors
+#    may be used to endorse or promote products derived from this software
+#    without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+import torch
+import torch.nn as nn
+
+from torch_harmonics.examples.models.sfno import SphericalFourierNeuralOperator
+
+
+class SFNO(nn.Module):
+    """
+    Configurable wrapper around the official SFNO (Spherical Fourier Neural
+    Operator, Bonev et al. 2023, ICML).
+
+    The underlying network is provided by ``torch-harmonics``
+    (``torch_harmonics.examples.models.sfno.SphericalFourierNeuralOperator``),
+    which replaces the plain FFT of a vanilla FNO by a Spherical Harmonic
+    Transform (SHT) so that the learned convolution respects the geometry of
+    the sphere.
+
+    Model inputs / outputs are deterministic global fields of shape
+    ``(Batch, C, H, W)``: a single 6-hour state ``u_t`` is mapped to the next
+    state ``u_{t+1}`` (trained with weighted L2 losses and 1-2 step rollout).
+    """
+
+    def __init__(
+        self,
+        img_size=(32, 64),
+        scale_factor=2,
+        in_chans=4,
+        out_chans=4,
+        embed_dim=16,
+        num_layers=2,
+        activation_function="gelu",
+        use_mlp=True,
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path_rate=0.0,
+        normalization_layer="instance_norm",
+        hard_thresholding_fraction=1.0,
+        residual_prediction=False,
+        pos_embed="none",
+        bias=False,
+    ):
+        super().__init__()
+        self.img_size = tuple(img_size)
+        self.in_chans = int(in_chans)
+        self.out_chans = int(out_chans)
+        self.model = SphericalFourierNeuralOperator(
+            img_size=self.img_size,
+            scale_factor=int(scale_factor),
+            in_chans=self.in_chans,
+            out_chans=self.out_chans,
+            embed_dim=int(embed_dim),
+            num_layers=int(num_layers),
+            activation_function=activation_function,
+            use_mlp=use_mlp,
+            mlp_ratio=mlp_ratio,
+            drop_rate=drop_rate,
+            drop_path_rate=drop_path_rate,
+            normalization_layer=normalization_layer,
+            hard_thresholding_fraction=hard_thresholding_fraction,
+            residual_prediction=residual_prediction,
+            pos_embed=pos_embed,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        return self.model(x)

scripts/fake_data.py

@@ -0,0 +1,108 @@
+import os
+import h5py
+import numpy as np
+import xarray as xr
+from onescience.utils.YParams import YParams
+
+
+# SFNO 单步输入、单步输出；T 保证每个年份至少产生一个完整 batch。
+def get_dims(cfg_model, cfg_data):
+    H, W = map(int, cfg_model.img_size)
+    if tuple(map(int, cfg_data.dataset.img_size)) != (H, W):
+        raise ValueError("model.img_size and datapipe.dataset.img_size must match")
+    batch_size = int(cfg_data.dataloader.batch_size)
+    input_steps, output_steps = 1, 1
+    T = input_steps + output_steps + batch_size - 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 新版读取方式对应。
+    """
+    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")
+
+
+def get_static(data_dir, var, name, H, W):
+    os.makedirs(data_dir, exist_ok=True)
+    ds = xr.Dataset(
+        data_vars={
+            f"{var}": (("valid_time", "latitude", "longitude"),
+                np.random.rand(1, H, W).astype(np.float32))
+        },
+        coords={
+            "valid_time": ["2015-12-31"],
+            "latitude": np.linspace(90, -90, H, dtype=np.float64),
+            "longitude": np.linspace(0, 360 - 360 / W, W, dtype=np.float64),
+            "number": 0,
+            "expver": "",
+        },
+        attrs={
+            "GRIB_centre": "ecmf",
+            "GRIB_centreDescription": "European Centre for Medium-Range Weather Forecasts",
+            "GRIB_subCentre": "0",
+            "Conventions": "CF-1.7",
+            "institution": "European Centre for Medium-Range Weather Forecasts",
+            "history": "Generated manually",
+        }
+    )
+
+    ds.to_netcdf(f"{data_dir}/{name}.nc")
+    arr = np.random.randn(H, W).astype(np.float32)
+    np.save(f'{data_dir}/land_mask.npy', arr)
+    np.save(f'{data_dir}/soil_type.npy', arr)
+    np.save(f'{data_dir}/topography.npy', arr)
+    print(f"✅ Static data: {arr.shape}, dtype: {arr.dtype}, save to {data_dir}")
+
+
+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) == 0:
+        raise ValueError("datapipe.dataset.channels must not be empty")
+
+    dims = get_dims(cfg_model, cfg_datapipe)
+    generate_fake_h5(cfg_datapipe.dataset.data_dir, atm_vars, years, dims)
+
+    static_dir = os.path.join(cfg_datapipe.dataset.data_dir, "static")
+    get_static(static_dir, 'z', 'geopotential', dims["H"], dims["W"])
+    get_static(static_dir, 'lsm', 'land_sea_mask', dims["H"], dims["W"])
+
+    print("\n✅ Fake datasets generated.")

scripts/inference.py

@@ -0,0 +1,89 @@
+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 glob
+import numpy as np
+import h5py
+from tqdm import tqdm
+from model.sfno import SFNO
+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,
+        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 = SFNO(
+        img_size=cfg.img_size,
+        scale_factor=cfg.scale_factor,
+        in_chans=cfg['N_in_channels'],
+        out_chans=cfg['N_out_channels'],
+        embed_dim=cfg.embed_dim,
+        num_layers=cfg.num_layers,
+        activation_function=cfg.activation_function,
+        use_mlp=cfg.use_mlp,
+        mlp_ratio=cfg.mlp_ratio,
+        drop_rate=cfg.drop_rate,
+        drop_path_rate=cfg.drop_path_rate,
+        normalization_layer=cfg.normalization_layer,
+        hard_thresholding_fraction=cfg.hard_thresholding_fraction,
+        residual_prediction=cfg.residual_prediction,
+        pos_embed=cfg.pos_embed,
+        bias=cfg.bias,
+    ).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)
+            filename = data[4][-1][0]
+            pred_var = model(invar).cpu().numpy()
+            pred_var = pred_var * stds + means
+            np.save(f"result/output/{filename}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,200 @@
+import numpy as np
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import os
+import sys
+import glob
+import h5py
+from datetime import datetime
+from tqdm import tqdm
+from onescience.utils.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()
+
+            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()
+
+    ##### 默认绘制 result/output 中第一个预测结果，用户可自行指定日期和变量 #####
+    eg_files = [total_files[0][:-4]]
+    channel_index = [cfg_data.dataset.channels.index(v) for v in ['2m_temperature', 'geopotential_500', 'temperature_850']]
+
+    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()
+        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,224 @@
+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.sfno import SFNO
+from onescience.datapipes.climate import ERA5Datapipe
+from onescience.utils.YParams import YParams
+from onescience.utils.fcn.darcy_loss import LpLoss
+
+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(),
+        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(),
+        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 = SFNO(
+        img_size=cfg.img_size,
+        scale_factor=cfg.scale_factor,
+        in_chans=cfg['N_in_channels'],
+        out_chans=cfg['N_out_channels'],
+        embed_dim=cfg.embed_dim,
+        num_layers=cfg.num_layers,
+        activation_function=cfg.activation_function,
+        use_mlp=cfg.use_mlp,
+        mlp_ratio=cfg.mlp_ratio,
+        drop_rate=cfg.drop_rate,
+        drop_path_rate=cfg.drop_path_rate,
+        normalization_layer=cfg.normalization_layer,
+        hard_thresholding_fraction=cfg.hard_thresholding_fraction,
+        residual_prediction=cfg.residual_prediction,
+        pos_embed=cfg.pos_embed,
+        bias=cfg.bias,
+    ).to(device)
+
+    if _FUSED_ADAM:
+        optimizer = optimizers.FusedAdam(model.parameters(), lr=cfg.lr)
+    else:
+        optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=5, mode='min')
+    loss_obj = LpLoss()
+
+    ## 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 = 1.0e6
+    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)
+            outvar = data[1].to(device, dtype=torch.float32)
+            outvar_pred = model(invar)
+            loss = loss_obj(outvar, outvar_pred)
+            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 = loss_obj(outvar, outvar_pred)
+
+                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'[{(time.time()-start_time)/(j+1): .02f}s/{cfg_data.dataloader.batch_size}batch] '
+                            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()