Add engineering reproduction package

.gitattributes

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

README.md

@@ -0,0 +1,148 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Precipitation Nowcasting
+- Spatiotemporal Sequence Prediction
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">ConvLSTM</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+ConvLSTM is a convolutional recurrent neural network for spatiotemporal sequence prediction. It replaces the input-to-state and state-to-state transformations of LSTM with spatial convolutions, thereby modeling local spatial correlations while preserving long-term temporal dependencies.
+
+Paper: Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting  
+https://papers.nips.cc/paper_files/paper/2015/hash/07563a3fe3bbe7e3ba84431ad9d055af-Abstract.html
+
+# Model Description
+
+ConvLSTM was proposed by researchers at the Hong Kong University of Science and Technology and the Hong Kong Observatory. The model is trained and evaluated using the 97 days with the highest rainfall in the 2011 to 2013 Hong Kong weather radar data, as well as synthetic Moving-MNIST sequences. The model is suitable for precipitation nowcasting, video prediction, and general spatiotemporal sequence prediction tasks.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Precipitation nowcasting | Predict the next 15 time steps from 5 historical radar maps. |
+| Spatiotemporal sequence modeling | Jointly learn spatial and temporal correlations with a convolutional gating structure. |
+| Multi-step image prediction | Continuously generate future images through an Encoder-Forecaster structure. |
+| Local workflow validation | Use synthetic radar sequences to check training, inference, evaluation, visualization, and checkpoint workflows. |
+| 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/ConvLSTM --local-dir ./ConvLSTM
+cd ConvLSTM
+```
+
+### 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
+
+This repository uses a small number of synthetic samples to validate the engineering workflow. The synthetic data contains 20 consecutive single-channel `float32` radar echo images of `100×100` pixels at 6-minute intervals. The first 5 frames are used as input and the following 15 frames as prediction targets. This data is only used to validate ConvLSTM's spatiotemporal encoding, 15-step prediction, training, inference, and evaluation workflows; it does not represent the official radar data distribution or training scale.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+The default configuration retains two Encoder layers, two Forecaster layers, `3×3` convolutions, channel-wise peephole connections, and the full prediction length; it only reduces the number of samples, hidden channels, and training epochs.
+
+```text
+result/checkpoints/convlstm.pt
+result/training/metrics.json
+```
+
+### Training Weights
+
+This repository does not include synthetic or trained weights, and the paper does not provide a directly downloadable original Theano pretrained checkpoint.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the training checkpoint and generates 15 future radar echo predictions from 5 historical radar maps. The inference results contain the input sequence, ground-truth target, predicted sequence, and corresponding lead-time information in minutes.
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Following the paper, the evaluation converts radar echoes to rainfall rates using the Z-R relationship and computes Rainfall-MSE, CSI, FAR, POD, and Correlation. The results include per-step metrics for all 15 forecast lead times and overall summary metrics, and generate comparison plots of targets, predictions, and absolute errors for selected lead times. Results on synthetic data are only for validating the engineering workflow and do not represent metrics on the paper's actual radar data.
+
+```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 publicly available specifications in the ConvLSTM paper.
+
+Use of the code and data in this repository remains subject to the licenses and terms of use of their respective projects.

README_zh.md

@@ -0,0 +1,151 @@
+---
+license: Apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- 地球科学
+- 降水临近预报
+- 时空序列预测
+frameworks: PyTorch
+datasets:
+- Hong Kong weather radar echo dataset
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">ConvLSTM</span>
+  </strong>
+</p>
+
+# 模型介绍
+
+ConvLSTM 是面向时空序列预测的卷积循环神经网络，将 LSTM 的输入到状态和状态到状态变换替换为空间卷积，从而在保留长期时间依赖的同时建模局部空间相关性。
+
+论文：Convolutional LSTM Network: A Machine Learning Approach for Precipitation Nowcasting  
+https://papers.nips.cc/paper_files/paper/2015/hash/07563a3fe3bbe7e3ba84431ad9d055af-Abstract.html
+
+# 模型描述
+
+ConvLSTM 由香港科技大学与香港天文台的研究人员提出。模型使用 2011 至 2013 年香港天气雷达数据中降雨量最高的 97 天，以及 Moving-MNIST 合成序列进行训练和评估。模型适用于降水临近预报、视频预测和一般时空序列预测任务。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 降水临近预报 | 根据 5 张历史雷达图预测未来 15 个时间步。 |
+| 时空序列建模 | 使用卷积门控结构联合学习空间和时间相关性。 |
+| 多步图像预测 | 通过 Encoder-Forecaster 结构连续生成未来图像。 |
+| 本地工程验证 | 使用虚拟雷达序列检查训练、推理、评估、可视化和 checkpoint 流程。 |
+| 多卡训练 | 通过 `torchrun` 启动分布式数据并行训练。 |
+
+# 使用说明
+
+## 1.OneCode
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/ConvLSTM --local_dir ./ConvLSTM
+cd ConvLSTM
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 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
+```
+
+### 训练数据介绍
+
+本仓库使用少量虚拟样本验证工程流程，虚拟数据包含连续 20 帧、时间间隔为 6 分钟的 `100×100` 单通道 `float32` 雷达回波图，其中前 5 帧作为输入，后 15 帧作为预测目标。该数据仅用于验证 ConvLSTM 的时空编码、15 步预测、训练、推理和评估流程，不代表官方雷达数据分布与训练规模。
+
+```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
+```
+
+默认配置保持两层 Encoder、两层 Forecaster、`3×3` 卷积、逐通道 peephole 和完整预测长度，只缩小样本数量、隐藏通道和训练周期。
+
+```text
+result/checkpoints/convlstm.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+本仓库不内置虚拟权重或训练权重，论文未提供可直接下载的原始 Theano 预训练 checkpoint。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载训练 checkpoint，根据 5 张历史雷达图生成未来 15 张雷达回波预测。推理结果包含输入序列、真实目标、预测序列及对应的分钟时效信息。
+
+```text
+result/output/predictions.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估按照论文的 Z-R 关系将雷达回波转换为降雨率，并计算 Rainfall-MSE、CSI、FAR、POD 和 Correlation。结果同时包含 15 个预测时效的分步指标及整体汇总指标，并生成部分时效的目标、预测和绝对误差对比图。虚拟数据结果仅用于验证工程流程，不代表论文真实雷达数据指标。
+
+```text
+result/evaluation/metrics.json
+result/evaluation/comparison.png
+```
+
+# OneScience官方信息
+
+| 平台 | OneScience 主仓库 | Skills 仓库 |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# 引用与许可证
+
+本仓库为 ConvLSTM 论文公开规格的独立工程复现版本。
+
+本仓库代码和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,45 @@
+seed: 42
+data:
+  root: data
+  format_version: convlstm_hko_synthetic_engineering_v1
+  train_samples: 4
+  test_samples: 2
+  input_frames: 5
+  output_frames: 15
+  channels: 1
+  height: 100
+  width: 100
+  interval_minutes: 6
+  patch_size: 2
+  rainfall_threshold_mm_h: 0.5
+  zr_a: 118.239
+  zr_b: 1.5241
+  radar_db_min: 0.0
+  radar_db_max: 70.0
+model:
+  input_channels: 1
+  patch_size: 2
+  hidden_channels: [16, 16]
+  kernel_size: 3
+  output_frames: 15
+paper_model:
+  input_channels: 1
+  patch_size: 2
+  hidden_channels: [64, 64]
+  kernel_size: 3
+  output_frames: 15
+train:
+  epochs: 1
+  batch_size: 1
+  learning_rate: 0.001
+  rmsprop_alpha: 0.9
+  weight_decay: 0.0
+  gradient_clip_norm: 10.0
+  num_workers: 0
+runtime:
+  device: auto
+paths:
+  checkpoint: result/checkpoints/convlstm.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,42 @@
+{
+  "model_name": "ConvLSTM",
+  "model_type": "convlstm",
+  "architectures": ["ConvLSTM", "ConvLSTMCell"],
+  "framework": "PyTorch",
+  "domain": "atmosphere",
+  "task": "precipitation-nowcasting",
+  "implementation": {
+    "entry_point": "model/convlstm.py",
+    "scope": "peephole ConvLSTM encoder-forecaster for fifteen-step radar echo prediction",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "family": "two-layer ConvLSTM encoder and two-layer unconditional forecaster",
+    "input_channels": 1,
+    "patch_size": 2,
+    "patched_channels": 4,
+    "hidden_channels": [16, 16],
+    "kernel_size": 3,
+    "peephole": true,
+    "output_frames": 15,
+    "output_head": "1x1 convolution over concatenated forecaster states"
+  },
+  "data": {
+    "datasets": ["Hong Kong weather radar echo dataset"],
+    "protocol": "convlstm_hko_synthetic_engineering_v1",
+    "format": "NPZ",
+    "train_file": "data/train.npz",
+    "test_file": "data/test.npz",
+    "input_shape": ["N", 5, 1, 100, 100],
+    "target_shape": ["N", 15, 1, 100, 100],
+    "interval_minutes": 6,
+    "forecast_horizon_minutes": 90,
+    "rainfall_threshold_mm_h": 0.5,
+    "zr_parameters": {"a": 118.239, "b": 1.5241},
+    "required_metadata": ["format_version", "data_source"]
+  },
+  "configuration_sources": ["conf/config.yaml", "model/convlstm.py", "scripts/fake_data.py", "scripts/train.py", "scripts/inference.py", "scripts/result.py"]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "precipitation_nowcasting",
+  "model": "ConvLSTM",
+  "input_format": "BTCHW",
+  "protocol": "convlstm_encoder_forecaster",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/convlstm.py

@@ -0,0 +1,95 @@
+"""Peephole ConvLSTM encoder-forecaster for precipitation nowcasting."""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def patchify(sequence, patch_size):
+    batch, steps, channels, height, width = sequence.shape
+    flattened = sequence.flatten(0, 1)
+    patched = F.pixel_unshuffle(flattened, patch_size)
+    return patched.unflatten(0, (batch, steps))
+
+
+def unpatchify(sequence, patch_size):
+    batch, steps = sequence.shape[:2]
+    images = F.pixel_shuffle(sequence.flatten(0, 1), patch_size)
+    return images.unflatten(0, (batch, steps))
+
+
+class ConvLSTMCell(nn.Module):
+    def __init__(self, input_channels, hidden_channels, kernel_size):
+        super().__init__()
+        padding = kernel_size // 2
+        self.hidden_channels = hidden_channels
+        self.input_conv = None if input_channels == 0 else nn.Conv2d(
+            input_channels, 4 * hidden_channels, kernel_size, padding=padding
+        )
+        self.hidden_conv = nn.Conv2d(hidden_channels, 4 * hidden_channels, kernel_size,
+                                     padding=padding, bias=False)
+        self.peephole_input = nn.Parameter(torch.zeros(1, hidden_channels, 1, 1))
+        self.peephole_forget = nn.Parameter(torch.zeros(1, hidden_channels, 1, 1))
+        self.peephole_output = nn.Parameter(torch.zeros(1, hidden_channels, 1, 1))
+        self.bias = nn.Parameter(torch.zeros(1, 4 * hidden_channels, 1, 1))
+
+    def forward(self, values, state):
+        hidden, cell = state
+        gates = self.hidden_conv(hidden) + self.bias
+        if values is not None:
+            if self.input_conv is None:
+                raise ValueError("this ConvLSTM cell has no external input projection")
+            gates = gates + self.input_conv(values)
+        input_gate, forget_gate, candidate, output_gate = gates.chunk(4, dim=1)
+        input_gate = torch.sigmoid(input_gate + self.peephole_input * cell)
+        forget_gate = torch.sigmoid(forget_gate + self.peephole_forget * cell)
+        cell = forget_gate * cell + input_gate * torch.tanh(candidate)
+        output_gate = torch.sigmoid(output_gate + self.peephole_output * cell)
+        hidden = output_gate * torch.tanh(cell)
+        return hidden, cell
+
+
+class ConvLSTM(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.patch_size = int(config["patch_size"])
+        self.output_frames = int(config["output_frames"])
+        patch_channels = int(config["input_channels"]) * self.patch_size ** 2
+        hidden = [int(value) for value in config["hidden_channels"]]
+        kernel = int(config["kernel_size"])
+        self.encoder = nn.ModuleList([
+            ConvLSTMCell(patch_channels, hidden[0], kernel),
+            ConvLSTMCell(hidden[0], hidden[1], kernel),
+        ])
+        self.forecaster = nn.ModuleList([
+            ConvLSTMCell(0, hidden[0], kernel),
+            ConvLSTMCell(hidden[0], hidden[1], kernel),
+        ])
+        self.output = nn.Conv2d(sum(hidden), patch_channels, 1)
+
+    @staticmethod
+    def _zero_state(batch, channels, height, width, reference):
+        zeros = reference.new_zeros(batch, channels, height, width)
+        return zeros, zeros.clone()
+
+    def forward(self, sequence, return_states=False):
+        patched = patchify(sequence, self.patch_size)
+        batch, _, _, height, width = patched.shape
+        states = [self._zero_state(batch, cell.hidden_channels, height, width, sequence)
+                  for cell in self.encoder]
+        for step in range(patched.shape[1]):
+            values = patched[:, step]
+            for index, cell in enumerate(self.encoder):
+                states[index] = cell(values, states[index])
+                values = states[index][0]
+        forecast_states = [(hidden.clone(), cell.clone()) for hidden, cell in states]
+        predictions, traces = [], []
+        for _ in range(self.output_frames):
+            forecast_states[0] = self.forecaster[0](None, forecast_states[0])
+            forecast_states[1] = self.forecaster[1](forecast_states[0][0], forecast_states[1])
+            hidden = torch.cat((forecast_states[0][0], forecast_states[1][0]), dim=1)
+            predictions.append(self.output(hidden))
+            traces.append([state[0] for state in forecast_states])
+        logits = torch.stack(predictions, dim=1)
+        images = unpatchify(logits.sigmoid(), self.patch_size)
+        return (images, logits, traces) if return_states else (images, logits)

scripts/fake_data.py

@@ -0,0 +1,50 @@
+"""Generate 5-to-15 synthetic radar sequences at the paper's 100x100 size."""
+
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def make_split(path, count, config, seed):
+    rng = np.random.default_rng(seed)
+    data = config["data"]
+    total = int(data["input_frames"]) + int(data["output_frames"])
+    height, width = int(data["height"]), int(data["width"])
+    y, x = np.mgrid[-1:1:complex(height), -1:1:complex(width)].astype(np.float32)
+    sequences = np.empty((count, total, 1, height, width), np.float32)
+    for sample in range(count):
+        centers = rng.uniform(-0.55, 0.55, (3, 2))
+        velocities = rng.uniform(-0.035, 0.035, (3, 2))
+        amplitudes = rng.uniform(0.25, 0.95, 3)
+        scales = rng.uniform(0.10, 0.28, 3)
+        for step in range(total):
+            field = np.zeros((height, width), np.float32)
+            for storm in range(3):
+                cy, cx = centers[storm] + velocities[storm] * step
+                distance = ((x - cx) ** 2 + (y - cy) ** 2) / (2 * scales[storm] ** 2)
+                field += amplitudes[storm] * np.exp(-distance)
+            sequences[sample, step, 0] = np.clip(field + rng.normal(0, 0.01, field.shape), 0, 1)
+    split = int(data["input_frames"])
+    np.savez_compressed(path, format_version=np.asarray(data["format_version"]),
+                        data_source=np.asarray("synthetic_hko_radar_like"),
+                        inputs=sequences[:, :split], targets=sequences[:, split:])
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    output = ROOT / config["data"]["root"]
+    output.mkdir(parents=True, exist_ok=True)
+    for offset, (filename, count) in enumerate((("train.npz", config["data"]["train_samples"]),
+                                                 ("test.npz", config["data"]["test_samples"]))):
+        target = output / filename
+        if not target.exists():
+            make_split(target, int(count), config, int(config["seed"]) + offset)
+        print(f"generated={target.relative_to(ROOT)} input=5x1x100x100 target=15x1x100x100")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,46 @@
+"""Forecast fifteen 6-minute radar frames from five observations."""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.convlstm import ConvLSTM
+from train import RadarDataset, device_from_config
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    device = device_from_config(config)
+    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=True)
+    model = ConvLSTM(checkpoint["model_config"]).to(device)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    loader = DataLoader(RadarDataset(ROOT / config["data"]["root"] / "test.npz", config), batch_size=1)
+    inputs_all, targets_all, predictions_all = [], [], []
+    with torch.no_grad():
+        for inputs, targets in loader:
+            prediction, _ = model(inputs.to(device))
+            inputs_all.append(inputs.numpy())
+            targets_all.append(targets.numpy())
+            predictions_all.append(prediction.cpu().numpy())
+    output = ROOT / config["paths"]["inference_dir"] / "predictions.npz"
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, inputs=np.concatenate(inputs_all), targets=np.concatenate(targets_all),
+                        predictions=np.concatenate(predictions_all),
+                        input_lead_minutes=np.arange(-24, 1, int(config["data"]["interval_minutes"]), dtype=np.int64),
+                        forecast_lead_minutes=np.arange(1, int(config["data"]["output_frames"]) + 1, dtype=np.int64)
+                                              * int(config["data"]["interval_minutes"]),
+                        normalized_value_range=np.asarray([0.0, 1.0], np.float32),
+                        data_type=np.asarray("normalized_radar_echo_grayscale"))
+    print(f"predictions={output.relative_to(ROOT)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,78 @@
+"""Compute the paper's rainfall metrics and visualize the 90-minute forecast."""
+
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def to_rainfall(values, config):
+    data = config["data"]
+    radar_db = values * (float(data["radar_db_max"]) - float(data["radar_db_min"])) + float(data["radar_db_min"])
+    return 10 ** ((radar_db - 10 * np.log10(float(data["zr_a"]))) / (10 * float(data["zr_b"])))
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
+    prediction, target = data["predictions"], data["targets"]
+    predicted_rain, target_rain = to_rainfall(prediction, config), to_rainfall(target, config)
+    threshold = float(config["data"]["rainfall_threshold_mm_h"])
+    forecast, observed = predicted_rain >= threshold, target_rain >= threshold
+    hits, misses = np.logical_and(forecast, observed).sum(), np.logical_and(~forecast, observed).sum()
+    false_alarms = np.logical_and(forecast, ~observed).sum()
+    eps = 1e-9
+    frame_correlation, per_step = [], []
+    for step in range(prediction.shape[1]):
+        left, right = prediction[:, step].reshape(-1), target[:, step].reshape(-1)
+        correlation = float(np.dot(left, right) / (np.sqrt(np.dot(left, left) * np.dot(right, right)) + eps))
+        frame_correlation.append(correlation)
+        step_forecast, step_observed = forecast[:, step], observed[:, step]
+        step_hits = np.logical_and(step_forecast, step_observed).sum()
+        step_misses = np.logical_and(~step_forecast, step_observed).sum()
+        step_false_alarms = np.logical_and(step_forecast, ~step_observed).sum()
+        per_step.append({
+            "lead_minutes": int(data["forecast_lead_minutes"][step]),
+            "rainfall_mse": float(np.mean((predicted_rain[:, step] - target_rain[:, step]) ** 2)),
+            "csi": float(step_hits / (step_hits + step_misses + step_false_alarms + eps)),
+            "far": float(step_false_alarms / (step_hits + step_false_alarms + eps)),
+            "pod": float(step_hits / (step_hits + step_misses + eps)),
+            "correlation": correlation,
+        })
+    metrics = {
+        "samples": int(len(prediction)),
+        "binary_cross_entropy": float(-(target * np.log(prediction.clip(1e-7, 1 - 1e-7)) +
+                                         (1 - target) * np.log((1 - prediction).clip(1e-7, 1))).mean()),
+        "rainfall_mse": float(np.mean((predicted_rain - target_rain) ** 2)),
+        "csi": float(hits / (hits + misses + false_alarms + eps)),
+        "far": float(false_alarms / (hits + false_alarms + eps)),
+        "pod": float(hits / (hits + misses + eps)),
+        "correlation": float(np.mean(frame_correlation)),
+        "per_forecast_step": 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")
+    steps = [0, 2, 5, 8, 11, 14]
+    figure, axes = plt.subplots(3, len(steps), figsize=(15, 7))
+    for column, step in enumerate(steps):
+        axes[0, column].imshow(target[0, step, 0], cmap="turbo", vmin=0, vmax=1)
+        axes[1, column].imshow(prediction[0, step, 0], cmap="turbo", vmin=0, vmax=1)
+        axes[2, column].imshow(np.abs(target[0, step, 0] - prediction[0, step, 0]), cmap="magma", vmin=0, vmax=1)
+        axes[0, column].set_title(f"+{(step + 1) * 6} min")
+        for axis in axes[:, column]:
+            axis.axis("off")
+    figure.tight_layout()
+    figure.savefig(output / "comparison.png", dpi=150)
+    plt.close(figure)
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,95 @@
+"""Train the ConvLSTM radar encoder-forecaster with full-sequence BPTT."""
+
+import json
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.convlstm import ConvLSTM
+
+
+class RadarDataset(Dataset):
+    def __init__(self, path, config):
+        self.data = np.load(path)
+        data = config["data"]
+        if str(self.data["format_version"]) != data["format_version"]:
+            raise ValueError("incompatible radar data format")
+        expected_input = (int(data["input_frames"]), int(data["channels"]), int(data["height"]), int(data["width"]))
+        expected_target = (int(data["output_frames"]), int(data["channels"]), int(data["height"]), int(data["width"]))
+        if self.data["inputs"].shape[1:] != expected_input or self.data["targets"].shape[1:] != expected_target:
+            raise ValueError("radar tensors do not preserve the paper dimensions")
+
+    def __len__(self):
+        return len(self.data["inputs"])
+
+    def __getitem__(self, index):
+        return torch.from_numpy(self.data["inputs"][index]).float(), torch.from_numpy(self.data["targets"][index]).float()
+
+
+def device_from_config(config, rank=0):
+    if config["runtime"]["device"] == "auto":
+        return torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu")
+    return torch.device(config["runtime"]["device"])
+
+
+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"))
+    if distributed:
+        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    rank = torch.distributed.get_rank() if distributed else 0
+    device = device_from_config(config, local_rank)
+    dataset = RadarDataset(ROOT / config["data"]["root"] / "train.npz", config)
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), sampler=sampler,
+                        shuffle=sampler is None, num_workers=int(config["train"]["num_workers"]))
+    model = ConvLSTM(config["model"]).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.RMSprop(model.parameters(), lr=float(config["train"]["learning_rate"]),
+                                    alpha=float(config["train"]["rmsprop_alpha"]),
+                                    weight_decay=float(config["train"]["weight_decay"]))
+    history = []
+    for epoch in range(int(config["train"]["epochs"])):
+        model.train()
+        total, steps = 0.0, 0
+        for inputs, targets in loader:
+            _, logits = model(inputs.to(device))
+            patched_target = torch.nn.functional.pixel_unshuffle(targets.to(device).flatten(0, 1),
+                                                                  int(config["model"]["patch_size"])).unflatten(0, targets.shape[:2])
+            loss = torch.nn.functional.binary_cross_entropy_with_logits(logits, patched_target)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), float(config["train"]["gradient_clip_norm"]))
+            optimizer.step()
+            total += float(loss.detach())
+            steps += 1
+        metrics = {"epoch": epoch + 1, "binary_cross_entropy": total / max(steps, 1)}
+        history.append(metrics)
+        if rank == 0:
+            print(f"epoch={epoch + 1} binary_cross_entropy={metrics['binary_cross_entropy']:.6f}")
+    if rank == 0:
+        checkpoint, metrics_path = ROOT / config["paths"]["checkpoint"], ROOT / config["paths"]["training_metrics"]
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        state = model.module.state_dict() if distributed else model.state_dict()
+        torch.save({"model": state, "model_config": config["model"],
+                    "format_version": config["data"]["format_version"]}, checkpoint)
+        metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()