Upload folder using huggingface_hub

README.md

@@ -0,0 +1,158 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth Science
+- Ocean Forecast
+- Subsurface Forecast
+- Sea Temperature and Salinity Forecast
+- CMEMS
+frameworks: PyTorch
+datasets:
+  - OneScience/CMEMS
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">XiHe</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+XiHe is the world's first data-driven, eddy-resolving global ocean forecast model, jointly proposed by the College of Meteorology and Oceanography at the National University of Defense Technology and multiple universities and research institutions.
+
+Paper: XiHe: A Data-Driven Model for Global Ocean Eddy-Resolving Forecasting
+
+https://arxiv.org/abs/2402.02995
+
+# Model Description
+
+XiHe is a Transformer model designed for high-resolution global ocean forecasting, suitable for global ocean eddy-resolving forecast research.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Global Ocean Forecast Research | Train an XiHe-style ocean forecast model using annual CMEMS HDF5 data. |
+| Local Quick Validation | Use synthetic data to verify data loading, training entry points, inference, and result scripts. |
+| ModelScope / OneCode Execution | Download as a standalone model package, install dependencies, and run scripts directly. |
+| Multi-GPU Training | Launch multi-process training via `torchrun`. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/XiHe --local_dir ./XiHe
+cd XiHe
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data Introduction
+
+The OneScience community provides CMEMS data for training (due to file size limits, the current repository contains a slice of the full dataset). Users can download it with the command below and confirm that the data path in `config/config.yaml` is set correctly:
+
+```bash
+modelscope download --dataset OneScience/CMEMS --local_dir ./data
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```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
+
+This repository provides weights trained on GLORYS12 global ocean reanalysis data in the `weights/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+
+### Inference
+
+Inference reads `data/checkpoints/model_bak.pth` by default:
+
+```bash
+python scripts/inference.py
+```
+
+Prediction results are output to:
+
+```text
+result/output/
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Output contents include:
+
+- `result/rmse.npy`
+- `result/acc.npy`
+- `result/loss.png`
+- Forecast comparison plots for specified dates and variables
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation & License
+
+- This repository is a reproduction of the original XiHe paper.

conf/config.yaml

@@ -0,0 +1,141 @@
+# ERA5数据集配置示例
+model:
+  start_epoch: 0
+  max_epoch: 9999
+  lr: 5E-5          #学习率
+  
+  betas: [0.9, 0.95]
+  weight_decay: 0.01
+  
+  embed_dim: 192
+  num_blocks: 8 
+  patch_size: [6, 12]
+  patience: 50 
+  checkpoint_dir: "./data/checkpoints" #相对路径
+
+  in_chans: 96
+  out_chans: 96
+  num_groups: 64
+  depth: 2
+  num_heads:  [6, 12, 12, 6]
+  img_size: [2041, 4320]
+  mask: "./data/land_mask.npy"
+  
+
+
+datapipe:
+  name: "CMEMS"
+  task: "ocean_forecasting"
+
+  dataset:
+    type: "hdf5"
+    stats_dir: './data/stats/' # "$ONESCIENCE_DATASETS_DIR/ERA5/newh5/stats/"  
+    static_dir: './data/static/' # "$ONESCIENCE_DATASETS_DIR/ERA5/newh5/static/" 
+    data_dir: './data/' # "$ONESCIENCE_DATASETS_DIR/ERA5/newh5/" 
+
+    
+    train_time: [2010]
+    val_time: [2013]
+    test_time: [2014]
+    img_size: [2041, 4320]
+    verbose: true
+    cache: false
+
+    # 气象变量
+
+    channels: ["sea_surface_height_above_geoid","10m_u_component_of_wind","10m_v_component_of_wind", "sea_surface_temperature",
+      "sea_water_potential_temperature_1",sea_water_salinity_1,"eastward_sea_water_velocity_1","northward_sea_water_velocity_1",
+      "sea_water_potential_temperature_2",sea_water_salinity_2,"eastward_sea_water_velocity_2","northward_sea_water_velocity_2",
+      "sea_water_potential_temperature_3",sea_water_salinity_3,"eastward_sea_water_velocity_3","northward_sea_water_velocity_3",
+      "sea_water_potential_temperature_4",sea_water_salinity_4,"eastward_sea_water_velocity_4","northward_sea_water_velocity_4",
+      "sea_water_potential_temperature_5",sea_water_salinity_5,"eastward_sea_water_velocity_5","northward_sea_water_velocity_5",
+      "sea_water_potential_temperature_6",sea_water_salinity_6,"eastward_sea_water_velocity_6","northward_sea_water_velocity_6",
+      "sea_water_potential_temperature_7",sea_water_salinity_7,"eastward_sea_water_velocity_7","northward_sea_water_velocity_7",
+      "sea_water_potential_temperature_8",sea_water_salinity_8,"eastward_sea_water_velocity_8","northward_sea_water_velocity_8",
+      "sea_water_potential_temperature_9",sea_water_salinity_9,"eastward_sea_water_velocity_9","northward_sea_water_velocity_9",
+      "sea_water_potential_temperature_10",sea_water_salinity_10,"eastward_sea_water_velocity_10","northward_sea_water_velocity_10",
+      "sea_water_potential_temperature_11",sea_water_salinity_11,"eastward_sea_water_velocity_11","northward_sea_water_velocity_11",
+      "sea_water_potential_temperature_12",sea_water_salinity_12,"eastward_sea_water_velocity_12","northward_sea_water_velocity_12",
+      # "sea_surface_height_above_geoid","10m_u_component_of_wind","10m_v_component_of_wind", "sea_surface_temperature",
+      "sea_water_potential_temperature_13",sea_water_salinity_13,"eastward_sea_water_velocity_13","northward_sea_water_velocity_13",
+      "sea_water_potential_temperature_14",sea_water_salinity_14,"eastward_sea_water_velocity_14","northward_sea_water_velocity_14",
+      "sea_water_potential_temperature_15",sea_water_salinity_15,"eastward_sea_water_velocity_15","northward_sea_water_velocity_15",
+      "sea_water_potential_temperature_16",sea_water_salinity_16,"eastward_sea_water_velocity_16","northward_sea_water_velocity_16",      
+      "sea_water_potential_temperature_17",sea_water_salinity_17,"eastward_sea_water_velocity_17","northward_sea_water_velocity_17",
+      "sea_water_potential_temperature_18",sea_water_salinity_18,"eastward_sea_water_velocity_18","northward_sea_water_velocity_18",
+      "sea_water_potential_temperature_19",sea_water_salinity_19,"eastward_sea_water_velocity_19","northward_sea_water_velocity_19",
+      "sea_water_potential_temperature_20",sea_water_salinity_20,"eastward_sea_water_velocity_20","northward_sea_water_velocity_20",
+      "sea_water_potential_temperature_21",sea_water_salinity_21,"eastward_sea_water_velocity_21","northward_sea_water_velocity_21",
+      "sea_water_potential_temperature_22",sea_water_salinity_22,"eastward_sea_water_velocity_22","northward_sea_water_velocity_22",
+      "sea_water_potential_temperature_23",sea_water_salinity_23,"eastward_sea_water_velocity_23","northward_sea_water_velocity_23",
+    ]
+
+     
+    
+
+    
+    variables:
+      - "u10"        # 10m U wind component
+      - "v10"        # 10m V wind component
+      - "uo"        #  'uo': 'eastward_sea_water_velocity',
+      - "vo"        # 'vo': 'northward_sea_water_velocity',
+      - "so"        #  'so': 'sea_water_salinity',
+      - "thetao"        #    'thetao': 'sea_water_potential_temperature',
+      - "ssh"       # 'zos': 'sea_surface_height_above_geoid'
+      - "sst"       # 'analysed_sst': 'sea_surface_temperature',
+             
+    
+    # 时间配置
+    time_range: ["2010-01-01", "2014-12-31"] 
+    time_steps: 1
+    time_res: 24
+    
+    # 空间配置
+    spatial_resolution: [1/12, 1/12]
+    
+    # 采样配置
+    num_samples: -1  # -1 表示使用全部数据
+    shuffle: true
+    random_seed: 42
+    
+    # 领域特定配置
+    extra:
+      levels: [0.495, 2.65, 5.08, 7.93, 11.41, 15.81, 21.60, 29.44,40.34, 55.76, 77.85, 92.32, 109.73, 130.67, 155.85,186.13, 222.48, 266.04, 318.13, 380.21, 453.94,541.09, 643.57]
+      lat_range: [-90, 90]
+      lon_range: [0, 360]
+  
+  # 数据转换配置
+  transforms:
+    - type: "Normalize"
+      params:
+        mean: [0.0, 0.0, 288.0, 101325.0, 50000.0, 270.0]
+        std: [5.0, 5.0, 15.0, 1000.0, 5000.0, 10.0]
+        keys: ["input", "target"]
+    
+    - type: "ToTensor"
+      params:
+        keys: null  # null表示转换所有numpy数组
+  
+  # 其他配置
+
+
+  # DataLoader配置
+  dataloader:
+    mask_dtype: "float32"
+    batch_size: 1
+    num_workers: 0
+    pin_memory: false
+    drop_last: true
+    shuffle: false  # 使用sampler时设为false
+    # prefetch_factor: 2
+    persistent_workers: false 
+  
+  # 分布式配置
+  distributed:
+    enabled: false #开关
+    sampler: "DistributedSampler"
+    rank: 0
+    world_size: 8
+    shuffle: true
+    seed: 42
+    drop_last: true

configuration.json

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

model/xihe.py

@@ -0,0 +1,180 @@
+import math
+import os
+import torch
+import sys
+import numpy as np
+import torch.nn as nn
+from dataclasses import dataclass
+from onescience.utils.YParams import YParams
+from onescience.models.meta import ModelMetaData
+from onescience.modules.embedding.xiheembedding import XiheEmbedding
+from onescience.modules.fuser.xihefuse import XiheFuser
+from onescience.modules.recovery.xihepatchrecovery import XihePatchRecovery
+from onescience.modules.sample.pangudownsample import PanguDownSample
+from onescience.modules.sample.xiheupsample import XiheUpSample
+
+
+
+@dataclass
+class MetaData(ModelMetaData):
+    name: str = "Xihe"
+    # Optimization
+    jit: bool = False  # ONNX Ops Conflict
+    cuda_graphs: bool = True
+    amp: bool = True
+    
+    # Inference
+    onnx_cpu: bool = False  # No FFT op on CPU
+    onnx_gpu: bool = True
+    onnx_runtime: bool = True
+    
+    # Physics informed
+    var_dim: int = 1
+    func_torch: bool = False
+    auto_grad: bool = False
+    
+class TensorWithMask:
+    def __init__(self, x, mask):
+        self.x = x
+        self.mask = mask
+        self.y=None
+
+   
+class Xihe(nn.Module):
+    """
+    Xihe A PyTorch impl of: `XiHe: A Data-Driven Model for Global Ocean Eddy-Resolving Forecasting`
+    https://arxiv.org/abs/2402.02995
+    """
+    def __init__(
+        self,
+        config,
+        img_size=(2041, 4320),     
+        patch_size=(6, 12),        
+        window_size=(6, 12),       
+        embed_dim=192,
+        num_heads=(6, 12, 12, 6),
+        in_chans=96,
+        depth=1,
+        mask_full=None,
+        out_chans=94,
+        num_groups=32,
+        
+    ):
+        super().__init__()
+        self.img_size = config.img_size
+        self.patch_size =config.patch_size
+        # 正确初始化 mask_full
+        self.mask=config.mask
+        mask_full = np.load(self.mask)
+        self.mask_full = mask_full if mask_full is not None else None
+        self.mask_h_w=None
+        self.out_chans=config.out_chans
+        self.in_chans=config.in_chans
+        self.num_groups=config.num_groups
+        self.embed_dim=config.embed_dim
+        
+        self.skip_proj = nn.Linear(2*self.embed_dim, self.embed_dim)
+
+
+        self.patchembed2d = XiheEmbedding()
+        self.patchrecovery2d = XihePatchRecovery()
+
+        # patch 后的 3D 分辨率: (Pl=1, Lat_out, Lon_out)
+        H_out = math.ceil(img_size[0] / patch_size[0])
+        W_out = math.ceil(img_size[1] / patch_size[1])
+        input_resolution = (1, H_out, W_out)
+        
+        self.mask_h_w=input_resolution
+        # 3D 窗口：把 2D 窗口扩成 (1, win_lat, win_lon)
+        window_size_3d = (1, window_size[0], window_size[1]) #  window_size (tuple[int]): Window size [pressure levels, latitude, longitude].
+
+        # 防止过拟合，随机丢弃一部分
+        if depth > 1:
+            drop_path = np.linspace(0, 0.2, depth).tolist()
+        else:
+            drop_path = 0.0
+
+        self.block1=XiheFuser(dim=self.embed_dim,input_resolution=input_resolution,num_local=1)
+
+        self.downsample = PanguDownSample(
+                                    in_dim=self.embed_dim,
+                                    input_resolution=(H_out, W_out),
+                                    output_resolution=(H_out // 2, W_out // 2))
+        
+        input_resolution = (1, H_out // 2, W_out // 2)
+        self.mask_h_w=input_resolution
+
+        self.block2=XiheFuser(dim=2*self.embed_dim,input_resolution=input_resolution,num_local=2)
+
+        self.block3=XiheFuser(dim=2*self.embed_dim,input_resolution=input_resolution,num_local=2)
+        self.block4=XiheFuser(dim=2*self.embed_dim,input_resolution=input_resolution,num_local=2)
+
+        self.upsample = XiheUpSample(in_dim=2*self.embed_dim,out_dim=embed_dim,input_resolution=(H_out // 2, W_out // 2),  output_resolution=(H_out, W_out), )
+        input_resolution = (1, H_out, W_out)
+        self.block5=XiheFuser(dim=self.embed_dim,input_resolution=input_resolution,num_local=1)
+    def change_mask(self,mask_full, x, h_out, w_out):
+        #根据当前层特征分辨率，自动生成掩码（海洋=1，陆地=0）
+            if not torch.is_tensor(mask_full):
+                mask_full = torch.tensor(mask_full, dtype=torch.float32)
+            else:
+                mask_full = mask_full
+
+            H, W = mask_full.shape
+            patch_h = math.ceil(H / h_out)
+            patch_w = math.ceil(W / w_out)
+
+            mask_coarse = torch.zeros((h_out, w_out), dtype=torch.float32)
+            for i in range(h_out):
+                for j in range(w_out):
+                    h0, h1 = i * patch_h, min((i + 1) * patch_h, H)
+                    w0, w1 = j * patch_w, min((j + 1) * patch_w, W)
+                    patch = mask_full[h0:h1, w0:w1]
+                    mask_coarse[i, j] = 1.0 if torch.any(patch > 0.5) else 0.0
+            
+            mask_coarse = mask_coarse.to(x.device, dtype=x.dtype) 
+            B = x.shape[0]                
+            mask_coarse = mask_coarse.unsqueeze(0).unsqueeze(0).repeat(B, 1, 1, 1) #broadcast
+            return mask_coarse  
+        
+    def forward(self, x: torch.Tensor):     
+        x = self.patchembed2d(x)                  # (B, C=embed_dim, H', W')
+        x = x.flatten(2).transpose(1, 2)          # (B, N=H'*W', C) 
+        B, N, C = x.shape     
+        mask_full=self.mask_full        
+        
+       
+        if mask_full is not None:              # mask1
+            H_out = math.ceil(self.img_size[0] / self.patch_size[0])
+            W_out = math.ceil(self.img_size[1] / self.patch_size[1])
+            mask1 = self.change_mask(mask_full, x, h_out=H_out, w_out=W_out)
+        else:
+            mask1 = None
+        
+        obj1 = TensorWithMask(x, mask1)
+        x=self.block1(obj1)          # (B, N, C) 经过 3D 全局注意力
+        x1=x
+        x=self.downsample(x)                 # (B, N, C) 经过 2D 下采样
+        
+        if mask_full is not None:            #  mask2
+            _,H_out,W_out = self.mask_h_w
+            mask2 = self.change_mask(mask_full, x, h_out=H_out, w_out=W_out)
+        else:
+            mask2 = None
+        obj2 = TensorWithMask(x, mask2)
+        x=self.block2(obj2)   
+        obj2 = TensorWithMask(x, mask2)
+        x=self.block3(obj2)  
+        obj2 = TensorWithMask(x, mask2)
+        x=self.block4(obj2) 
+        x=self.upsample(x) 
+        obj1 = TensorWithMask(x, mask1)
+        x=self.block5(obj1)
+        x_out = torch.cat([x, x1], dim=-1)         # (B, N, 2C)
+        x_out = self.skip_proj(x_out)
+        # B, N, C = x.shape
+        # H_, W_ = 341, 360   # 对应 patch grid 尺寸
+        H_ = math.ceil(self.img_size[0] / self.patch_size[0])
+        W_ = math.ceil(self.img_size[1] / self.patch_size[1])
+        x_out = x_out.transpose(1, 2).reshape(B, C, H_, W_)
+        x=self.patchrecovery2d(x_out)
+        return x

scripts/fake_data.py

@@ -0,0 +1,88 @@
+import json
+import os
+
+import h5py
+import numpy as np
+
+from onescience.utils.YParams import YParams
+
+
+DATASET_DIMS = {"T": 10, "H": 2041, "W": 4320, "time_step": 24}
+
+
+def generate_fake_h5(data_dir, var_names, years, dims):
+    os.makedirs(os.path.join(data_dir, "data"), exist_ok=True)
+    t_dim, c_dim = dims["T"], len(var_names)
+    h_dim, w_dim = dims["H"], dims["W"]
+
+    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_dim, c_dim, h_dim, w_dim),
+                dtype="float32",
+                chunks=(1, c_dim, h_dim, w_dim),
+                fillvalue=0.0,
+            )
+            ds.attrs["variables"] = var_names
+            ds.attrs["time_step"] = dims["time_step"]
+
+        size_kb = os.path.getsize(path) / 1024
+        print(
+            f"  {year}.h5  shape=({t_dim},{c_dim},{h_dim},{w_dim})  "
+            f"logical={t_dim * c_dim * h_dim * w_dim * 4 / 1024**3:.1f}GB  actual={size_kb:.1f}KB"
+        )
+
+
+def generate_metadata(data_dir, var_names, years):
+    metadata = {
+        "years": [str(year) for year in years],
+        "variables": var_names,
+    }
+    with open(os.path.join(data_dir, "metadata.json"), "w") as f:
+        json.dump(metadata, f, indent=2)
+    print(f"  metadata saved -> {os.path.join(data_dir, 'metadata.json')}")
+
+
+def generate_stats(stats_dir, n_vars):
+    os.makedirs(stats_dir, exist_ok=True)
+    shape = (1, n_vars, 1, 1)
+    np.save(os.path.join(stats_dir, "global_means.npy"), np.zeros(shape, dtype=np.float32))
+    np.save(os.path.join(stats_dir, "global_stds.npy"), np.ones(shape, dtype=np.float32))
+    print(f"  stats saved -> {stats_dir}")
+
+
+def generate_mask(save_path, shape, one_ratio=0.7, seed=42):
+    np.random.seed(seed)
+    mask = (np.random.rand(*shape) < one_ratio).astype(np.float32)
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    np.save(save_path, mask)
+    print(f"  mask saved -> {save_path}")
+
+
+def main():
+    cfg_model = YParams("conf/config.yaml", "model")
+    cfg_datapipe = YParams("conf/config.yaml", "datapipe")
+
+    if cfg_datapipe.dataset.data_dir.startswith("/public/onestore"):
+        print("Please check config and ensure local test paths are used.")
+        exit()
+
+    years = (
+        cfg_datapipe.dataset.train_time
+        + cfg_datapipe.dataset.val_time
+        + cfg_datapipe.dataset.test_time
+    )
+    channels = cfg_datapipe.dataset.channels
+
+    generate_fake_h5(cfg_datapipe.dataset.data_dir, channels, years, DATASET_DIMS)
+    generate_metadata(cfg_datapipe.dataset.data_dir, channels, years)
+    generate_stats(cfg_datapipe.dataset.stats_dir, len(channels))
+    generate_mask(cfg_model.mask, (DATASET_DIMS["H"], DATASET_DIMS["W"]))
+
+    print("\n✅ Fake CMEMS datasets generated.")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,75 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+
+import glob
+import os
+
+import h5py
+import numpy as np
+import torch
+
+from tqdm import tqdm
+
+from onescience.datapipes.climate import CMEMSDatapipe
+from model.xihe import Xihe
+from onescience.utils.YParams import YParams
+
+
+def get_stats(data_dir, stats_dir, channels):
+    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"]]
+
+    channel_indices = [all_variables.index(v) for v in channels]
+    mu = np.load(os.path.join(stats_dir, "global_means.npy"))
+    std = np.load(os.path.join(stats_dir, "global_stds.npy"))
+    means = mu[:, channel_indices, :, :]
+    stds = std[:, channel_indices, :, :]
+    return means, stds
+
+
+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")
+
+    means, stds = get_stats(
+        cfg_data.dataset.data_dir,
+        cfg_data.dataset.stats_dir,
+        cfg_data.dataset.channels,
+    )
+
+    datapipe = CMEMSDatapipe(
+        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=cfg_data.dataloader.num_workers,
+    )
+    test_dataloader, _ = datapipe.get_dataloader("test")
+
+    ckpt = torch.load(f"{cfg.checkpoint_dir}/model_bak.pth", map_location="cuda:0")
+    model = Xihe(cfg).to("cuda:0")
+    model.load_state_dict(ckpt["model_state_dict"])
+
+    model.eval()
+    os.makedirs("result/output/", exist_ok=True)
+    print("📂 samples 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("cuda:0", dtype=torch.float32)
+            filename = data[4][-1][0]
+
+            outvar_pred = model(invar)
+            pred_var = outvar_pred.cpu().numpy()
+            pred_var = pred_var * stds + means
+            np.save(f"result/output/{filename}.npy", pred_var)

scripts/result.py

@@ -0,0 +1,208 @@
+import glob
+import os
+import sys
+
+from datetime import datetime
+
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+
+from matplotlib import rcParams
+from tqdm import tqdm
+
+from onescience.utils.fcn.YParams import YParams
+
+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_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 = sorted(f for f in os.listdir("./result/output/") if f.endswith(".npy"))
+    return total_files, channel_indices, time_step
+
+
+def filename_to_index(filename, time_step):
+    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]
+            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]
+                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 = 36
+
+    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)
+
+    mu = np.load(os.path.join(cfg_data.dataset.stats_dir, "global_means.npy"))
+    clim_mean = mu[:, channel_indices, :, :]
+    get_result(total_files, channel_indices, time_step, data_dir, clim_mean)
+    show_result()
+
+    test_year = cfg_data.dataset.test_time[0]
+    eg_files = [f"{test_year}010200"]
+    channel_index = [
+        cfg_data.dataset.channels.index(v)
+        for v in [
+            "sea_surface_height_above_geoid",
+            "sea_water_potential_temperature_1",
+            "sea_water_salinity_4",
+        ]
+    ]
+
+    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]
+            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,221 @@
+import sys
+from pathlib import Path
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.append(str(root_path))
+
+import logging
+import os
+import time
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from torch.nn.parallel import DistributedDataParallel
+
+from onescience.datapipes.climate import CMEMSDatapipe
+from model.xihe import Xihe
+from onescience.utils.YParams import YParams
+from onescience.utils.fcn.darcy_loss import LpLoss
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    logger = logging.getLogger()
+
+    config_file_path = os.path.join(current_path, "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+
+    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:
+        dist.init_process_group(backend="nccl", init_method="env://")
+        local_rank = int(os.environ["LOCAL_RANK"])
+        world_rank = dist.get_rank()
+
+    cfg_data = YParams(config_file_path, "datapipe")
+    datapipe = CMEMSDatapipe(
+        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 = CMEMSDatapipe(
+        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 = Xihe(config=cfg).to(local_rank)
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=cfg.lr,
+        betas=tuple(cfg.betas),
+        weight_decay=cfg.weight_decay,
+    )
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, factor=0.2, patience=5, mode="min"
+    )
+    loss_obj = LpLoss()
+
+    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)
+
+    if cfg.world_size == 1:
+        total_params = sum(p.numel() for p in model.parameters())
+        print("\n")
+        print("-" * 50)
+        print(f"📂 now params is {total_params}, {total_params / 1e6:.2f}M, {total_params / 1e9:.2f}B")
+        print("-" * 50, "\n")
+
+    if os.path.exists(f"{cfg.checkpoint_dir}/model_bak.pth"):
+        if world_rank == 0:
+            print("\n")
+            print("-" * 50)
+            print("✅ 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=f"cuda:{local_rank}",
+            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)
+
+    if cfg.world_size > 1:
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+
+    world_rank == 0 and logger.info("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.0
+        start_time = time.time()
+        for j, data in enumerate(train_dataloader):
+            invar = data[0].to(local_rank, dtype=torch.float32)
+            outvar = data[1].to(local_rank, 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.0
+        val_start_time = time.time()
+        with torch.no_grad():
+            for j, data in enumerate(val_dataloader):
+                invar = data[0].to(local_rank, dtype=torch.float32)
+                outvar = data[1].to(local_rank, dtype=torch.float32)
+                outvar_pred = model(invar)
+                loss = loss_obj(outvar, outvar_pred)
+
+                if cfg.world_size > 1:
+                    loss_tensor = loss.detach().to(local_rank)
+                    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() - val_start_time) // 60):02}:{int((time.time() - val_start_time) % 60):02}] "
+                        f"[{(time.time() - val_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")
+    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 @@
+