Publish RainNet model package

.gitattributes

@@ -1,35 +1,14 @@
-*.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
+*.hdf5 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
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl 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

README.md

@@ -0,0 +1,187 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth Science
+- precipitation nowcasting
+- weather radar
+- RYDL
+frameworks: PyTorch
+datasets:
+- RYDL
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">RainNet</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Paper: RainNet v1.0: a convolutional neural network for radar-based precipitation nowcasting  
+https://doi.org/10.5194/gmd-13-2631-2020
+
+RainNet is designed for radar-based precipitation nowcasting. It takes four consecutive radar precipitation fields at 5 min intervals as input, predicts the precipitation field for the next 5 min, and can be extended recursively to a lead time of approximately 60 min.
+
+# Model Description
+
+RainNet was proposed by the authors of the original paper and trained on the German Weather Service (DWD) RY radar precipitation product. It performs radar nowcasting as a regression task for continuous precipitation intensity.
+
+The current implementation takes four consecutive historical frames as input and uses the immediately following time step, `i+4`, as the target. Precipitation values are transformed with `x -> log(x + 0.01)` before entering the model. Each raw `900x900` radar field is expanded to `928x928` with reflect/mirror padding and cropped back to `900x900` after prediction. RainNet has approximately 31.4M parameters; the validated parameter count is 31,380,613. The decoder uses nearest-neighbor upsampling.
+
+# Applicable Scenarios
+
+| Scenario | Description |
+| :---: | :--- |
+| Continuous precipitation regression training | Train RainNet with consecutive radar precipitation fields. |
+| Radar precipitation nowcasting | Predict future precipitation from consecutive historical radar fields. |
+| Local pipeline validation | Use Fake Data to validate data loading, training, inference, evaluation, and visualization. |
+| ModelScope / OneCode execution | Run the project as a standalone model package. |
+| Multi-GPU training | Launch distributed training processes with `torchrun`. |
+
+# Usage
+
+Run the following commands from the root of the model package. The default smoke-test configuration preserves the full `900x900` spatial grid and uses 1 epoch with at most 1 batch per stage to validate the engineering pipeline.
+
+## 1. OneCode
+
+Use the OneCode online environment for intelligent one-click AI4S programming:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+Python 3.11, PyTorch, NumPy, h5py, PyYAML, and Matplotlib are required. After installing the environment, run Fake Data generation, training, inference, and evaluation in sequence.
+
+### Hardware Requirements
+
+The model has approximately 31.4M parameters. Training with full `928x928` internal tensors requires substantial accelerator memory, so a CUDA/HIP-compatible GPU or DCU with sufficient memory is recommended. CPU execution is supported but substantially slower.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/RainNet --local_dir ./RainNet
+cd RainNet
+```
+
+### Install the Runtime Environment
+
+#### 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
+
+#### Real Data
+
+The real dataset is RYDL, available from https://doi.org/10.5281/zenodo.3629951. It uses HDF5, with a raw frame size of `900x900`, a spatial resolution of 1 km, and a temporal resolution of 5 min. Each top-level HDF5 timestamp key corresponds to one two-dimensional precipitation field. This repository does not include or automatically download the complete real dataset.
+
+#### Fake Data
+
+```bash
+python scripts/fake_data.py
+```
+
+This command generates `data/rainnet_fake.hdf5`. Fake Data preserves the real `900x900` frame size, simulates the RYDL HDF5 timestamp key-value organization, maintains a continuous sequence at 5 min intervals, and only reduces the number of time frames.
+
+Fake Data is only used to validate the engineering pipeline. It does not represent real precipitation forecasting performance and does not reproduce the accuracy reported in the paper.
+
+### Training
+
+This reproduction uses Log-Cosh Loss and the Adam optimizer, with a default learning rate of `1e-4`.
+
+#### Single-Accelerator Training
+
+```bash
+python scripts/fake_data.py
+python scripts/train.py
+```
+
+#### Distributed Training
+
+```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 Weights
+
+This repository plans to provide weights trained on DWD RY/RYDL radar precipitation data under `weight/`. The weight files will be uploaded in a future update.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference results are saved under `result/output/`. The script loads a training checkpoint, performs one-step inference, and runs a 12-step autoregressive rollout while updating the four-frame sliding window.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+The script computes MAE, CSI, FSS, and the Persistence baseline from actual inference outputs, and generates forecast comparisons, a training-loss curve, and metric plots.
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- The OneScience RainNet model package is licensed under Apache License 2.0.
+- The original RainNet source code is licensed under the MIT License.
+- The RainNet paper was published by Copernicus Publications under the Creative Commons Attribution 4.0 License (CC BY 4.0).
+- Attribution to the original paper and authors must be retained when using or redistributing this package.
+
+```bibtex
+@article{ayzel2020rainnet,
+  title={RainNet v1.0: a convolutional neural network for radar-based precipitation nowcasting},
+  author={Ayzel, Georgy and Scheffer, Tobias and Heistermann, Maik},
+  journal={Geoscientific Model Development},
+  volume={13},
+  pages={2631--2644},
+  year={2020},
+  doi={10.5194/gmd-13-2631-2020}
+}
+```

conf/config.yaml

@@ -0,0 +1,59 @@
+seed: 42
+device: auto
+
+data:
+  mode: synthetic
+  path: ./data/rainnet_fake.hdf5
+  input_steps: 4
+  forecast_steps: 1
+  interval_minutes: 5
+  raw_height: 900
+  raw_width: 900
+  padded_height: 928
+  padded_width: 928
+  pixel_size_km: 1
+  train_frames: 8
+  val_frames: 6
+  test_frames: 16
+  num_workers: 0
+
+model:
+  in_channels: 4
+  out_channels: 1
+  dropout: 0.5
+  mode: regression
+
+train:
+  epochs: 1
+  batch_size: 1
+  learning_rate: 0.0001
+  max_train_batches: 1
+  max_valid_batches: 1
+  checkpoint_last: ./weight/rainnet_last.pth
+  checkpoint_best: ./weight/rainnet_best.pth
+
+paper_reference:
+  epochs: 10
+  batch_size: 2
+  learning_rate: 0.0001
+  raw_size: 900
+  padded_size: 928
+
+inference:
+  checkpoint: ./weight/rainnet_best.pth
+  rollout_steps: 12
+  output_dir: ./result/output
+
+evaluation:
+  thresholds_mm_h:
+    - 0.125
+    - 1.0
+    - 5.0
+    - 10.0
+    - 15.0
+  fss_windows_km:
+    - 1
+    - 5
+    - 10
+    - 20
+  result_dir: ./result

configuration.json

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

model/rainnet.py

@@ -0,0 +1,80 @@
+"""PyTorch implementation of RainNet v1.0."""
+
+import torch
+from torch import nn
+
+
+class RainNet(nn.Module):
+    """The 20-convolution encoder-decoder from Ayzel et al. (2020)."""
+
+    def __init__(self, in_channels=4, out_channels=1, dropout=0.5, mode="regression"):
+        super().__init__()
+        if mode != "regression":
+            raise ValueError("This package implements RainNet regression mode only")
+        self.in_channels = in_channels
+        self.pool = nn.MaxPool2d(2)
+        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv1f = nn.Conv2d(in_channels, 64, 3, padding=1)
+        self.conv1s = nn.Conv2d(64, 64, 3, padding=1)
+        self.conv2f = nn.Conv2d(64, 128, 3, padding=1)
+        self.conv2s = nn.Conv2d(128, 128, 3, padding=1)
+        self.conv3f = nn.Conv2d(128, 256, 3, padding=1)
+        self.conv3s = nn.Conv2d(256, 256, 3, padding=1)
+        self.conv4f = nn.Conv2d(256, 512, 3, padding=1)
+        self.conv4s = nn.Conv2d(512, 512, 3, padding=1)
+        self.drop4 = nn.Dropout(dropout)
+        self.conv5f = nn.Conv2d(512, 1024, 3, padding=1)
+        self.conv5s = nn.Conv2d(1024, 1024, 3, padding=1)
+        self.drop5 = nn.Dropout(dropout)
+
+        self.conv6f = nn.Conv2d(1536, 512, 3, padding=1)
+        self.conv6s = nn.Conv2d(512, 512, 3, padding=1)
+        self.conv7f = nn.Conv2d(768, 256, 3, padding=1)
+        self.conv7s = nn.Conv2d(256, 256, 3, padding=1)
+        self.conv8f = nn.Conv2d(384, 128, 3, padding=1)
+        self.conv8s = nn.Conv2d(128, 128, 3, padding=1)
+        self.conv9f = nn.Conv2d(192, 64, 3, padding=1)
+        self.conv9s = nn.Conv2d(64, 64, 3, padding=1)
+        self.conv9out = nn.Conv2d(64, 2, 3, padding=1)
+        self.output = nn.Conv2d(2, out_channels, 1)
+        self.apply(self._initialize)
+
+    @staticmethod
+    def _initialize(module):
+        if isinstance(module, nn.Conv2d):
+            nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu")
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+    def forward(self, x):
+        if x.ndim != 4:
+            raise ValueError(f"RainNet expects BCHW input, got shape {tuple(x.shape)}")
+        if x.shape[1] != self.in_channels:
+            raise ValueError(f"RainNet expects {self.in_channels} channels, got {x.shape[1]}")
+        if x.shape[-2] % 16 or x.shape[-1] % 16:
+            raise ValueError(
+                f"Spatial dimensions must be divisible by 16, got {tuple(x.shape[-2:])}"
+            )
+
+        c1 = self.relu(self.conv1s(self.relu(self.conv1f(x))))
+        c2 = self.relu(self.conv2s(self.relu(self.conv2f(self.pool(c1)))))
+        c3 = self.relu(self.conv3s(self.relu(self.conv3f(self.pool(c2)))))
+        c4 = self.relu(self.conv4s(self.relu(self.conv4f(self.pool(c3)))))
+        c5 = self.relu(self.conv5s(self.relu(self.conv5f(self.pool(self.drop4(c4))))))
+        c5 = self.drop5(c5)
+
+        c6 = torch.cat((self.upsample(c5), c4), dim=1)
+        c6 = self.relu(self.conv6s(self.relu(self.conv6f(c6))))
+        c7 = torch.cat((self.upsample(c6), c3), dim=1)
+        c7 = self.relu(self.conv7s(self.relu(self.conv7f(c7))))
+        c8 = torch.cat((self.upsample(c7), c2), dim=1)
+        c8 = self.relu(self.conv8s(self.relu(self.conv8f(c8))))
+        c9 = torch.cat((self.upsample(c8), c1), dim=1)
+        c9 = self.relu(self.conv9s(self.relu(self.conv9f(c9))))
+        return self.output(self.relu(self.conv9out(c9)))
+
+
+def build_rainnet(in_channels=4, out_channels=1, dropout=0.5, mode="regression"):
+    return RainNet(in_channels, out_channels, dropout, mode)

scripts/fake_data.py

@@ -0,0 +1,71 @@
+"""Generate a compact, full-resolution synthetic RYDL-like sequence."""
+
+from datetime import datetime, timedelta
+from pathlib import Path
+import sys
+
+import h5py
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_config():
+    with (ROOT / "conf/config.yaml").open(encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def precipitation_frame(y, x, step, rng):
+    cells = (
+        (220 + 4.2 * step, 260 + 6.0 * step, 70 + 0.5 * step, 45, 1.25),
+        (610 - 3.2 * step, 590 + 2.7 * step, 58, 82 - 0.4 * step, 0.85),
+        (430 + 1.5 * step, 710 - 4.0 * step, 42 + 0.6 * step, 55, 0.55),
+    )
+    field = np.zeros_like(x, dtype=np.float32)
+    for cy, cx, sy, sx, amplitude in cells:
+        evolution = 1.0 + 0.12 * np.sin((step + amplitude) / 3.0)
+        field += amplitude * evolution * np.exp(
+            -0.5 * (((x - cx) / sx) ** 2 + ((y - cy) / sy) ** 2)
+        )
+    # A smooth perturbation evolves with the cells without decorrelating frames.
+    phase = rng.uniform(-0.03, 0.03)
+    field *= 1.0 + 0.025 * np.sin(x / 35.0 + step / 4.0 + phase) * np.cos(y / 47.0)
+    return np.maximum(field, 0).astype(np.float32)
+
+
+def main():
+    config = load_config()
+    data = config["data"]
+    np.random.seed(config["seed"])
+    rng = np.random.default_rng(config["seed"])
+    total = data["train_frames"] + data["val_frames"] + data["test_frames"]
+    height, width = data["raw_height"], data["raw_width"]
+    yy, xx = np.mgrid[:height, :width].astype(np.float32)
+    path = ROOT / data["path"]
+    path.parent.mkdir(parents=True, exist_ok=True)
+    start = datetime(2017, 1, 1)
+    keys = []
+    with h5py.File(path, "w") as handle:
+        handle.attrs["units"] = "mm/5min"
+        handle.attrs["interval_minutes"] = data["interval_minutes"]
+        for index in range(total):
+            key = (start + timedelta(minutes=index * data["interval_minutes"])).strftime(
+                "%Y%m%d%H%M"
+            )
+            keys.append(key)
+            handle.create_dataset(key, data=precipitation_frame(yy, xx, index, rng))
+    with h5py.File(path, "r") as handle:
+        first = handle[keys[0]][...]
+    print(f"Number of frames: {len(keys)}")
+    print(f"First key: {keys[0]}")
+    print(f"Last key: {keys[-1]}")
+    print(f"Frame shape: {first.shape}")
+    print(f"dtype: {first.dtype}")
+    print(f"min: {first.min():.8f}")
+    print(f"max: {first.max():.8f}")
+    print(f"mean: {first.mean():.8f}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,107 @@
+"""Checkpoint-backed one-step and autoregressive RainNet inference."""
+
+import json
+from pathlib import Path
+import sys
+
+import h5py
+import numpy as np
+import torch
+import torch.nn.functional as F
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.rainnet import build_rainnet
+
+
+def load_config():
+    with (ROOT / "conf/config.yaml").open(encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def inverse_transform(tensor):
+    return torch.clamp(torch.exp(tensor) - 0.01, min=0.0)
+
+
+def describe(name, array):
+    print(
+        f"{name}: shape={array.shape}, dtype={array.dtype}, "
+        f"min={array.min():.8f}, max={array.max():.8f}, mean={array.mean():.8f}"
+    )
+
+
+def main():
+    config = load_config()
+    torch.manual_seed(config["seed"])
+    requested_device = config["device"]
+    if requested_device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(requested_device)
+    checkpoint_path = ROOT / config["inference"]["checkpoint"]
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Required checkpoint does not exist: {checkpoint_path}")
+    model = build_rainnet(**config["model"]).to(device)
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+    print("checkpoint_loaded: True")
+    print(f"checkpoint path: {checkpoint_path}")
+
+    data = config["data"]
+    data_path = ROOT / data["path"]
+    with h5py.File(data_path, "r") as handle:
+        keys = sorted(handle.keys())
+        test_start = data["train_frames"] + data["val_frames"]
+        test_keys = keys[test_start : test_start + data["test_frames"]]
+        raw = np.stack([handle[key][...] for key in test_keys]).astype(np.float32)
+    needed = data["input_steps"] + config["inference"]["rollout_steps"]
+    if len(raw) < needed:
+        raise ValueError(f"Inference needs {needed} test frames, found {len(raw)}")
+
+    inputs = torch.from_numpy(raw[: data["input_steps"]]).unsqueeze(0).to(device)
+    targets = raw[data["input_steps"] : needed]
+    log_window = torch.log(inputs + 0.01)
+    pad_h = data["padded_height"] - data["raw_height"]
+    pad_w = data["padded_width"] - data["raw_width"]
+    pad = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
+    crop = (pad_h // 2, pad_w // 2)
+    predictions = []
+    with torch.inference_mode():
+        for _ in range(config["inference"]["rollout_steps"]):
+            padded = F.pad(log_window, pad, mode="reflect")
+            padded_prediction = model(padded)
+            prediction = padded_prediction[
+                :, :, crop[0] : crop[0] + data["raw_height"], crop[1] : crop[1] + data["raw_width"]
+            ]
+            predictions.append(inverse_transform(prediction).squeeze(0).squeeze(0).cpu().numpy())
+            log_window = torch.cat((log_window[:, 1:], prediction), dim=1)
+    predictions = np.stack(predictions).astype(np.float32)
+    persistence = np.repeat(raw[data["input_steps"] - 1][None], len(predictions), axis=0).astype(np.float32)
+    output_dir = ROOT / config["inference"]["output_dir"]
+    output_dir.mkdir(parents=True, exist_ok=True)
+    np.save(output_dir / "inputs.npy", raw[: data["input_steps"]])
+    np.save(output_dir / "predictions.npy", predictions)
+    np.save(output_dir / "targets.npy", targets)
+    np.save(output_dir / "persistence.npy", persistence)
+    metadata = {
+        "units": "mm/5min",
+        "interval_minutes": data["interval_minutes"],
+        "rollout_steps": len(predictions),
+        "input_keys": test_keys[: data["input_steps"]],
+        "target_keys": test_keys[data["input_steps"] : needed],
+        "checkpoint": str(checkpoint_path),
+    }
+    with (output_dir / "metadata.json").open("w", encoding="utf-8") as handle:
+        json.dump(metadata, handle, indent=2)
+    describe("inputs", raw[: data["input_steps"]])
+    describe("predictions", predictions)
+    describe("targets", targets)
+    describe("persistence", persistence)
+    print(f"One-step inference shape: {predictions[:1].shape}")
+    print(f"Autoregressive rollout steps: {len(predictions)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,145 @@
+"""Evaluate actual RainNet rollout outputs and create diagnostic figures."""
+
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn.functional as F
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_config():
+    with (ROOT / "conf/config.yaml").open(encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def csi(prediction, target, threshold):
+    pred_event, target_event = prediction >= threshold, target >= threshold
+    hits = np.logical_and(pred_event, target_event).sum(dtype=np.float64)
+    false_alarms = np.logical_and(pred_event, ~target_event).sum(dtype=np.float64)
+    misses = np.logical_and(~pred_event, target_event).sum(dtype=np.float64)
+    denominator = hits + false_alarms + misses
+    return float(hits / denominator) if denominator else 0.0
+
+
+def fss(prediction, target, threshold, window):
+    pred = torch.from_numpy((prediction >= threshold).astype(np.float32))[None, None]
+    obs = torch.from_numpy((target >= threshold).astype(np.float32))[None, None]
+    if window > 1:
+        padding = window // 2
+        pred = F.avg_pool2d(pred, window, stride=1, padding=padding)
+        obs = F.avg_pool2d(obs, window, stride=1, padding=padding)
+        pred = pred[..., : prediction.shape[0], : prediction.shape[1]]
+        obs = obs[..., : target.shape[0], : target.shape[1]]
+    numerator = torch.sum((pred - obs) ** 2)
+    denominator = torch.sum(pred**2) + torch.sum(obs**2)
+    return float(1.0 - numerator / denominator) if denominator > 0 else 0.0
+
+
+def metric_series(predictions, targets, thresholds, windows):
+    output = []
+    for index, (prediction, target) in enumerate(zip(predictions, targets)):
+        pred_rate, target_rate = prediction * 12.0, target * 12.0
+        output.append(
+            {
+                "lead_minutes": (index + 1) * 5,
+                "mae_mm_h": float(np.mean(np.abs(pred_rate - target_rate), dtype=np.float64)),
+                "csi": {str(t): csi(pred_rate, target_rate, t) for t in thresholds},
+                "fss": {
+                    str(t): {str(w): fss(pred_rate, target_rate, t, w) for w in windows}
+                    for t in thresholds
+                },
+            }
+        )
+    return output
+
+
+def main():
+    config = load_config()
+    output_dir = ROOT / config["inference"]["output_dir"]
+    required = ["inputs.npy", "predictions.npy", "targets.npy", "persistence.npy"]
+    missing = [name for name in required if not (output_dir / name).exists()]
+    if missing:
+        raise FileNotFoundError(f"Missing inference outputs: {missing}")
+    inputs = np.load(output_dir / "inputs.npy")
+    predictions = np.load(output_dir / "predictions.npy")
+    targets = np.load(output_dir / "targets.npy")
+    persistence = np.load(output_dir / "persistence.npy")
+    for name, array in (("prediction", predictions), ("target", targets), ("persistence", persistence)):
+        if array.shape != predictions.shape or not np.isfinite(array).all():
+            raise ValueError(f"Invalid {name}: shape={array.shape}, finite={np.isfinite(array).all()}")
+    thresholds = config["evaluation"]["thresholds_mm_h"]
+    windows = config["evaluation"]["fss_windows_km"]
+    rainnet_metrics = metric_series(predictions, targets, thresholds, windows)
+    persistence_metrics = metric_series(persistence, targets, thresholds, windows)
+    metrics = {
+        "units": "mm/h",
+        "rainnet": rainnet_metrics,
+        "persistence": persistence_metrics,
+    }
+    result_dir = ROOT / config["evaluation"]["result_dir"]
+    result_dir.mkdir(parents=True, exist_ok=True)
+    with (result_dir / "metrics.json").open("w", encoding="utf-8") as handle:
+        json.dump(metrics, handle, indent=2)
+
+    history_path = result_dir / "train_history.json"
+    with history_path.open(encoding="utf-8") as handle:
+        history = json.load(handle)
+    fig, ax = plt.subplots(figsize=(6, 4))
+    ax.plot(history["train_loss"], marker="o", label="Train")
+    ax.plot(history["validation_loss"], marker="o", label="Validation")
+    ax.set(xlabel="Epoch", ylabel="Log-Cosh loss", title="RainNet smoke training")
+    ax.legend()
+    fig.tight_layout()
+    fig.savefig(result_dir / "loss.png", dpi=150)
+    plt.close(fig)
+
+    selected = [0, 5, 11]
+    fig, axes = plt.subplots(3, 5, figsize=(16, 10))
+    for row, index in enumerate(selected):
+        panels = [inputs[-1], targets[index], predictions[index], persistence[index], predictions[index] - targets[index]]
+        titles = ["Last Input", "Truth", "RainNet Prediction", "Persistence", "Prediction Error"]
+        for axis, panel, title in zip(axes[row], panels, titles):
+            image = axis.imshow(panel, cmap="RdBu_r" if title == "Prediction Error" else "Blues")
+            axis.set_title(f"{title}\n{(index + 1) * 5} min")
+            axis.axis("off")
+            fig.colorbar(image, ax=axis, fraction=0.046)
+    fig.tight_layout()
+    fig.savefig(result_dir / "forecast_comparison.png", dpi=120)
+    plt.close(fig)
+
+    leads = [item["lead_minutes"] for item in rainnet_metrics]
+    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+    axes[0].plot(leads, [item["mae_mm_h"] for item in rainnet_metrics], label="RainNet")
+    axes[0].plot(leads, [item["mae_mm_h"] for item in persistence_metrics], label="Persistence")
+    axes[0].set(title="MAE", xlabel="Lead time (min)", ylabel="mm/h")
+    for threshold in thresholds:
+        axes[1].plot(leads, [item["csi"][str(threshold)] for item in rainnet_metrics], label=str(threshold))
+        axes[2].plot(leads, [item["fss"][str(threshold)]["20"] for item in rainnet_metrics], label=str(threshold))
+    axes[1].set(title="RainNet CSI", xlabel="Lead time (min)", ylabel="CSI")
+    axes[2].set(title="RainNet FSS (20 km)", xlabel="Lead time (min)", ylabel="FSS")
+    axes[0].legend()
+    axes[1].legend(title="mm/h", fontsize=7)
+    axes[2].legend(title="mm/h", fontsize=7)
+    fig.tight_layout()
+    fig.savefig(result_dir / "metrics.png", dpi=150)
+    plt.close(fig)
+
+    print(f"Prediction shape: {predictions.shape}")
+    print(f"Target shape: {targets.shape}")
+    print(f"Persistence shape: {persistence.shape}")
+    for index in selected:
+        item = rainnet_metrics[index]
+        print(f"Lead {item['lead_minutes']} min MAE: {item['mae_mm_h']:.8f} mm/h")
+        print(f"Lead {item['lead_minutes']} min CSI: {item['csi']}")
+        print(f"Lead {item['lead_minutes']} min FSS: {item['fss']}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,214 @@
+"""Train RainNet on contiguous windows from an RYDL-style HDF5 file."""
+
+import json
+import math
+import os
+from pathlib import Path
+import random
+import sys
+
+import h5py
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+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.rainnet import build_rainnet
+
+
+def load_config():
+    with (ROOT / "conf/config.yaml").open(encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def seed_everything(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def setup_device(config):
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if distributed:
+        torch.distributed.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo")
+    if torch.cuda.is_available() and config["device"] in ("auto", "cuda"):
+        device = torch.device("cuda", local_rank)
+        torch.cuda.set_device(device)
+    else:
+        device = torch.device("cpu")
+    return device, distributed, local_rank
+
+
+class RainNetDataset(Dataset):
+    def __init__(self, path, keys, input_steps=4):
+        self.path = path
+        self.keys = list(keys)
+        self.input_steps = input_steps
+        if len(self.keys) <= input_steps:
+            raise ValueError("A split needs at least input_steps + 1 frames")
+
+    def __len__(self):
+        return len(self.keys) - self.input_steps
+
+    def __getitem__(self, index):
+        with h5py.File(self.path, "r") as handle:
+            inputs = np.stack(
+                [handle[key][...] for key in self.keys[index : index + self.input_steps]]
+            )
+            target_key = self.keys[index + self.input_steps]
+            target = handle[target_key][...][None]
+        return torch.from_numpy(inputs), torch.from_numpy(target), target_key
+
+
+class LogCoshLoss(nn.Module):
+    def forward(self, prediction, target):
+        error = torch.abs(prediction - target)
+        return (error + F.softplus(-2.0 * error) - math.log(2.0)).mean()
+
+
+def transform_and_pad(tensor, pad):
+    return F.pad(torch.log(tensor + 0.01), pad, mode="reflect")
+
+
+def run_epoch(model, loader, criterion, device, pad, max_batches, optimizer=None):
+    training = optimizer is not None
+    model.train(training)
+    losses = []
+    parameter_updated = False
+    first_shapes = None
+    context = torch.enable_grad() if training else torch.no_grad()
+    with context:
+        for batch_index, (inputs, targets, target_keys) in enumerate(loader):
+            if batch_index >= max_batches:
+                break
+            inputs, targets = inputs.to(device), targets.to(device)
+            padded_inputs = transform_and_pad(inputs, pad)
+            padded_targets = transform_and_pad(targets, pad)
+            if training:
+                optimizer.zero_grad(set_to_none=True)
+            output = model(padded_inputs)
+            loss = criterion(output, padded_targets)
+            if not torch.isfinite(loss):
+                raise RuntimeError(f"Non-finite loss: {loss.item()}")
+            if first_shapes is None:
+                first_shapes = (inputs.shape, targets.shape, padded_inputs.shape, output.shape, target_keys[0])
+            if training:
+                tracked = next(model.parameters()).detach().clone()
+                loss.backward()
+                optimizer.step()
+                parameter_updated = parameter_updated or not torch.equal(tracked, next(model.parameters()).detach())
+            losses.append(loss.item())
+    return float(np.mean(losses)), parameter_updated, first_shapes
+
+
+def save_checkpoint(path, model, optimizer, epoch, val_loss, config):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    state_model = model.module if isinstance(model, DistributedDataParallel) else model
+    torch.save(
+        {
+            "model_state_dict": state_model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "epoch": epoch,
+            "validation_loss": val_loss,
+            "config": config,
+        },
+        path,
+    )
+
+
+def main():
+    config = load_config()
+    seed_everything(config["seed"])
+    device, distributed, local_rank = setup_device(config)
+    is_main = local_rank == 0
+    data = config["data"]
+    path = ROOT / data["path"]
+    if not path.exists():
+        raise FileNotFoundError(f"Fake data not found: {path}; run scripts/fake_data.py")
+    with h5py.File(path, "r") as handle:
+        keys = sorted(handle.keys())
+    train_end = data["train_frames"]
+    val_end = train_end + data["val_frames"]
+    train_set = RainNetDataset(path, keys[:train_end], data["input_steps"])
+    val_set = RainNetDataset(path, keys[train_end:val_end], data["input_steps"])
+    train_sampler = DistributedSampler(train_set, shuffle=True) if distributed else None
+    train_loader = DataLoader(
+        train_set,
+        batch_size=config["train"]["batch_size"],
+        shuffle=train_sampler is None,
+        sampler=train_sampler,
+        num_workers=data["num_workers"],
+    )
+    val_loader = DataLoader(val_set, batch_size=1, shuffle=False, num_workers=data["num_workers"])
+    model = build_rainnet(**config["model"]).to(device)
+    parameter_count = sum(parameter.numel() for parameter in model.parameters())
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    criterion = LogCoshLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=config["train"]["learning_rate"])
+    pad_h = data["padded_height"] - data["raw_height"]
+    pad_w = data["padded_width"] - data["raw_width"]
+    pad = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
+    history = {"train_loss": [], "validation_loss": [], "learning_rate": []}
+    best_loss = float("inf")
+    any_update = False
+    for epoch in range(config["train"]["epochs"]):
+        if train_sampler:
+            train_sampler.set_epoch(epoch)
+        train_loss, updated, shapes = run_epoch(
+            model, train_loader, criterion, device, pad, config["train"]["max_train_batches"], optimizer
+        )
+        val_loss, _, _ = run_epoch(
+            model, val_loader, criterion, device, pad, config["train"]["max_valid_batches"]
+        )
+        any_update = any_update or updated
+        history["train_loss"].append(train_loss)
+        history["validation_loss"].append(val_loss)
+        history["learning_rate"].append(optimizer.param_groups[0]["lr"])
+        if is_main:
+            last_path = ROOT / config["train"]["checkpoint_last"]
+            best_path = ROOT / config["train"]["checkpoint_best"]
+            save_checkpoint(last_path, model, optimizer, epoch + 1, val_loss, config)
+            if val_loss < best_loss:
+                best_loss = val_loss
+                save_checkpoint(best_path, model, optimizer, epoch + 1, val_loss, config)
+            result_dir = ROOT / config["evaluation"]["result_dir"]
+            result_dir.mkdir(parents=True, exist_ok=True)
+            with (result_dir / "train_history.json").open("w", encoding="utf-8") as handle:
+                json.dump(history, handle, indent=2)
+            print(f"Device: {device}")
+            print(f"Input shape: {tuple(shapes[0])}")
+            print(f"Target shape: {tuple(shapes[1])}")
+            print(f"Target key (i+4): {shapes[4]}")
+            print(f"Padded input shape: {tuple(shapes[2])}")
+            print(f"Model output shape: {tuple(shapes[3])}")
+            print(f"Parameter count: {parameter_count}")
+            print(f"Epoch: {epoch + 1}")
+            print(f"Train loss: {train_loss:.8f}")
+            print(f"Validation loss: {val_loss:.8f}")
+            print(f"Learning rate: {optimizer.param_groups[0]['lr']}")
+            print(f"parameter_update_detected: {any_update}")
+            print(f"Checkpoint path: {best_path}")
+    if not any_update:
+        raise RuntimeError("No model parameter changed after optimizer.step()")
+    if is_main:
+        reload_model = build_rainnet(**config["model"])
+        checkpoint = torch.load(
+            ROOT / config["train"]["checkpoint_best"], map_location="cpu", weights_only=False
+        )
+        reload_model.load_state_dict(checkpoint["model_state_dict"])
+        print("checkpoint_reload_after_training: True")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()