Publish CRAI-ClimateExtremes reproduction

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,149 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Climate Extremes
+- Spatial Reconstruction
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">CRAI-ClimateExtremes</span></strong>
+</p>
+
+# Model Introduction
+
+CRAI-ClimateExtremes reconstructs missing regions from incomplete climate-extreme index fields and validity masks for historical-extreme analysis and sparse-observation reconstruction.
+
+Paper: Artificial intelligence reveals past climate extremes by reconstructing historical records  
+https://doi.org/10.1038/s41467-024-53464-2
+
+# Model Description
+
+The method was proposed by research teams from the German Climate Computing Center, the Met Office, the University of Hamburg, and collaborating institutions. The paper uses 45 historical simulations from eight CMIP6 models, HadEX-CAM observations, and reanalysis products including ERA5. The model reconstructs monthly TX90p, TN90p, TX10p, and TN10p temperature-extreme indices.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Climate-extreme reconstruction | Reconstruct four monthly temperature-extreme indices from incomplete fields and validity masks. |
+| Irregular-gap handling | Validate partial convolution, mask propagation, and missing-region loss. |
+| Local engineering validation | Validate the end-to-end workflow with structured synthetic data on the complete `144×192` grid. |
+| ModelScope/OneCode execution | Validate structured data, training, inference, climate-reconstruction metrics, and visualization in ModelScope or OneCode environments. |
+| Multi-GPU training | Validate distributed training and the checkpoint workflow through `torchrun`. |
+
+# Usage Instructions
+
+## 1.OneCode
+
+Experience intelligent, one-click AI4S programming through the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Download and Installation
+
+```bash
+hf download OneScience-Group/CRAI-ClimateExtremes --local-dir ./CRAI-ClimateExtremes
+cd CRAI-ClimateExtremes
+```
+
+### Environment Dependencies
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for connectivity validation with the default small-sample configuration.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+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
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+The paper uses monthly CMIP6 extreme indices from 1901-2014 and applies HadEX-CAM missing-data masks to construct training samples. Inputs contain an index field and validity mask `[B,2,144,192]`, while targets are complete `[B,1,144,192]` index fields. This repository uses a small structured synthetic dataset for engineering validation only and does not represent the real CMIP6 or HadEX-CAM distributions, training scale, or paper performance.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+For single-GPU training, use:
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```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
+```
+
+The default configuration reduces samples, network width, epochs, and ensemble members without reducing the `144×192` grid or four-index protocol. Formal experiments require real CMIP6 and HadEX-CAM data and paper-scale training, with artifacts saved to:
+
+```text
+result/checkpoints/crai_climateextremes.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+The paper does not provide directly loadable official model weights, and this repository bundles no weights under `weight/`. The locally trained checkpoint is saved to `result/checkpoints/crai_climateextremes.pt` and must not be represented as an official pretrained weight.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the unified ensemble checkpoint, reconstructs missing extreme indices, and saves the member mean and spread. Complete numerical outputs are saved to:
+
+```text
+result/output/predictions.npz
+result/output/metadata.json
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation computes missing-region RMSE, Spearman correlation, bias, and neighborhood spatial correlation and generates a target, input, and reconstruction comparison. Synthetic-data results validate engineering only and do not represent paper performance; outputs are saved to:
+
+```text
+result/evaluation/metrics.json
+result/evaluation/comparison.png
+```
+
+# 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 engineering reproduction of the public CRAI-ClimateExtremes specifications.
+
+Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.

README_zh.md

@@ -0,0 +1,150 @@
+---
+license: apache-2.0
+language:
+- zh
+- en
+tags:
+- OneScience
+- 地球科学
+- 气候极端
+- 空间重建
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">CRAI-ClimateExtremes</span></strong>
+</p>
+
+# 模型介绍
+
+CRAI-ClimateExtremes 根据不完整的气候极端指数场及其有效掩码重建缺测区域，主要用于历史气候极端事件分析和稀疏观测资料重建。
+
+论文：Artificial intelligence reveals past climate extremes by reconstructing historical records  
+https://doi.org/10.1038/s41467-024-53464-2
+
+# 模型描述
+
+该方法由德国气候计算中心、英国气象局和汉堡大学等机构的研究团队提出。论文使用 8 个 CMIP6 模式的 45 个历史模拟、HadEX-CAM 观测资料以及 ERA5 等再分析数据。模型适用于重建 TX90p、TN90p、TX10p 和 TN10p 月尺度温度极端指数。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 气候极端指数重建 | 从不完整指数场和有效掩码重建四类月尺度温度极端指数。 |
+| 不规则缺测处理 | 验证 partial convolution、掩码传播和缺失区域损失。 |
+| 本地工程验证 | 使用完整 `144×192` 网格的结构化虚拟数据验证端到端流程。 |
+| ModelScope/OneCode 运行 | 在 ModelScope 或 OneCode 环境中验证结构化数据、训练、推理、气候重建指标和可视化流程。 |
+| 多卡训练 | 通过 `torchrun` 验证分布式训练和 checkpoint 流程。 |
+
+# 使用说明
+
+## 1.OneCode
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/CRAI-ClimateExtremes --local_dir ./CRAI-ClimateExtremes
+cd CRAI-ClimateExtremes
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 GPU 或 DCU 运行。
+- CPU 可用于默认小样本配置的连通性验证。
+- DCU 用户需预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。
+
+**DCU环境**
+
+```bash
+# 请首先激活 DTK 及 Conda
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU环境**
+
+```bash
+# 请首先激活 Conda
+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
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### 训练数据介绍
+
+论文使用 1901-2014 年 CMIP6 月尺度极端指数，并应用 HadEX-CAM 缺测掩码构造训练样本。输入为指数场与有效掩码 `[B,2,144,192]`，目标为 `[B,1,144,192]` 的完整指数场。本仓库使用少量结构化虚拟样本验证工程流程，不代表 CMIP6 或 HadEX-CAM 的真实数据分布、训练规模或论文性能。
+
+```bash
+python scripts/fake_data.py
+```
+
+### 训练
+
+单卡训练可使用：
+
+```bash
+python scripts/train.py
+```
+
+多卡训练可使用：
+
+```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
+```
+
+默认配置仅缩小样本数、网络宽度、训练轮数和 ensemble 成员数，不缩小 `144×192` 空间网格和四类指数协议。正式实验需要真实 CMIP6 与 HadEX-CAM 数据和论文规模训练配置，训练产物保存到：
+
+```text
+result/checkpoints/crai_climateextremes.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+论文未提供可直接加载的官方模型权重，本仓库不在 `weight/` 中内置权重。本地训练生成的 checkpoint 保存到 `result/checkpoints/crai_climateextremes.pt`，不得描述为官方预训练权重。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载统一 ensemble checkpoint，重建缺测极端指数并保存成员均值和离散度。完整数值结果保存到：
+
+```text
+result/output/predictions.npz
+result/output/metadata.json
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估计算缺失区域 RMSE、Spearman 相关系数、bias 和邻域空间相关，并生成目标、输入和重建结果对比图。虚拟数据结果仅用于工程验证，不代表论文正式性能；结果保存到：
+
+```text
+result/evaluation/metrics.json
+result/evaluation/comparison.png
+```
+
+# OneScience官方信息
+
+| 平台 | OneScience 主仓库 | Skills 仓库 |
+| --- | --- | --- |
+| 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 |
+
+# 引用与许可证
+
+本仓库为 CRAI-ClimateExtremes 公开规格的独立工程复现版本��
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,18 @@
+seed: 42
+data_path: data/crai_fake.npz
+height: 144
+width: 192
+base_channels: 8
+batch_size: 2
+epochs: 1
+max_iterations: 8
+learning_rate: 0.00005
+ensemble_members: 2
+paper_batch_size: 16
+paper_iterations: 1000000
+paper_ensemble_members: 20
+paper_model: false
+num_samples: 8
+checkpoint_path: result/checkpoints/crai_climateextremes.pt
+output_dir: result/output
+evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,38 @@
+{
+  "model_name": "CRAI-ClimateExtremes",
+  "model_type": "crai_climateextremes",
+  "architectures": ["CRAIClimateExtremes"],
+  "framework": "PyTorch",
+  "domain": "earth-science",
+  "task": "climate-extreme-index-reconstruction",
+  "implementation": {
+    "entry_point": "model/crai_climateextremes.py",
+    "scope": "core-method and full-grid reduced-sample engineering reproduction",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "input_shape": ["B", 2, 144, 192],
+    "output_shape": ["B", 1, 144, 192],
+    "core": "partial-convolution U-Net with mask propagation and skip connections",
+    "output_range_percent": [0, 100]
+  },
+  "data": {
+    "dataset": "HadEX-CAM and CMIP6-derived extreme indices",
+    "format_version": "1.0",
+    "resolution_degrees": [1.25, 1.875],
+    "indices": ["TX90p", "TN90p", "TX10p", "TN10p"],
+    "input_layout": "BCHW",
+    "synthetic": true
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/crai_climateextremes.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "climate extreme index reconstruction",
+  "model": "CRAI-ClimateExtremes",
+  "input_format": "BCHW: incomplete index field and validity mask [B,2,144,192]",
+  "protocol": "partial-convolution reconstruction of masked HadEX-resolution climate-extreme fields",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/crai_climateextremes.py

@@ -0,0 +1,93 @@
+"""Independent PyTorch implementation of the CRAI reconstruction architecture."""
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class PartialConv2d(nn.Module):
+    """Convolution normalized by valid pixels, with the corresponding mask update."""
+
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
+        super().__init__()
+        padding = kernel_size // 2
+        self.input_conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
+        self.kernel_area = float(kernel_size * kernel_size)
+        self.stride = stride
+        self.padding = padding
+        self.kernel_size = kernel_size
+        self.register_buffer("mask_kernel", torch.ones(1, 1, kernel_size, kernel_size))
+
+    def forward(self, values, mask):
+        if mask.shape[1] != 1:
+            mask = mask.amax(dim=1, keepdim=True)
+        with torch.no_grad():
+            valid_count = F.conv2d(
+                mask, self.mask_kernel, stride=self.stride, padding=self.padding
+            )
+            new_mask = (valid_count > 0).to(values.dtype)
+            scale = self.kernel_area / valid_count.clamp_min(1.0)
+        raw = self.input_conv(values * mask)
+        if self.input_conv.bias is not None:
+            bias = self.input_conv.bias.view(1, -1, 1, 1)
+            raw = (raw - bias) * scale + bias
+        else:
+            raw = raw * scale
+        return raw * new_mask, new_mask
+
+
+class PConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super().__init__()
+        self.partial = PartialConv2d(in_channels, out_channels, 3, stride)
+        self.norm = nn.GroupNorm(1, out_channels)
+
+    def forward(self, values, mask):
+        values, mask = self.partial(values, mask)
+        return F.leaky_relu(self.norm(values), 0.2), mask
+
+
+class CRAIClimateExtremes(nn.Module):
+    """Partial-convolution U-Net producing bounded extreme-index percentages."""
+
+    def __init__(self, base_channels=8):
+        super().__init__()
+        c = base_channels
+        self.enc1 = PConvBlock(1, c)
+        self.enc2 = PConvBlock(c, 2 * c, stride=2)
+        self.enc3 = PConvBlock(2 * c, 4 * c, stride=2)
+        self.bottleneck = PConvBlock(4 * c, 8 * c, stride=2)
+        self.dec3 = PConvBlock(8 * c + 4 * c, 4 * c)
+        self.dec2 = PConvBlock(4 * c + 2 * c, 2 * c)
+        self.dec1 = PConvBlock(2 * c + c, c)
+        self.output = nn.Conv2d(c, 1, 1)
+
+    @staticmethod
+    def _upsample(values, mask, size):
+        return (
+            F.interpolate(values, size=size, mode="bilinear", align_corners=False),
+            F.interpolate(mask, size=size, mode="nearest"),
+        )
+
+    def forward(self, inputs):
+        if inputs.ndim != 4 or inputs.shape[1] != 2:
+            raise ValueError("expected [B, 2, H, W] containing index and valid mask")
+        values = inputs[:, :1] / 100.0
+        mask = (inputs[:, 1:2] > 0.5).to(values.dtype)
+        e1, m1 = self.enc1(values, mask)
+        e2, m2 = self.enc2(e1, m1)
+        e3, m3 = self.enc3(e2, m2)
+        x, m = self.bottleneck(e3, m3)
+        x, m = self._upsample(x, m, e3.shape[-2:])
+        x, m = self.dec3(torch.cat((x, e3), 1), torch.maximum(m, m3))
+        x, m = self._upsample(x, m, e2.shape[-2:])
+        x, m = self.dec2(torch.cat((x, e2), 1), torch.maximum(m, m2))
+        x, m = self._upsample(x, m, e1.shape[-2:])
+        x, _ = self.dec1(torch.cat((x, e1), 1), torch.maximum(m, m1))
+        return 100.0 * torch.sigmoid(self.output(x))
+
+
+if __name__ == "__main__":
+    model = CRAIClimateExtremes()
+    output = model(torch.rand(2, 2, 144, 192))
+    print(tuple(output.shape), float(output.min()), float(output.max()))

scripts/fake_data.py

@@ -0,0 +1,71 @@
+"""Create structurally realistic full-grid samples with irregular HadEX-style masks."""
+
+from pathlib import Path
+import json
+import numpy as np
+
+
+INDICES = np.array(["TX90p", "TN90p", "TX10p", "TN10p"])
+H, W = 144, 192
+
+
+def europe_mask(lat, lon):
+    yy, xx = np.meshgrid(lat, lon, indexing="ij")
+    broad = (yy >= 30) & (yy <= 72) & (xx >= -25) & (xx <= 45)
+    # A coarse geographic silhouette keeps the scientific global-grid contract.
+    atlantic_cut = (xx < -10) & (yy < 44)
+    southeast_cut = (xx > 30) & (yy < 40)
+    north_cut = (yy > 68) & ((xx < 5) | (xx > 30))
+    return (broad & ~atlantic_cut & ~southeast_cut & ~north_cut).astype(np.float32)
+
+
+def main():
+    rng = np.random.default_rng(42)
+    root = Path(__file__).resolve().parents[1]
+    output = root / "data"
+    output.mkdir(exist_ok=True)
+    lat = np.linspace(-89.375, 89.375, H, dtype=np.float32)
+    lon = np.linspace(-179.0625, 179.0625, W, dtype=np.float32)
+    land = europe_mask(lat, lon)
+    yy, xx = np.meshgrid(lat, lon, indexing="ij")
+    n = 8
+    target = np.zeros((n, 1, H, W), dtype=np.float32)
+    valid = np.zeros_like(target)
+    index_ids = np.arange(n, dtype=np.int64) % 4
+    for sample in range(n):
+        phase = 0.55 * sample
+        field = 50 + 21 * np.sin(np.deg2rad(2.3 * xx) + phase)
+        field += 16 * np.cos(np.deg2rad(3.2 * yy) - 0.4 * phase)
+        field += 5 * np.sin(np.deg2rad(xx + yy) * 4 + phase)
+        field += rng.normal(0, 1.2, (H, W))
+        if index_ids[sample] >= 2:
+            field = 100 - field
+        target[sample, 0] = np.clip(field, 0, 100) * land
+        observed = land.copy()
+        observed[rng.random((H, W)) < (0.35 + 0.04 * (sample % 3))] = 0
+        for _ in range(5):
+            cy, cx = rng.integers(45, 99), rng.integers(78, 121)
+            ry, rx = rng.integers(3, 11), rng.integers(4, 15)
+            hole = ((np.arange(H)[:, None] - cy) / ry) ** 2
+            hole = hole + ((np.arange(W)[None, :] - cx) / rx) ** 2
+            observed[hole < 1] = 0
+        valid[sample, 0] = observed
+    observed_values = target * valid
+    np.savez_compressed(
+        output / "crai_fake.npz", target=target, observed=observed_values,
+        valid_mask=valid, europe_mask=land, index_ids=index_ids,
+        index_names=INDICES, latitude=lat, longitude=lon,
+    )
+    metadata = {
+        "kind": "structured_synthetic",
+        "shape": [n, 1, H, W],
+        "grid_resolution": {"longitude_degrees": 1.875, "latitude_degrees": 1.25},
+        "indices": INDICES.tolist(),
+        "mask": "global grid with coarse Europe land support and irregular missing regions",
+    }
+    (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
+    print(f"wrote {output / 'crai_fake.npz'} with shape {target.shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,58 @@
+"""Run bounded ensemble reconstruction from a unified checkpoint."""
+
+from pathlib import Path
+import argparse
+import json
+import sys
+import numpy as np
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.crai_climateextremes import CRAIClimateExtremes
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    args = parser.parse_args()
+    config_path = args.config if args.config.is_absolute() else ROOT / args.config
+    with open(config_path, encoding="utf-8") as handle:
+        cfg = yaml.safe_load(handle)
+    data = np.load(ROOT / cfg["data_path"])
+    inputs = torch.from_numpy(np.concatenate((data["observed"], data["valid_mask"]), axis=1))
+    checkpoint_path = ROOT / cfg["checkpoint_path"]
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}; run scripts/train.py first")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=True)
+    if checkpoint.get("format_version") != "1.0" or not isinstance(checkpoint.get("model"), list):
+        raise ValueError(f"unsupported checkpoint format: {checkpoint_path}")
+    model_config = checkpoint.get("model_config", {})
+    predictions = []
+    for state in checkpoint["model"]:
+        model = CRAIClimateExtremes(**model_config).to(device)
+        model.load_state_dict(state); model.eval()
+        with torch.no_grad():
+            predictions.append(model(inputs.to(device)).cpu().numpy())
+    if not predictions:
+        raise ValueError(f"checkpoint contains no ensemble members: {checkpoint_path}")
+    members = np.stack(predictions)
+    output = ROOT / cfg["output_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        output / "predictions.npz", prediction=members.mean(0),
+        ensemble_std=members.std(0), target=data["target"],
+        observed=data["observed"], valid_mask=data["valid_mask"],
+        europe_mask=data["europe_mask"], index_ids=data["index_ids"],
+        index_names=data["index_names"],
+    )
+    metadata = {"ensemble_members": len(predictions), "checkpoint_semantics": "member state list in one checkpoint", "output_range": [0, 100]}
+    (output / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
+    print(f"predicted {inputs.shape[0]} samples with {len(predictions)} members")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,81 @@
+"""Evaluate missing-region reconstruction and create a comparison figure."""
+
+from pathlib import Path
+import argparse
+import json
+import numpy as np
+import torch
+import torch.nn.functional as F
+import yaml
+
+
+def correlation(a, b):
+    if a.size < 2 or np.std(a) == 0 or np.std(b) == 0:
+        return 0.0
+    return float(np.corrcoef(a, b)[0, 1])
+
+
+def rankdata(values):
+    order = np.argsort(values, kind="mergesort")
+    ranks = np.empty(len(values), dtype=np.float64)
+    ranks[order] = np.arange(len(values), dtype=np.float64)
+    unique, inverse, counts = np.unique(values, return_inverse=True, return_counts=True)
+    del unique
+    for group, count in enumerate(counts):
+        if count > 1:
+            positions = np.flatnonzero(inverse == group)
+            ranks[positions] = ranks[positions].mean()
+    return ranks
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    root = Path(__file__).resolve().parents[1]
+    parser.add_argument("--config", type=Path, default=root / "conf/config.yaml")
+    args = parser.parse_args()
+    config_path = args.config if args.config.is_absolute() else root / args.config
+    with open(config_path, encoding="utf-8") as handle:
+        cfg = yaml.safe_load(handle)
+    archive = np.load(root / cfg["output_dir"] / "predictions.npz")
+    pred, target = archive["prediction"], archive["target"]
+    missing = archive["europe_mask"][None, None] * (1 - archive["valid_mask"])
+    selected = missing.astype(bool)
+    error = pred[selected] - target[selected]
+    sample_spearman = []
+    for i in range(len(pred)):
+        mask = selected[i, 0]
+        sample_spearman.append(correlation(rankdata(pred[i, 0][mask]), rankdata(target[i, 0][mask])))
+    kernel = torch.ones(1, 1, 3, 3) / 8
+    kernel[0, 0, 1, 1] = 0
+    pred_neighbor = F.conv2d(torch.from_numpy(pred), kernel, padding=1).numpy()
+    target_neighbor = F.conv2d(torch.from_numpy(target), kernel, padding=1).numpy()
+    metrics = {
+        "missing_rmse": float(np.sqrt(np.mean(error ** 2))),
+        "sample_spearman_mean": float(np.mean(sample_spearman)),
+        "sample_spearman": [float(x) for x in sample_spearman],
+        "missing_bias": float(np.mean(error)),
+        "neighborhood_spatial_correlation_prediction": correlation(pred[selected], pred_neighbor[selected]),
+        "neighborhood_spatial_correlation_target": correlation(target[selected], target_neighbor[selected]),
+        "missing_points": int(selected.sum()),
+    }
+    output = root / cfg["evaluation_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    fig, axes = plt.subplots(1, 3, figsize=(13, 4), constrained_layout=True)
+    sample = 0
+    fields = [archive["observed"][sample, 0], target[sample, 0], pred[sample, 0]]
+    titles = ["Irregular observations", "Synthetic truth", "CRAI reconstruction"]
+    for axis, field, title in zip(axes, fields, titles):
+        image = axis.imshow(np.where(archive["europe_mask"] > 0, field, np.nan), origin="lower", vmin=0, vmax=100, cmap="RdYlBu_r")
+        axis.set_title(title); axis.set_axis_off()
+    fig.colorbar(image, ax=axes, label="Extreme index (%)", shrink=0.8)
+    fig.suptitle(f"Missing RMSE={metrics['missing_rmse']:.3f} | Spearman={metrics['sample_spearman_mean']:.3f} | Bias={metrics['missing_bias']:.3f}")
+    fig.savefig(output / "comparison.png", dpi=160); plt.close(fig)
+    print(json.dumps(metrics))
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,113 @@
+"""Train independent CRAI ensemble members, optionally under torchrun DDP."""
+
+from pathlib import Path
+import argparse
+import json
+import os
+import random
+import sys
+import numpy as np
+import torch
+from torch import distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.crai_climateextremes import CRAIClimateExtremes
+
+
+class ClimateDataset(Dataset):
+    def __init__(self, archive, limit):
+        self.observed = torch.from_numpy(archive["observed"][:limit])
+        self.valid = torch.from_numpy(archive["valid_mask"][:limit])
+        self.target = torch.from_numpy(archive["target"][:limit])
+        land = torch.from_numpy(archive["europe_mask"])[None, None]
+        self.missing = land * (1.0 - self.valid)
+
+    def __len__(self):
+        return len(self.target)
+
+    def __getitem__(self, index):
+        return torch.cat((self.observed[index], self.valid[index]), 0), self.target[index], self.missing[index]
+
+
+def load_config(path):
+    with open(path, encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    parser.add_argument("--paper-model", action="store_true")
+    args = parser.parse_args()
+    config_path = args.config if args.config.is_absolute() else ROOT / args.config
+    cfg = load_config(config_path)
+    use_paper = args.paper_model or cfg["paper_model"]
+    batch_size = cfg["paper_batch_size"] if use_paper else cfg["batch_size"]
+    iterations = cfg["paper_iterations"] if use_paper else cfg["max_iterations"]
+    members = cfg["paper_ensemble_members"] if use_paper else cfg["ensemble_members"]
+    rank, world = int(os.getenv("RANK", 0)), int(os.getenv("WORLD_SIZE", 1))
+    if world > 1:
+        dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', 0))}" if torch.cuda.is_available() else "cpu")
+    if device.type == "cuda":
+        torch.cuda.set_device(device)
+    archive = np.load(ROOT / cfg["data_path"])
+    dataset = ClimateDataset(archive, cfg["num_samples"])
+    checkpoint_path = ROOT / cfg["checkpoint_path"]
+    if rank == 0:
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+        (ROOT / "result/training").mkdir(parents=True, exist_ok=True)
+    if world > 1:
+        dist.barrier()
+    records, member_states = [], []
+    for member in range(members):
+        seed = int(cfg["seed"]) + member
+        random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
+        sampler = DistributedSampler(dataset, shuffle=True, seed=seed) if world > 1 else None
+        loader = DataLoader(dataset, batch_size=batch_size, shuffle=sampler is None, sampler=sampler)
+        model = CRAIClimateExtremes(cfg["base_channels"]).to(device)
+        if world > 1:
+            model = DistributedDataParallel(model, device_ids=[device.index] if device.type == "cuda" else None)
+        optimizer = torch.optim.Adam(model.parameters(), lr=cfg["learning_rate"])
+        step, losses = 0, []
+        for epoch in range(cfg["epochs"] if not use_paper else 10**9):
+            if sampler is not None:
+                sampler.set_epoch(epoch)
+            for inputs, target, missing in loader:
+                inputs, target, missing = inputs.to(device), target.to(device), missing.to(device)
+                prediction = model(inputs)
+                loss = (torch.abs(prediction - target) * missing).sum() / missing.sum().clamp_min(1)
+                optimizer.zero_grad(); loss.backward(); optimizer.step()
+                losses.append(float(loss.detach()))
+                step += 1
+                if step >= iterations:
+                    break
+            if step >= iterations:
+                break
+        raw_model = model.module if isinstance(model, DistributedDataParallel) else model
+        if rank == 0:
+            member_states.append({key: value.detach().cpu() for key, value in raw_model.state_dict().items()})
+            records.append({"member": member, "iterations": step, "final_missing_mae": losses[-1]})
+    if rank == 0:
+        torch.save(
+            {
+                "format_version": "1.0",
+                "model_config": {"base_channels": cfg["base_channels"]},
+                "model": member_states,
+            },
+            checkpoint_path,
+        )
+        payload = {"paper_model": bool(use_paper), "world_size": world, "members": records}
+        (ROOT / "result/training/metrics.json").write_text(json.dumps(payload, indent=2) + "\n")
+        print(json.dumps(payload))
+    if world > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()