Publish ScaleAdaptiveCM 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,130 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Probabilistic Downscaling
+- Consistency Model
+frameworks: PyTorch
+---
+
+<p align="center"><strong><span style="font-size: 30px;">ScaleAdaptiveCM</span></strong></p>
+
+# Model Introduction
+
+ScaleAdaptiveCM uses a single-step consistency model to transform coarse Earth-system-model precipitation into high-resolution probabilistic fields for fast, scale-adaptive climate downscaling and uncertainty analysis.
+
+Paper: Fast, scale-adaptive and uncertainty-aware downscaling of Earth system model fields with generative machine learning  
+https://doi.org/10.1038/s42256-025-00980-5
+
+# Model Description
+
+The method was proposed by teams from the Potsdam Institute for Climate Impact Research, the Technical University of Munich, Nanjing University of Information Science and Technology, and collaborating institutions. The paper trains on daily ERA5 precipitation and evaluates POEM, GFDL-ESM4, and SpeedyWeather.jl simulations. The model generates `240×384` precipitation ensembles from `60×96` coarse fields.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Probabilistic precipitation downscaling | Convert coarse ESM precipitation fields into high-resolution ensembles. |
+| Scale-adaptive generation | Control retained large scales and generated small-scale structure through noise scale. |
+| Local engineering validation | Validate consistency training, single-step sampling, ensemble mean, and spread. |
+| ModelScope/OneCode execution | Validate structured data, training, inference, downscaling metrics, and visualization in ModelScope or OneCode environments. |
+| Multi-GPU training | Validate distributed training and the checkpoint workflow through `torchrun`. |
+
+# Usage Instructions
+
+## 1.OneCode
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Download and Installation
+
+```bash
+hf download OneScience-Group/ScaleAdaptiveCM --local-dir ./ScaleAdaptiveCM
+cd ScaleAdaptiveCM
+```
+
+### 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 daily ERA5 precipitation from 1940-2018 on a `240×384` target grid and ESM fields on a `60×96` native grid. This repository generates structured precipitation with ITCZ, weather-system, and intermittent small-scale patterns while preserving the fourfold downscaling relation. Synthetic data validate engineering only and do not represent ERA5 or ESM 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 reduces samples, network width, and epochs without reducing the `60×96 → 240×384` protocol. Training artifacts are saved to `result/checkpoints/scale_adaptive_cm.pt` and `result/training/metrics.json`.
+
+### Trained Weights
+
+The paper does not provide directly loadable official model weights, and this repository bundles no weights under `weight/`.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference adds scale-guided noise to the coarse field and generates a high-resolution ensemble in one network evaluation per member. Results are saved to `result/output/predictions.npz`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation computes MAE, RMSE, large-scale correlation, power-spectrum error, and ensemble CRPS and writes `result/evaluation/metrics.json` and `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 ScaleAdaptiveCM 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,136 @@
+---
+license: apache-2.0
+language:
+- zh
+- en
+tags:
+- OneScience
+- 地球科学
+- 概率降尺度
+- 一致性模型
+frameworks: PyTorch
+---
+
+<p align="center"><strong><span style="font-size: 30px;">ScaleAdaptiveCM</span></strong></p>
+
+# 模型介绍
+
+ScaleAdaptiveCM 使用单步一致性生成模型将粗分辨率地球系统模式降水场转换为高分辨率概率降水场，用于快速、尺度自适应的气候降尺度和不确定性分析。
+
+论文：Fast, scale-adaptive and uncertainty-aware downscaling of Earth system model fields with generative machine learning  
+https://doi.org/10.1038/s42256-025-00980-5
+
+# 模型描述
+
+该方法由波茨坦气候影响研究所、慕尼黑工业大学和南京信息工程大学等机构的研究团队提出。论文使用 ERA5 逐日降水训练一致性模型，并使用 POEM、GFDL-ESM4 和 SpeedyWeather.jl 等模式数据评估。模型适用于从 `60×96` 粗网格生成 `240×384` 高分辨率降水集合。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 概率降水降尺度 | 将粗分辨率 ESM 降水场转换为高分辨率集合。 |
+| 尺度自适应生成 | 通过噪声尺度控制保留的大尺度模式和生成的小尺度结构。 |
+| 本地工程验证 | 验证一致性训练、单步采样、集合均值和离散度。 |
+| ModelScope/OneCode 运行 | 在 ModelScope 或 OneCode 环境中验证结构化数据、训练、推理、降尺度指标和可视化流程。 |
+| 多卡训练 | 通过 `torchrun` 验证分布式训练和 checkpoint 流程。 |
+
+# 使用说明
+
+## 1.OneCode
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/ScaleAdaptiveCM --local_dir ./ScaleAdaptiveCM
+cd ScaleAdaptiveCM
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 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
+```
+
+### 训练数据介绍
+
+论文使用 1940-2018 年 ERA5 逐日降水，目标网格为 `240×384`，ESM 原生网格为 `60×96`。本仓库生成具有 ITCZ、天气系统和小尺度间歇性的结构化虚拟降水场，并保持真实四倍降尺度关系。虚拟数据仅用于验证工程流程，不代表 ERA5 或 ESM 的真实分布、训练规模或论文性能。
+
+```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
+```
+
+默认配置缩小样本数、网络宽度和训练轮数，但不缩小 `60×96 → 240×384` 空间协议。正式实验需要真实 ERA5 与 ESM 数据和论文规模模型，训练产物保存到：
+
+```text
+result/checkpoints/scale_adaptive_cm.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+论文未提供可直接加载的官方模型权重，本仓库不在 `weight/` 中内置权重。本地 checkpoint 不得描述为官方预训练权重。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理对粗分辨率降水场施加尺度引导噪声，并在单次网络调用中生成高分辨率集合。结果保存到 `result/output/predictions.npz`。
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估计算 MAE、RMSE、大尺度相关性、功率谱误差和 ensemble CRPS，并生成输入、目标、集合均值与误差图。结果保存到 `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 |
+
+# 引用与许可证
+
+本仓库为 ScaleAdaptiveCM 公开规格的独立工程复现版本。
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,40 @@
+seed: 42
+data:
+  path: data/scale_adaptive_cm_fake.npz
+  format_version: scale_adaptive_cm_v1
+  samples: 4
+  low_grid: [60, 96]
+  high_grid: [240, 384]
+  scale_factor: 4
+  log_epsilon: 0.0001
+model:
+  channels: [4, 8, 16]
+  time_dim: 16
+  sigma_data: 0.5
+train:
+  epochs: 1
+  batch_size: 1
+  learning_rate: 0.0002
+  ema_decay: 0.9
+runtime:
+  device: cpu
+paths:
+  checkpoint: result/checkpoints/scale_adaptive_cm.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png
+evaluation:
+  ensemble_members: 3
+  guidance_sigma: 0.468
+paper_model:
+  high_grid: [240, 384]
+  low_grid: [60, 96]
+  channels: [128, 128, 256, 256]
+  epochs: 150
+  batch_size: 1
+  learning_rate: 0.0002
+  t_min: 0.002
+  t_max: 80.0
+  rho: 7
+  ema_initial_decay: 0.9

config.json

@@ -0,0 +1,19 @@
+{
+  "model_name": "ScaleAdaptiveCM",
+  "model_type": "scale_adaptive_cm",
+  "architectures": ["ScaleAdaptiveCM"],
+  "framework": "PyTorch",
+  "domain": "climate",
+  "task": "probabilistic-precipitation-downscaling",
+  "implementation": {
+    "entry_point": "model/scale_adaptive_cm.py",
+    "scope": "core-method, full-spatial-dimension reduced-model 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_channels": 1, "output_channels": 1, "engineering_channels": [4, 8, 16], "paper_channels": [128, 128, 256, 256], "scale_factor": 4},
+  "data": {"dataset": "ERA5 and ESM precipitation", "format_version": "scale_adaptive_cm_v1", "low_shape": ["B", 1, 60, 96], "high_shape": ["B", 1, 240, 384], "unit": "mm day-1", "synthetic": true},
+  "configuration_sources": ["conf/config.yaml", "model/scale_adaptive_cm.py", "scripts/fake_data.py", "scripts/train.py", "scripts/inference.py", "scripts/result.py"]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "probabilistic precipitation downscaling",
+  "model": "ScaleAdaptiveCM",
+  "input_format": "BCHW: low-resolution [B,1,60,96] and target [B,1,240,384]",
+  "protocol": "unconditional consistency training and scale-guided single-step downscaling",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/scale_adaptive_cm.py

@@ -0,0 +1,85 @@
+"""Compact consistency model for full-grid precipitation downscaling."""
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+import yaml
+
+
+def load_config(root):
+    return yaml.safe_load((Path(root) / "conf/config.yaml").read_text())
+
+
+class TimeBlock(nn.Module):
+    def __init__(self, cin, cout, time_dim):
+        super().__init__()
+        groups = min(4, cout)
+        self.conv = nn.Conv2d(cin, cout, 3, padding=1)
+        self.norm = nn.GroupNorm(groups, cout)
+        self.time = nn.Linear(time_dim, cout)
+
+    def forward(self, x, embedding):
+        return F.silu(self.norm(self.conv(x)) + self.time(embedding)[:, :, None, None])
+
+
+class ScaleAdaptiveCM(nn.Module):
+    def __init__(self, channels=(4, 8, 16), time_dim=16, sigma_data=0.5):
+        super().__init__()
+        self.sigma_data = float(sigma_data)
+        self.time_dim = int(time_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(time_dim, time_dim), nn.SiLU(), nn.Linear(time_dim, time_dim))
+        self.enc1 = TimeBlock(1, channels[0], time_dim)
+        self.enc2 = TimeBlock(channels[0], channels[1], time_dim)
+        self.mid = TimeBlock(channels[1], channels[2], time_dim)
+        self.dec2 = TimeBlock(channels[2] + channels[1], channels[1], time_dim)
+        self.dec1 = TimeBlock(channels[1] + channels[0], channels[0], time_dim)
+        self.out = nn.Conv2d(channels[0], 1, 1)
+        self.model_config = {"channels": list(channels), "time_dim": time_dim, "sigma_data": sigma_data}
+
+    def embed_time(self, t):
+        half = self.time_dim // 2
+        freq = torch.exp(torch.linspace(0, -7, half, device=t.device))
+        emb = torch.cat((torch.sin(t[:, None] * freq), torch.cos(t[:, None] * freq)), 1)
+        return self.time_mlp(emb)
+
+    def forward(self, noisy, t):
+        if noisy.ndim != 4 or noisy.shape[1] != 1:
+            raise ValueError("expected [B,1,H,W]")
+        emb = self.embed_time(t.float())
+        e1 = self.enc1(noisy, emb)
+        e2 = self.enc2(F.avg_pool2d(e1, 2), emb)
+        mid = self.mid(F.avg_pool2d(e2, 2), emb)
+        d2 = self.dec2(torch.cat((F.interpolate(mid, e2.shape[-2:], mode="bilinear", align_corners=False), e2), 1), emb)
+        d1 = self.dec1(torch.cat((F.interpolate(d2, e1.shape[-2:], mode="bilinear", align_corners=False), e1), 1), emb)
+        raw = self.out(d1)
+        sigma2 = self.sigma_data ** 2
+        cskip = sigma2 / ((t[:, None, None, None] - 0.002).square() + sigma2)
+        cout = self.sigma_data * t[:, None, None, None] / torch.sqrt(t[:, None, None, None].square() + sigma2)
+        return cskip * noisy + cout * raw
+
+
+def structured_fields(samples, high_h, high_w, seed):
+    rng = np.random.default_rng(seed)
+    yy, xx = np.mgrid[-1:1:complex(high_h), -1:1:complex(high_w)]
+    fields = []
+    for i in range(samples):
+        phase = 2 * np.pi * i / max(samples, 4)
+        itcz = 9 * np.exp(-((yy - .12 * np.sin(phase)) / .16) ** 2)
+        storms = 18 * np.exp(-((xx - .45 * np.cos(phase)) ** 2 + (yy - .3 * np.sin(phase)) ** 2) / .025)
+        texture = 2 * np.maximum(0, np.sin(18 * xx + phase) * np.cos(13 * yy - phase))
+        fields.append(np.maximum(0, itcz + storms + texture + rng.normal(0, .15, yy.shape)))
+    return np.asarray(fields, np.float32)[:, None]
+
+
+def radial_spectrum(field):
+    power = np.abs(np.fft.fftshift(np.fft.fft2(field))) ** 2
+    y, x = np.indices(field.shape); r = np.sqrt((y-field.shape[0]/2)**2 + (x-field.shape[1]/2)**2).astype(int)
+    return np.bincount(r.ravel(), power.ravel()) / np.maximum(np.bincount(r.ravel()), 1)
+
+
+def write_json(path, value):
+    path = Path(path); path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(value, indent=2) + "\n")

scripts/fake_data.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+import sys
+import numpy as np
+import torch.nn.functional as F
+import torch
+ROOT = Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT))
+from model.scale_adaptive_cm import load_config, structured_fields
+
+c = load_config(ROOT); h,w=c["data"]["high_grid"]; lh,lw=c["data"]["low_grid"]
+high=structured_fields(c["data"]["samples"],h,w,c["seed"])
+low=F.avg_pool2d(torch.from_numpy(high),c["data"]["scale_factor"]).numpy()
+assert low.shape[-2:]==(lh,lw)
+path=ROOT/c["data"]["path"]; path.parent.mkdir(parents=True,exist_ok=True)
+np.savez_compressed(path,format_version=np.array(c["data"]["format_version"]),low=low,high=high,split=np.array(["train"]*(len(high)-1)+["test"]),unit=np.array("mm day-1"))
+print(path,low.shape,high.shape)

scripts/inference.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+import sys
+import numpy as np,torch
+import torch.nn.functional as F
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.scale_adaptive_cm import ScaleAdaptiveCM,load_config
+c=load_config(ROOT);d=np.load(ROOT/c["data"]["path"]);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);assert ck["format_version"]==str(d["format_version"])
+m=ScaleAdaptiveCM(**ck["model_config"]);m.load_state_dict(ck["model"]);m.eval(); low=torch.from_numpy(d["low"][d["split"]=="test"]).float(); up=F.interpolate(low,size=c["data"]["high_grid"],mode="bilinear",align_corners=False)
+base=(torch.log(up+c["data"]["log_epsilon"])-np.log(c["data"]["log_epsilon"])-ck["normalization"]["mean"])/ck["normalization"]["std"];members=[];t=torch.full((len(base),),c["evaluation"]["guidance_sigma"])
+with torch.no_grad():
+ for i in range(c["evaluation"]["ensemble_members"]): members.append(m(base+t[:,None,None,None]*torch.randn_like(base),t))
+z=torch.stack(members)*ck["normalization"]["std"]+ck["normalization"]["mean"];pred=torch.exp(z+np.log(c["data"]["log_epsilon"]))-c["data"]["log_epsilon"];pred=pred.clamp_min(0).numpy()
+path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,low=low.numpy(),target=d["high"][d["split"]=="test"],members=pred,mean=pred.mean(0),std=pred.std(0),unit=d["unit"]);print(path)

scripts/result.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+import sys
+import numpy as np
+import matplotlib;matplotlib.use("Agg");import matplotlib.pyplot as plt
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.scale_adaptive_cm import load_config,radial_spectrum,write_json
+c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);p,t=d["mean"],d["target"];err=p-t
+mae=float(np.mean(abs(err)));rmse=float(np.sqrt(np.mean(err**2)));corr=float(np.corrcoef(p.ravel(),t.ravel())[0,1]);psp,psT=radial_spectrum(p[0,0]),radial_spectrum(t[0,0]);n=min(len(psp),len(psT));psd=float(np.mean(abs(np.log1p(psp[:n])-np.log1p(psT[:n]))))
+members=d["members"][:,0,0];obs=t[0,0];crps=float(np.mean(abs(members-obs))-0.5*np.mean(abs(members[:,None]-members[None,:])))
+write_json(ROOT/c["paths"]["evaluation"],{"mae_mm_day":mae,"rmse_mm_day":rmse,"large_scale_correlation":corr,"log_psd_mae":psd,"ensemble_crps":crps,"synthetic":True})
+fig,ax=plt.subplots(1,4,figsize=(14,3.4));fields=(d["low"][0,0],t[0,0],p[0,0],abs(err[0,0]));titles=("Low resolution","Target","CM mean","Absolute error")
+for a,f,title in zip(ax,fields,titles):im=a.imshow(f,cmap="viridis");a.set_title(title);a.axis("off");fig.colorbar(im,ax=a,shrink=.7)
+fig.tight_layout();path=ROOT/c["paths"]["figure"];path.parent.mkdir(parents=True,exist_ok=True);fig.savefig(path,dpi=140);plt.close(fig);print(path)

scripts/train.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+import sys,copy,os
+import numpy as np, torch
+from torch.nn.parallel import DistributedDataParallel
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.scale_adaptive_cm import ScaleAdaptiveCM,load_config,write_json
+c=load_config(ROOT);rank=int(os.environ.get("RANK",0));world=int(os.environ.get("WORLD_SIZE",1));local=int(os.environ.get("LOCAL_RANK",0));distributed=world>1;use_cuda=torch.cuda.is_available() and torch.cuda.device_count()>=world and c["runtime"]["device"]!="cpu"
+if distributed: torch.distributed.init_process_group(backend="nccl" if use_cuda else "gloo")
+device=torch.device(f"cuda:{local}" if use_cuda else "cpu");torch.manual_seed(c["seed"]);torch.set_num_threads(2)
+d=np.load(ROOT/c["data"]["path"]); x=torch.from_numpy(np.log(d["high"][d["split"]=="train"]+c["data"]["log_epsilon"])-np.log(c["data"]["log_epsilon"])).float()
+mean,std=x.mean(),x.std().clamp_min(1e-6);x=(x-mean)/std
+m=ScaleAdaptiveCM(**c["model"]).to(device);target=copy.deepcopy(m).eval();wrapped=DistributedDataParallel(m,device_ids=[local] if use_cuda else None) if distributed else m;opt=torch.optim.RAdam(wrapped.parameters(),lr=c["train"]["learning_rate"]);hist=[];x=x.to(device)
+for e in range(c["train"]["epochs"]):
+    noise=torch.randn_like(x);t1=torch.full((len(x),),.35);t2=torch.full((len(x),),.55)
+    with torch.no_grad(): y=target(x+t1[:,None,None,None]*noise,t1)
+    pred=wrapped(x+t2[:,None,None,None]*noise,t2);loss=(pred-y).abs().mean();opt.zero_grad();loss.backward();opt.step()
+    with torch.no_grad():
+        for p,q in zip(target.parameters(),m.parameters()):p.mul_(c["train"]["ema_decay"]).add_(q,alpha=1-c["train"]["ema_decay"])
+    hist.append({"epoch":e+1,"loss":float(loss)})
+summary=torch.tensor([sum(v["loss"] for v in hist),len(hist)],dtype=torch.float64,device=device)
+if distributed: torch.distributed.all_reduce(summary)
+if rank==0:
+ path=ROOT/c["paths"]["checkpoint"];path.parent.mkdir(parents=True,exist_ok=True);torch.save({"model":target.state_dict(),"model_config":c["model"],"format_version":c["data"]["format_version"],"normalization":{"mean":float(mean),"std":float(std)}},path)
+ write_json(ROOT/c["paths"]["training_metrics"],{"history":hist,"global_mean_loss":float(summary[0]/summary[1]),"world_size":world});print(path)
+if distributed: torch.distributed.destroy_process_group()