Publish PrecipExtremes-GAN 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
+- Extreme Precipitation
+- Generative Downscaling
+frameworks: PyTorch
+---
+
+<p align="center"><strong><span style="font-size: 30px;">PrecipExtremes-GAN</span></strong></p>
+
+# Model Introduction
+
+PrecipExtremes-GAN uses a residual generative adversarial network to map coarse atmospheric fields to high-resolution daily precipitation and assess extrapolation of extreme changes to warmer climates.
+
+Paper: On the Extrapolation of Generative Adversarial Networks for downscaling precipitation extremes in warmer climates  
+https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2024GL112492
+
+# Model Description
+
+The method was proposed by teams from New Zealand's National Institute of Water and Atmospheric Research, the University of New South Wales, and collaborating institutions. The paper trains on ACCESS-CM2-driven CCAM simulations and evaluates four independent GCM-driven historical and SSP3-7.0 simulations. The model supports daily precipitation downscaling and extreme-precipitation climate-change analysis over New Zealand.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Daily precipitation downscaling | Generate high-resolution precipitation from eight coarse atmospheric predictors. |
+| Extreme extrapolation | Compare historical and future training for 99.5th-percentile precipitation change. |
+| Local engineering validation | Validate deterministic U-Net, residual GAN, and ensemble inference. |
+| ModelScope/OneCode execution | Validate structured data, training, inference, extreme-precipitation metrics, and visualization in ModelScope or OneCode environments. |
+| Multi-GPU training | Validate distributed training and checkpoint workflows 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/PrecipExtremes-GAN --local-dir ./PrecipExtremes-GAN
+cd PrecipExtremes-GAN
+```
+
+### 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 predicts approximately 12 km CCAM daily precipitation from eight 1.5-degree U, V, T, and Q fields at 500 and 850 hPa. Since final tensor dimensions are not reported, this repository uses an explicitly documented `24×24 → 96×96` engineering grid. Synthetic data validate engineering only and do not represent CCAM or GCM distributions, 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
+```
+
+Training first fits a deterministic U-Net and then trains a residual generator with MSE, adversarial, and intensity losses. The default reduces samples, width, ensemble size, and epochs, writing `result/checkpoints/precip_extremes_gan.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 writes deterministic, residual-GAN member, and ensemble-mean outputs to `result/output/predictions.npz`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation computes daily MAE, 99.5th-percentile error, and future-versus-historical climate-signal error 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 PrecipExtremes-GAN 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,131 @@
+---
+license: apache-2.0
+language:
+- zh
+- en
+tags:
+- OneScience
+- 地球科学
+- 极端降水
+- 生成式降尺度
+frameworks: PyTorch
+---
+
+<p align="center"><strong><span style="font-size: 30px;">PrecipExtremes-GAN</span></strong></p>
+
+# 模型介绍
+
+PrecipExtremes-GAN 使用残差生成对抗网络将粗分辨率大气场转换为高分辨率日降水，并评估模型向更暖气候外推极端降水变化的能力。
+
+论文：On the Extrapolation of Generative Adversarial Networks for downscaling precipitation extremes in warmer climates  
+https://agupubs.onlinelibrary.wiley.com/doi/10.1029/2024GL112492
+
+# 模型描述
+
+该方法由新西兰国家水与大气研究所和新南威尔士大学等机构的研究团队提出。论文使用 ACCESS-CM2 驱动的 CCAM 模拟训练，并在四个独立 GCM 驱动的历史和 SSP3-7.0 模拟上评估。模型适用于新西兰区域的日降水降尺度及极端降水气候变化信号分析。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 日降水降尺度 | 根据 8 个粗分辨率大气变量生成高分辨率日降水。 |
+| 极端降水外推 | 比较历史和未来训练对 99.5 百分位降水变化的影响。 |
+| 本地工程验证 | 验证 deterministic U-Net、残差 GAN 和集合推理。 |
+| ModelScope/OneCode 运行 | 在 ModelScope 或 OneCode 环境中验证结构化数据、训练、推理、极端降水指标和可视化流程。 |
+| 多卡训练 | 通过 `torchrun` 验证分布式训练和 checkpoint 流程。 |
+
+# 使用说明
+
+## 1.OneCode
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/PrecipExtremes-GAN --local_dir ./PrecipExtremes-GAN
+cd PrecipExtremes-GAN
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 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
+```
+
+### 训练数据介绍
+
+论文使用 1.5° 粗分辨率的 U、V、T 和 Q 在 500/850 hPa 的 8 个通道预测约 12 km CCAM 日降水。公开正文未给最终张量网格点数，本仓库使用 `24×24 → 96×96` 工程网格并明确标记为假设。虚拟数据仅用于验证工程流程，不代表 CCAM 或 GCM 的真实分布、训练规模或论文性能。
+
+```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
+```
+
+训练先拟合 deterministic U-Net，再使用 MSE、对抗损失和强度约束训练残差生成器。默认缩小样本、网络宽度、ensemble 和训练轮数，训练产物保存到 `result/checkpoints/precip_extremes_gan.pt` 和 `result/training/metrics.json`。
+
+### 训练权重
+
+论文未提供可直接加载的官方模型权重，本仓库不在 `weight/` 中内置权重。本地 checkpoint 不得描述为官方预训练权重。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理输出 deterministic baseline、残差 GAN 成员及集合均值，保存到 `result/output/predictions.npz`。
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估计算日降水 MAE、99.5 百分位误差和未来相对历史气候变化信号误差，并生成对比图。结果保存到 `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 |
+
+# 引用与许可证
+
+本仓库为 PrecipExtremes-GAN 公开规格的独立工程复现版本。
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,43 @@
+seed: 42
+data:
+  path: data/precip_extremes_gan_fake.npz
+  format_version: precip_extremes_gan_v1
+  samples_per_period: 16
+  low_grid: [24, 24]
+  high_grid: [96, 96]
+  input_channels: 8
+  log_epsilon: 0.001
+model:
+  filters: [8, 16, 32]
+  noise_channels: 1
+train:
+  baseline_epochs: 1
+  gan_epochs: 1
+  batch_size: 2
+  baseline_lr: 0.0007
+  gan_lr: 0.0002
+  adversarial_weight: 0.01
+  intensity_weight: 1.0
+runtime:
+  device: cpu
+paths:
+  checkpoint: result/checkpoints/precip_extremes_gan.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
+  extreme_percentile: 99.5
+  wet_day_threshold: 1.0
+paper_model:
+  filters: [32, 64, 128]
+  parameters_baseline: 1500000
+  epochs: 250
+  batch_size: 32
+  adversarial_weight: 0.01
+  ensemble_members: 10
+  historical_period: 1960-2014
+  future_period: 2044-2099
+  evaluation_historical: 1985-2014
+  evaluation_future: 2070-2099

config.json

@@ -0,0 +1,12 @@
+{
+  "model_name": "PrecipExtremes-GAN",
+  "model_type": "precip_extremes_gan",
+  "architectures": ["PrecipExtremesGAN"],
+  "framework": "PyTorch",
+  "domain": "climate",
+  "task": "extreme-precipitation-downscaling",
+  "implementation": {"entry_point": "model/precip_extremes_gan.py", "scope": "core-method 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_channels": 8, "output_channels": 1, "baseline": "residual U-Net", "generator": "stochastic residual U-Net", "critic": "conditional patch critic", "paper_filters": [32,64,128]},
+  "data": {"dataset": "CCAM New Zealand perfect-framework simulations", "format_version": "precip_extremes_gan_v1", "input_variables": ["U500","U850","V500","V850","T500","T850","Q500","Q850"], "target": "daily precipitation", "target_unit": "mm day-1", "engineering_low_grid": [24,24], "engineering_high_grid": [96,96], "grid_status": "engineering assumption because final tensor dimensions are not reported", "synthetic": true},
+  "configuration_sources": ["conf/config.yaml", "model/precip_extremes_gan.py", "scripts/fake_data.py", "scripts/train.py", "scripts/inference.py", "scripts/result.py"]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "extreme precipitation downscaling",
+  "model": "PrecipExtremes-GAN",
+  "input_format": "BCHW: eight coarse atmospheric predictor fields",
+  "protocol": "deterministic U-Net plus residual conditional GAN for historical-to-future extrapolation",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/precip_extremes_gan.py

@@ -0,0 +1,47 @@
+"""Residual U-Net and conditional GAN for precipitation extremes."""
+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 ResidualBlock(nn.Module):
+    def __init__(self, cin, cout):
+        super().__init__(); self.c1=nn.Conv2d(cin,cout,3,padding=1);self.c2=nn.Conv2d(cout,cout,3,padding=1);self.skip=nn.Conv2d(cin,cout,1) if cin!=cout else nn.Identity()
+    def forward(self,x): return F.relu(self.c2(F.relu(self.c1(x)))+self.skip(x))
+
+class ResidualUNet(nn.Module):
+    def __init__(self,in_channels=8,filters=(8,16,32),out_channels=1):
+        super().__init__();a,b,c=filters;self.e1=ResidualBlock(in_channels,a);self.e2=ResidualBlock(a,b);self.mid=ResidualBlock(b,c);self.d2=ResidualBlock(c+b,b);self.d1=ResidualBlock(b+a,a);self.out=nn.Conv2d(a,out_channels,1)
+    def forward(self,x,out_size=(96,96)):
+        e1=self.e1(x);e2=self.e2(F.avg_pool2d(e1,2));m=self.mid(F.avg_pool2d(e2,2));d2=self.d2(torch.cat((F.interpolate(m,e2.shape[-2:],mode="bilinear",align_corners=False),e2),1));d1=self.d1(torch.cat((F.interpolate(d2,e1.shape[-2:],mode="bilinear",align_corners=False),e1),1));return F.interpolate(self.out(d1),size=out_size,mode="bilinear",align_corners=False)
+
+class PrecipExtremesGAN(nn.Module):
+    def __init__(self,input_channels=8,filters=(8,16,32),noise_channels=1):
+        super().__init__();self.baseline=ResidualUNet(input_channels,filters);self.generator=ResidualUNet(input_channels+noise_channels,filters);self.model_config={"input_channels":input_channels,"filters":list(filters),"noise_channels":noise_channels}
+    def forward(self,x,noise=None,out_size=(96,96)):
+        base=self.baseline(x,out_size)
+        if noise is None: noise=torch.randn(x.shape[0],1,*x.shape[-2:],device=x.device)
+        return base,base+self.generator(torch.cat((x,noise),1),out_size)
+
+class PatchCritic(nn.Module):
+    def __init__(self,input_channels=8):
+        super().__init__();self.net=nn.Sequential(nn.Conv2d(input_channels+1,16,4,2,1),nn.LeakyReLU(.2),nn.Conv2d(16,32,4,2,1),nn.LeakyReLU(.2),nn.Conv2d(32,1,3,1,1))
+    def forward(self,x,y): return self.net(torch.cat((F.interpolate(x,y.shape[-2:],mode="bilinear",align_corners=False),y),1))
+
+def make_data(n,h,w,lh,lw,seed):
+    rng=np.random.default_rng(seed); yy,xx=np.mgrid[-1:1:complex(h),-1:1:complex(w)]; periods=[]
+    for future in (0,1):
+      for i in range(n):
+        phase=2*np.pi*i/n; coarse=np.stack([np.sin((k+1)*xx+phase)+np.cos((k+2)*yy-phase)+future*.2 for k in range(8)]).astype(np.float32)
+        low=torch.from_numpy(coarse[None]);low=F.interpolate(low,(lh,lw),mode="area")[0].numpy(); ups=F.interpolate(torch.from_numpy(low[None]),(h,w),mode="bilinear",align_corners=False)[0]
+        moisture=torch.sigmoid(ups[7:8]*1.2+ups[6:7]*.7);lift=torch.relu(ups[2:3]+.4*ups[0:1]);extreme=(1+future*.35)*8*np.exp(-((xx-.35*np.sin(phase))**2+(yy-.25*np.cos(phase))**2)/.018)
+        pr=np.maximum(0,(5*moisture*lift).numpy()[0]+extreme+np.maximum(0,rng.normal(0,.4,(h,w)))).astype(np.float32)
+        periods.append((low,pr[None],future))
+    return periods
+
+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,5 @@
+from pathlib import Path
+import sys,numpy as np
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.precip_extremes_gan import load_config,make_data
+c=load_config(ROOT);d=c["data"];records=make_data(d["samples_per_period"],*d["high_grid"],*d["low_grid"],c["seed"]);x=np.stack([r[0] for r in records]);y=np.stack([r[1] for r in records]);period=np.array([r[2] for r in records]);split=np.array(["train" if i%d["samples_per_period"]<d["samples_per_period"]-4 else "test" for i in range(len(records))]);path=ROOT/d["path"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,format_version=np.array(d["format_version"]),predictors=x,precipitation=y,period=period,split=split,unit=np.array("mm day-1"));print(path,x.shape,y.shape)

scripts/inference.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+import sys,numpy as np,torch
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.precip_extremes_gan import *
+c=load_config(ROOT);d=np.load(ROOT/c["data"]["path"]);ck=torch.load(ROOT/c["paths"]["checkpoint"],map_location="cpu",weights_only=True);net=PrecipExtremesGAN(**ck["model_config"]);net.load_state_dict(ck["model"]);net.eval();mask=d["split"]=="test";x=torch.from_numpy(d["predictors"][mask]);members=[]
+with torch.no_grad():
+ base=net.baseline(x,c["data"]["high_grid"])
+ for _ in range(c["evaluation"]["ensemble_members"]):members.append(base+net.generator(torch.cat((x,torch.randn(len(x),1,*x.shape[-2:])),1),c["data"]["high_grid"]))
+z=torch.stack(members);p=(torch.exp(z)-c["data"]["log_epsilon"]).clamp_min(0).numpy();b=(torch.exp(base)-c["data"]["log_epsilon"]).clamp_min(0).numpy();path=ROOT/c["paths"]["predictions"];path.parent.mkdir(parents=True,exist_ok=True);np.savez_compressed(path,predictors=x.numpy(),target=d["precipitation"][mask],baseline=b,members=p,mean=p.mean(0),period=d["period"][mask],unit=d["unit"]);print(path)

scripts/result.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+import sys,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.precip_extremes_gan import load_config,write_json
+c=load_config(ROOT);d=np.load(ROOT/c["paths"]["predictions"]);q=c["evaluation"]["extreme_percentile"];metrics={}
+for name,p in (("deterministic",d["baseline"]),("gan",d["mean"])):
+ metrics[name]={"mae":float(np.mean(abs(p-d["target"]))),"extreme_percentile_mae":float(np.mean(abs(np.percentile(p,q,axis=0)-np.percentile(d["target"],q,axis=0))))}
+def signal(a):
+ h=a[d["period"]==0].mean(0);f=a[d["period"]==1].mean(0);return 100*(f-h)/np.maximum(h,.1)
+truth=signal(d["target"]);metrics["climate_signal_mae_percent"]={"deterministic":float(np.mean(abs(signal(d["baseline"])-truth))),"gan":float(np.mean(abs(signal(d["mean"])-truth)))};metrics["synthetic"]=True;write_json(ROOT/c["paths"]["evaluation"],metrics)
+fig,ax=plt.subplots(1,4,figsize=(14,3.5));fields=(d["target"][0,0],d["baseline"][0,0],d["mean"][0,0],abs(d["mean"][0,0]-d["target"][0,0]));titles=("CCAM target","Deterministic","Residual GAN","Absolute error")
+for a,f,t in zip(ax,fields,titles):im=a.imshow(f,cmap="Blues");a.set_title(t);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,22 @@
+from pathlib import Path
+import sys,os
+import numpy as np,torch
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel
+ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
+from model.precip_extremes_gan import *
+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"]+rank);torch.set_num_threads(2);d=np.load(ROOT/c["data"]["path"]);m=d["split"]=="train";x=torch.from_numpy(d["predictors"][m]).to(device);y=torch.log(torch.from_numpy(d["precipitation"][m])+c["data"]["log_epsilon"]).to(device);net=PrecipExtremesGAN(c["data"]["input_channels"],c["model"]["filters"],c["model"]["noise_channels"]).to(device);critic=PatchCritic(c["data"]["input_channels"]).to(device)
+baseline=DistributedDataParallel(net.baseline,device_ids=[local] if use_cuda else None) if distributed else net.baseline;generator=DistributedDataParallel(net.generator,device_ids=[local] if use_cuda else None) if distributed else net.generator;critic_train=DistributedDataParallel(critic,device_ids=[local] if use_cuda else None) if distributed else critic
+ob=torch.optim.Adam(baseline.parameters(),lr=c["train"]["baseline_lr"])
+for _ in range(c["train"]["baseline_epochs"]): pred=baseline(x,c["data"]["high_grid"]);loss=F.mse_loss(pred,y);ob.zero_grad();loss.backward();ob.step()
+og=torch.optim.Adam(generator.parameters(),lr=c["train"]["gan_lr"]);oc=torch.optim.Adam(critic_train.parameters(),lr=c["train"]["gan_lr"])
+for _ in range(c["train"]["gan_epochs"]):
+ with torch.no_grad():base=baseline(x,c["data"]["high_grid"])
+ fake=base+generator(torch.cat((x,torch.randn(len(x),1,*x.shape[-2:],device=device)),1),c["data"]["high_grid"]);cl=F.softplus(critic_train(x,fake.detach())).mean()+F.softplus(-critic_train(x,y)).mean();oc.zero_grad();cl.backward();oc.step();adv=F.softplus(-critic_train(x,fake)).mean();intensity=F.mse_loss(fake.mean((2,3)),y.mean((2,3)));gl=F.mse_loss(fake.mean(0),y.mean(0))+c["train"]["adversarial_weight"]*adv+intensity;og.zero_grad();gl.backward();og.step()
+vals=torch.tensor([float(loss),float(gl),float(cl)],device=device); 
+if distributed:torch.distributed.all_reduce(vals);vals/=world
+if rank==0:
+ path=ROOT/c["paths"]["checkpoint"];path.parent.mkdir(parents=True,exist_ok=True);torch.save({"model":net.state_dict(),"critic":critic.state_dict(),"model_config":net.model_config,"format_version":c["data"]["format_version"]},path);write_json(ROOT/c["paths"]["training_metrics"],{"baseline_mse":float(vals[0]),"generator_loss":float(vals[1]),"critic_loss":float(vals[2]),"world_size":world});print(path)
+if distributed:torch.distributed.destroy_process_group()