Upload folder using huggingface_hub

.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,147 @@
+---
+license: cc-by-nc-nd-4.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Rainfall Correction
+- UnetDif
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">UnetDif</span></strong>
+</p>
+
+# Model Introduction
+
+UnetDif uses the paper's bias-targeted multitask U-Net to correct short-range heavy rainfall forecasts.
+
+Paper: Bias-targeted deep learning enhances short-range heavy rainfall forecasts  
+https://doi.org/10.1038/s41612-026-01366-z
+
+# Model Description
+
+The method reproduced by UnetDif was proposed by research teams from Zhejiang Meteorological Observatory, the National Meteorological Center, and Zhejiang University. The paper constructs Yangtze River Delta samples from European Centre for Medium-Range Weather Forecasts (ECMWF) fields, China Meteorological Administration Multi-source Precipitation Analysis (CMPA), and terrain data. The model learns `ECMWF - CMPA` rainfall bias and corrects three-hour precipitation through a shared U-Net backbone with four heads for dry areas, false alarms, positive bias, and negative bias.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Three-hour rainfall-bias correction | Generate corrected ECMWF rainfall from 39-channel inputs on the fixed Yangtze River Delta grid. |
+| ModelScope/OneCode execution | Validate synthetic-data generation, six-loss training, inference, evaluation, and visualization in ModelScope or OneCode environments. |
+| Hugging Face/OneCode execution | Run the complete engineering workflow in Hugging Face or OneCode environments. |
+| Multi-GPU training | Launch distributed data-parallel training with `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/UnetDif --local-dir ./UnetDif
+cd UnetDif
+```
+
+### 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 first. DTK 25.04.2 or later, or the OneScience-recommended version matching the 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
+
+Training data contain ECMWF weather forecasts, CMPA precipitation observations, and static terrain. Each sample uses 39 input channels, with corresponding CMPA precipitation and ECMWF precipitation bias as targets. The data use a `56×56` Yangtze River Delta grid at `0.125°` resolution. Each rain day contains eight consecutive three-hour periods, with precipitation in `mm/3h`. Synthetic data validate training, inference, and evaluation only and do not represent the distributions of ECMWF, CMPA, real terrain, or formal paper performance.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+For single-device 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 learns dry-area, false-alarm, and signed precipitation-bias tasks and saves the model checkpoint and training metrics. Results are saved to:
+
+```text
+result/checkpoints/unetdif.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not include weights under `weight/`. The paper does not provide a confirmed official checkpoint, and the engineering checkpoint trained on synthetic data is for workflow validation only and is not claimed to be compatible with external weights.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the checkpoint and applies all four output heads under `bias = ECMWF - CMPA` to produce nonnegative corrected rainfall. Results are saved to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation produces categorical and spatial-skill metrics for individual three-hour periods and 24-hour accumulated precipitation and saves structured results. It also generates comparisons of ECMWF, CMPA, corrected precipitation, and errors. Synthetic-data results validate the engineering workflow only and do not represent formal paper performance.
+
+```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 UnetDif specifications.
+
+Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.

conf/config.yaml

@@ -0,0 +1,46 @@
+seed: 42
+data:
+  root: data
+  format_version: biastarget_rainfall_y_rd_v1
+  region: YRD
+  grid_resolution_degrees: 0.125
+  boundary_mode: half_open_engineering_assumption
+  channels: 39
+  height: 56
+  width: 56
+  time_steps: 8
+  interval_hours: 3
+  lead_hours: [15, 18, 21, 24, 27, 30, 33, 36]
+  precipitation_channel: 37
+  rain_threshold_mm_3h: 0.1
+  rain_days: {train: 2, val: 1, test: 1}
+model:
+  in_channels: 39
+  base_channels: 8
+  depth: 3
+paper_model:
+  name: UnetDif
+  base_channels: null
+  depth: null
+train:
+  epochs: 2
+  batch_size: 16
+  learning_rate: 0.0001
+  optimizer: Adam
+  optimizer_source: engineering_assumption
+  num_workers: 0
+loss:
+  focal_alpha: 0.25
+  focal_gamma: 2.0
+  first_epoch_component_normalization: true
+inference:
+  batch_size: 16
+  probability_threshold: 0.5
+evaluation:
+  thresholds_mm_24h: [10, 25, 50]
+  fss_windows: [1, 3, 7, 11, 15, 19, 23, 27]
+paths:
+  checkpoint: result/checkpoints/unetdif.pt
+  training_metrics: result/training/metrics.json
+  inference: result/output/predictions.npz
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,18 @@
+{
+  "model_name": "UnetDif",
+  "model_type": "unetdif",
+  "architectures": ["UnetDif"],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "ECMWF precipitation bias correction",
+  "license": "cc-by-nc-nd-4.0",
+  "implementation": {
+    "entry_point": "model/unetdif.py",
+    "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": 39, "output_heads": 4, "base_channels": 8, "depth": 3, "paper_base_channels": null, "paper_depth": null},
+  "data": {"region": "YRD", "resolution_degrees": 0.125, "shape": ["rain_day", 8, 39, 56, 56], "lead_hours": [15, 18, 21, 24, 27, 30, 33, 36]}
+}

configuration.json

@@ -0,0 +1,11 @@
+{
+  "framework": "PyTorch",
+  "task": "precipitation_bias_correction",
+  "model": "UnetDif",
+  "input_format": "BCHW",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "distributed_training": "torchrun"
+}

model/unetdif.py

@@ -0,0 +1,115 @@
+"""Bias-target U-Net (UnetDif) for ECMWF precipitation correction."""
+
+from typing import Dict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class ConvBlock(nn.Sequential):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__(
+            nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+
+class UnetDif(nn.Module):
+    """Paper model name; paper width/depth are unavailable and configurable."""
+
+    def __init__(self, in_channels: int = 39, base_channels: int = 8, depth: int = 3):
+        super().__init__()
+        if depth < 1:
+            raise ValueError("depth must be at least one")
+        widths = [base_channels * (2**level) for level in range(depth + 1)]
+        self.encoders = nn.ModuleList()
+        current = in_channels
+        for width in widths[:-1]:
+            self.encoders.append(ConvBlock(current, width))
+            current = width
+        self.pool = nn.MaxPool2d(2)
+        self.bottleneck = ConvBlock(widths[-2], widths[-1])
+        self.upconvs = nn.ModuleList()
+        self.decoders = nn.ModuleList()
+        for level in range(depth - 1, -1, -1):
+            self.upconvs.append(nn.ConvTranspose2d(widths[level + 1], widths[level], 2, stride=2))
+            self.decoders.append(ConvBlock(widths[level] * 2, widths[level]))
+        self.head = nn.Conv2d(widths[0], 4, 1)
+
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        skips = []
+        for encoder in self.encoders:
+            x = encoder(x)
+            skips.append(x)
+            x = self.pool(x)
+        x = self.bottleneck(x)
+        for upconv, decoder, skip in zip(self.upconvs, self.decoders, reversed(skips)):
+            x = upconv(x)
+            if x.shape[-2:] != skip.shape[-2:]:
+                x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False)
+            x = decoder(torch.cat((skip, x), dim=1))
+        output = self.head(x)
+        return {
+            "dry_logits": output[:, 0],
+            "false_alarm_logits": output[:, 1],
+            "positive_bias": F.softplus(output[:, 2]),
+            "negative_bias": F.softplus(output[:, 3]),
+        }
+
+
+def make_masks(ecmwf: torch.Tensor, cmpa: torch.Tensor, rain_threshold: float = 0.1) -> Dict[str, torch.Tensor]:
+    """Build four disjoint masks from bias = ECMWF - CMPA."""
+    ecmwf_wet = ecmwf >= rain_threshold
+    cmpa_wet = cmpa >= rain_threshold
+    bias = ecmwf - cmpa
+    return {
+        "dry": (~ecmwf_wet & ~cmpa_wet).float(),
+        "false_alarm": (ecmwf_wet & ~cmpa_wet).float(),
+        "positive": (cmpa_wet & (bias >= 0)).float(),
+        "negative": (cmpa_wet & (bias < 0)).float(),
+    }
+
+
+def apply_correction(
+    ecmwf: torch.Tensor, heads: Dict[str, torch.Tensor], probability_threshold: float = 0.5
+) -> torch.Tensor:
+    """Apply predicted signed bias and dry masks, then enforce non-negativity."""
+    corrected = (ecmwf - heads["positive_bias"] + heads["negative_bias"]).clamp_min(0)
+    dry = torch.sigmoid(heads["dry_logits"]) >= probability_threshold
+    false_alarm = torch.sigmoid(heads["false_alarm_logits"]) >= probability_threshold
+    return corrected.masked_fill(dry | false_alarm, 0)
+
+
+def _focal(logits: torch.Tensor, target: torch.Tensor, alpha: float, gamma: float) -> torch.Tensor:
+    bce = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
+    probability = torch.sigmoid(logits)
+    pt = probability * target + (1 - probability) * (1 - target)
+    alpha_t = alpha * target + (1 - alpha) * (1 - target)
+    return (alpha_t * (1 - pt).pow(gamma) * bce).mean()
+
+
+def _masked_mse(prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    return (((prediction - target).square() * mask).sum() / mask.sum().clamp_min(1))
+
+
+def loss_components(
+    heads: Dict[str, torch.Tensor], ecmwf: torch.Tensor, cmpa: torch.Tensor,
+    rain_threshold: float = 0.1, focal_alpha: float = 0.25, focal_gamma: float = 2.0,
+) -> Dict[str, torch.Tensor]:
+    masks = make_masks(ecmwf, cmpa, rain_threshold)
+    bias = ecmwf - cmpa
+    continuous = (ecmwf - heads["positive_bias"] + heads["negative_bias"]).clamp_min(0)
+    dry_target = masks["dry"] + masks["false_alarm"]
+    return {
+        "dry_focal": _focal(heads["dry_logits"], masks["dry"], focal_alpha, focal_gamma),
+        "false_alarm_focal": _focal(heads["false_alarm_logits"], masks["false_alarm"], focal_alpha, focal_gamma),
+        "positive_mse": _masked_mse(heads["positive_bias"], bias.clamp_min(0), masks["positive"]),
+        "negative_mse": _masked_mse(heads["negative_bias"], (-bias).clamp_min(0), masks["negative"]),
+        "dry_mae": (continuous.abs() * dry_target).sum() / dry_target.sum().clamp_min(1),
+        "all_mae": F.l1_loss(continuous, cmpa),
+    }

scripts/fake_data.py

@@ -0,0 +1,79 @@
+"""Generate grouped synthetic rain days at the real 39x56x56 input shape."""
+
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def smooth(field: np.ndarray, rounds: int = 2) -> np.ndarray:
+    for _ in range(rounds):
+        field = sum(np.roll(np.roll(field, y, -2), x, -1) for y, x in ((0, 0), (1, 0), (-1, 0), (0, 1), (0, -1))) / 5
+    return field
+
+
+def make_day(rng: np.random.Generator, steps: int, height: int, width: int):
+    yy, xx = np.mgrid[-1:1:complex(height), -1:1:complex(width)].astype(np.float32)
+    terrain = np.clip(0.55 + 0.25 * yy + 0.16 * np.sin(2.5 * xx) - 0.12 * np.cos(3 * yy), 0, 1)
+    center = rng.uniform(-0.45, 0.45, 2)
+    velocity = rng.uniform(-0.10, 0.10, 2)
+    width0 = rng.uniform(0.13, 0.27)
+    amplitude = rng.uniform(12, 28)
+    inputs = np.empty((steps, 39, height, width), np.float32)
+    targets = np.empty((steps, height, width), np.float32)
+    for step in range(steps):
+        cy, cx = center + velocity * step
+        rain = amplitude * np.exp(-((xx - cx) ** 2 + (yy - cy) ** 2) / (2 * width0**2))
+        rain += 0.35 * amplitude * np.exp(-((xx + cx * 0.5) ** 2 + (yy - cy * 0.6) ** 2) / (3 * width0**2))
+        rain = np.clip(rain * (0.8 + 0.3 * terrain) + rng.normal(0, 0.12, rain.shape), 0, None)
+        forecast = np.clip(np.roll(rain, shift=(1, -1), axis=(0, 1)) * (1.12 + 0.08 * terrain) - 0.25, 0, None)
+        channels = np.empty((39, height, width), np.float32)
+        for variable in range(8):
+            for level in range(4):
+                index = variable * 4 + level
+                noise = smooth(rng.normal(0, 0.05, rain.shape).astype(np.float32))
+                if variable == 1:  # humidity follows rain and terrain
+                    channels[index] = 0.45 + 0.02 * rain + 0.12 * terrain - 0.025 * level + noise
+                elif variable == 4:  # vertical velocity is strongest near moving rain
+                    channels[index] = -0.08 * rain / (level + 1) + 0.04 * (xx * velocity[1] + yy * velocity[0]) + noise
+                else:
+                    channels[index] = (variable + 1) * 0.12 + level * 0.04 + 0.08 * xx - 0.05 * yy + noise
+        for index in range(32, 37):
+            channels[index] = 0.1 * (index - 31) + 0.03 * rain + smooth(rng.normal(0, 0.04, rain.shape).astype(np.float32))
+        channels[37] = forecast
+        channels[38] = terrain
+        inputs[step], targets[step] = channels, rain
+    return inputs, targets, terrain
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = config["data"]
+    expected = (8, 39, 56, 56)
+    actual = (int(data["time_steps"]), int(data["channels"]), int(data["height"]), int(data["width"]))
+    if actual != expected or len(data["lead_hours"]) != expected[0]:
+        raise ValueError(f"fixed data contract is {expected} with eight lead hours, got {actual}")
+    counts = data["rain_days"]
+    output = ROOT / data["root"]
+    output.mkdir(parents=True, exist_ok=True)
+    seed = int(config["seed"])
+    group_id = 0
+    for split_index, split in enumerate(("train", "val", "test")):
+        arrays, targets, terrains, groups = [], [], [], []
+        for day in range(int(counts[split])):
+            x, y, terrain = make_day(np.random.default_rng(seed + group_id), int(data["time_steps"]), int(data["height"]), int(data["width"]))
+            arrays.append(x); targets.append(y); terrains.append(terrain); groups.append(group_id)
+            group_id += 1
+        np.savez_compressed(
+            output / f"{split}.npz", inputs=np.stack(arrays), targets=np.stack(targets),
+            terrain=np.stack(terrains), group_id=np.asarray(groups), split=np.asarray(split),
+            format_version=np.asarray(data["format_version"]), lead_hours=np.asarray(data["lead_hours"]),
+        )
+        print(f"split={split} rain_days={counts[split]} inputs={np.stack(arrays).shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,62 @@
+"""Apply all four UnetDif heads and save 3-hour and 24-hour products."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.unetdif import UnetDif, apply_correction
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=True)
+    required_checkpoint = {"epoch", "model", "optimizer", "model_config", "loss_normalizers", "loss_components", "format_version"}
+    if not required_checkpoint.issubset(checkpoint):
+        raise ValueError(f"checkpoint is missing {sorted(required_checkpoint - checkpoint.keys())}")
+    if checkpoint["format_version"] != config["data"]["format_version"]:
+        raise ValueError("checkpoint and configured data format versions differ")
+    model = UnetDif(**checkpoint["model_config"]).to(device)
+    model.load_state_dict(checkpoint["model"]); model.eval()
+    archive = np.load(ROOT / config["data"]["root"] / "test.npz", allow_pickle=False)
+    inputs = archive["inputs"].astype(np.float32, copy=False)
+    expected = (int(config["data"]["time_steps"]), int(config["data"]["channels"]),
+                int(config["data"]["height"]), int(config["data"]["width"]))
+    if inputs.shape[1:] != expected or archive["targets"].shape[1:] != (expected[0], expected[2], expected[3]):
+        raise ValueError(f"test NPZ does not satisfy the fixed day contract {expected}")
+    if archive["lead_hours"].tolist() != config["data"]["lead_hours"]:
+        raise ValueError("test NPZ lead_hours do not match configuration")
+    flat = torch.from_numpy(inputs.reshape(-1, *inputs.shape[2:])).to(device)
+    heads_all = {key: [] for key in ("dry_logits", "false_alarm_logits", "positive_bias", "negative_bias")}
+    corrected = []
+    batch_size = int(config["inference"]["batch_size"])
+    with torch.no_grad():
+        for start in range(0, len(flat), batch_size):
+            batch = flat[start:start + batch_size]
+            heads = model(batch)
+            corrected.append(apply_correction(batch[:, int(config["data"]["precipitation_channel"])], heads,
+                                              float(config["inference"]["probability_threshold"])).cpu())
+            for key in heads_all:
+                heads_all[key].append(heads[key].cpu())
+    day_shape = inputs.shape[:2] + inputs.shape[-2:]
+    corrected_3h = torch.cat(corrected).numpy().reshape(day_shape)
+    raw_3h = inputs[:, :, int(config["data"]["precipitation_channel"])]
+    output = ROOT / config["paths"]["inference"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    payload = {key: torch.cat(value).numpy().reshape(day_shape) for key, value in heads_all.items()}
+    np.savez_compressed(
+        output, inputs=inputs, group_id=archive["group_id"], lead_hours=archive["lead_hours"],
+        raw_3h=raw_3h, target_3h=archive["targets"], corrected_3h=corrected_3h,
+        raw_24h=raw_3h.sum(1), target_24h=archive["targets"].sum(1), corrected_24h=corrected_3h.sum(1), **payload,
+    )
+    print(f"saved={output.relative_to(ROOT)} corrected_3h={corrected_3h.shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,86 @@
+"""Evaluate 24-hour rainfall and produce per-step metrics and maps."""
+
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def categorical(prediction, target, threshold):
+    forecast, observed = prediction >= threshold, target >= threshold
+    hits = np.logical_and(forecast, observed).sum(); false = np.logical_and(forecast, ~observed).sum()
+    misses = np.logical_and(~forecast, observed).sum(); correct_negative = np.logical_and(~forecast, ~observed).sum()
+    total = hits + false + misses + correct_negative; eps = 1e-8
+    random_hits = (hits + misses) * (hits + false) / max(total, 1)
+    return {
+        "ACC": float((hits + correct_negative) / (total + eps)),
+        "TS": float(hits / (hits + misses + false + eps)),
+        "ETS": float((hits - random_hits) / (hits + misses + false - random_hits + eps)),
+        "BS": float((hits + false) / (hits + misses + eps)),
+        "FAR": float(false / (hits + false + eps)),
+    }
+
+
+def fss(prediction, target, threshold, window):
+    forecast = torch.from_numpy((prediction >= threshold).astype(np.float32))[:, None]
+    observed = torch.from_numpy((target >= threshold).astype(np.float32))[:, None]
+    pf = torch.nn.functional.avg_pool2d(forecast, window, stride=1, padding=window // 2)
+    po = torch.nn.functional.avg_pool2d(observed, window, stride=1, padding=window // 2)
+    return float(1 - (pf - po).square().sum() / (pf.square().sum() + po.square().sum() + 1e-8))
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    archive = np.load(ROOT / config["paths"]["inference"], allow_pickle=False)
+    steps = int(config["data"]["time_steps"])
+    lead_hours = [int(value) for value in archive["lead_hours"]]
+    expected_3h = (steps, int(config["data"]["height"]), int(config["data"]["width"]))
+    for key in ("raw_3h", "target_3h", "corrected_3h"):
+        if archive[key].shape[1:] != expected_3h:
+            raise ValueError(f"{key} must have per-day shape {expected_3h}")
+    if lead_hours != config["data"]["lead_hours"]:
+        raise ValueError("inference NPZ lead_hours do not match configuration")
+    thresholds = [float(value) for value in config["evaluation"]["thresholds_mm_24h"]]
+    windows = [int(value) for value in config["evaluation"]["fss_windows"]]
+    metrics = {"units_24h": "mm/24h", "units_per_step": "mm/3h", "per_step_length": steps,
+               "lead_hours": lead_hours, "raw_24h": {}, "corrected_24h": {}, "per_step": {}}
+    for label in ("raw", "corrected"):
+        prediction_24h = archive[f"{label}_24h"]
+        metrics[f"{label}_24h"] = {str(int(t)): categorical(prediction_24h, archive["target_24h"], t) for t in thresholds}
+        metrics[f"{label}_24h"]["FSS@50"] = {str(w): fss(prediction_24h, archive["target_24h"], 50, w) for w in windows}
+        per_step = []
+        for step in range(steps):
+            item = {"step": step, "lead_hour": lead_hours[step]}
+            item.update({str(int(t)): categorical(archive[f"{label}_3h"][:, step], archive["target_3h"][:, step], t) for t in thresholds})
+            item["FSS@50"] = {str(w): fss(archive[f"{label}_3h"][:, step], archive["target_3h"][:, step], 50, w) for w in windows}
+            per_step.append(item)
+        metrics["per_step"][label] = per_step
+    output = ROOT / config["paths"]["evaluation_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    figure, axes = plt.subplots(4, steps, figsize=(18, 9), constrained_layout=True)
+    rows = (archive["raw_3h"][0], archive["target_3h"][0], archive["corrected_3h"][0], archive["corrected_3h"][0] - archive["target_3h"][0])
+    labels = ("ECMWF", "CMPA target", "Corrected", "Error")
+    vmax = max(float(rows[0].max()), float(rows[1].max()), 1)
+    for row, (fields, label) in enumerate(zip(rows, labels)):
+        for step in range(steps):
+            cmap = "RdBu_r" if row == 3 else "Blues"
+            limit = vmax if row < 3 else max(vmax / 2, 1)
+            axes[row, step].imshow(fields[step], cmap=cmap, vmin=-limit if row == 3 else 0, vmax=limit)
+            axes[row, step].axis("off")
+            if row == 0: axes[row, step].set_title(f"+{archive['lead_hours'][step]} h")
+            if step == 0: axes[row, step].set_ylabel(label)
+    figure.savefig(output / "comparison.png", dpi=150); plt.close(figure)
+    print(f"metrics={output / 'metrics.json'} figure={output / 'comparison.png'}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,115 @@
+"""Train UnetDif with single-process or torchrun DDP execution."""
+
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.unetdif import UnetDif, loss_components
+
+
+class RainfallDataset(Dataset):
+    def __init__(self, path: Path, config):
+        archive = np.load(path, allow_pickle=False)
+        expected = (int(config["time_steps"]), int(config["channels"]), int(config["height"]), int(config["width"]))
+        if archive["inputs"].shape[1:] != expected or archive["targets"].shape[1:] != (expected[0], expected[2], expected[3]):
+            raise ValueError(f"expected day shapes {expected} and {(expected[0], expected[2], expected[3])}")
+        if len(np.unique(archive["group_id"])) != len(archive["group_id"]):
+            raise ValueError("each rain day must belong to one split group")
+        if archive["lead_hours"].tolist() != config["lead_hours"]:
+            raise ValueError("NPZ lead_hours do not match the configured eight 3-hour steps")
+        self.inputs = archive["inputs"].astype(np.float32, copy=False).reshape(-1, *expected[1:])
+        self.targets = archive["targets"].astype(np.float32, copy=False).reshape(-1, expected[2], expected[3])
+
+    def __len__(self):
+        return len(self.inputs)
+
+    def __getitem__(self, index):
+        return torch.from_numpy(self.inputs[index]), torch.from_numpy(self.targets[index])
+
+
+def get_device(local_rank=0, local_world_size=1):
+    use_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= local_world_size
+    return torch.device("cuda", local_rank) if use_cuda else torch.device("cpu")
+
+
+def reduce_values(values, device, distributed):
+    tensor = torch.tensor(values, device=device)
+    if distributed:
+        torch.distributed.all_reduce(tensor)
+    return tensor.cpu().tolist()
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    torch.manual_seed(int(config["seed"]))
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "1"))
+    device = get_device(local_rank, local_world_size)
+    if distributed:
+        torch.distributed.init_process_group("nccl" if device.type == "cuda" else "gloo")
+    rank = torch.distributed.get_rank() if distributed else 0
+    dataset = RainfallDataset(ROOT / config["data"]["root"] / "train.npz", config["data"])
+    if distributed and len(dataset) < torch.distributed.get_world_size():
+        raise ValueError("DDP requires at least one training item per process")
+    sampler = DistributedSampler(dataset) if distributed else None
+    loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), shuffle=sampler is None,
+                        sampler=sampler, num_workers=int(config["train"]["num_workers"]))
+    model = UnetDif(**config["model"]).to(device)
+    if distributed:
+        model = DDP(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(config["train"]["learning_rate"]))
+    names = ["dry_focal", "false_alarm_focal", "positive_mse", "negative_mse", "dry_mae", "all_mae"]
+    normalizers = {name: 1.0 for name in names}
+    history = []
+    for epoch in range(int(config["train"]["epochs"])):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train(); sums = {name: 0.0 for name in names}; batches = 0
+        for inputs, target in loader:
+            inputs, target = inputs.to(device), target.to(device)
+            heads = model(inputs)
+            raw = loss_components(heads, inputs[:, int(config["data"]["precipitation_channel"])], target,
+                                  float(config["data"]["rain_threshold_mm_3h"]),
+                                  float(config["loss"]["focal_alpha"]), float(config["loss"]["focal_gamma"]))
+            loss = sum(raw[name] * normalizers[name] for name in names)
+            optimizer.zero_grad(set_to_none=True); loss.backward(); optimizer.step()
+            for name in names:
+                sums[name] += float(raw[name].detach())
+            batches += 1
+        reduced = reduce_values([sums[name] for name in names] + [batches], device, distributed)
+        total_batches = max(reduced[-1], 1)
+        averages = {name: reduced[index] / total_batches for index, name in enumerate(names)}
+        if epoch == 0:
+            normalizers = {name: 1.0 / max(value, 1e-6) for name, value in averages.items()}
+        record = {"epoch": epoch + 1, "components": averages, "normalizers": normalizers}
+        history.append(record)
+        if rank == 0:
+            print(json.dumps(record))
+            checkpoint = ROOT / config["paths"]["checkpoint"]
+            checkpoint.parent.mkdir(parents=True, exist_ok=True)
+            state = model.module.state_dict() if distributed else model.state_dict()
+            torch.save({"epoch": epoch + 1, "model": state, "optimizer": optimizer.state_dict(),
+                         "model_config": config["model"], "loss_normalizers": normalizers,
+                        "loss_components": names,
+                         "format_version": config["data"]["format_version"]}, checkpoint)
+    if rank == 0:
+        metrics = ROOT / config["paths"]["training_metrics"]
+        metrics.parent.mkdir(parents=True, exist_ok=True)
+        metrics.write_text(json.dumps({"history": history}, indent=2) + "\n")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()