Publish MetNet-2 reproduction

.gitattributes

@@ -1,35 +1,3 @@
-*.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
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,148 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- earth-science
+- precipitation-forecasting
+- MetNet-2
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">MetNet-2</span></strong>
+</p>
+
+# Model Introduction
+
+MetNet-2 generates high-resolution probabilistic precipitation forecasts up to 12 hours ahead from radar, satellite, and atmospheric-state inputs for short-range forecasting and extreme-precipitation risk analysis.
+
+Paper: Deep learning for twelve hour precipitation forecasts  
+https://doi.org/10.1038/s41467-022-32483-x
+
+# Model Description
+
+The method was proposed by the Google Research team. The paper constructs its 2017-2020 training and test data from MRMS radar, GOES satellite imagery, and HRRR assimilated atmospheric states. The model predicts grid-cell precipitation probability distributions at two-minute intervals up to 12 hours ahead.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Categorical precipitation forecasting | Validate the 512-category conditional distribution and 12-hour lead-time protocol. |
+| Core-method validation | Validate ConvLSTM, lead-time FiLM, and multiscale dilated residual stacks. |
+| Local engineering validation | Exercise the complete logical `641×512×512` contract through a deterministic procedural field without materializing the full input. |
+| ModelScope/OneCode execution | Validate structured data, training, inference, probabilistic precipitation metrics, and visualization in ModelScope or OneCode environments. |
+| Multi-GPU training | Validate distributed training and the checkpoint workflow through `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/MetNet-2 --local-dir ./MetNet-2
+cd MetNet-2
+```
+
+### Environment Dependencies
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for connectivity validation with the default small-sample configuration.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+The paper training inputs combine MRMS, HRRR, GOES, static geography, and temporal information under a 641-channel, `512×512` spatial-domain protocol. This repository creates eight deterministic window records and constructs only the selected `32×32` windows plus their halos at runtime while retaining the complete logical shape and channel grouping. The synthetic data validate engineering connectivity only and do not represent real meteorological distributions, paper-scale training, or paper performance.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+For single-GPU training, use:
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+The default reduces sample count, network width, residual blocks, and training steps without reducing the 641 channels, 512 precipitation categories, or 12-hour lead protocol. Formal experiments require real MRMS, GOES, and HRRR data and full computing resources, with artifacts saved to:
+
+```text
+result/checkpoints/metnet_2.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+The paper does not provide directly loadable official model weights, and this repository bundles no weights under `weight/`. The locally trained checkpoint is saved to `result/checkpoints/metnet_2.pt` and must not be represented as an official pretrained weight.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the training checkpoint and generates 512-category precipitation probabilities and their CDF for selected spatial windows while preserving coverage and completeness metadata. Complete numerical results are saved to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation computes discrete CRPS, Brier Score, and CSI at multiple precipitation thresholds and generates target, expected-rate, and error comparisons. Synthetic-data results validate engineering only and do not represent paper performance; outputs are saved to:
+
+```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 MetNet-2 specifications.
+
+Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.

README_zh.md

@@ -0,0 +1,149 @@
+---
+license: apache-2.0
+language:
+- zh
+- en
+tags:
+- OneScience
+- 地球科学
+- 降水预报
+- MetNet-2
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">MetNet-2</span></strong>
+</p>
+
+# 模型介绍
+
+MetNet-2 根据雷达、卫星和大气状态生成未来 12 小时的高分辨率概率降水预报，主要用于短时降水预测和极端降水风险分析。
+
+论文：Deep learning for twelve hour precipitation forecasts  
+https://doi.org/10.1038/s41467-022-32483-x
+
+# 模型描述
+
+该方法由 Google Research 团队提出。论文使用 MRMS 雷达、GOES 卫星和 HRRR 同化大气状态构建 2017-2020 年训练与测试数据。模型适用于以 2 分钟间隔预测未来 12 小时的逐网格降水概率分布。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 降水分类预报 | 验证 512 类条件降水分布及 12 小时时效协议。 |
+| 核心方法验证 | 验证 ConvLSTM、lead-time FiLM 和多尺度膨胀残差栈。 |
+| 本地工程验证 | 使用确定性 procedural field 验证完整 `641×512×512` 逻辑协议，不物化完整输入。 |
+| ModelScope/OneCode 运行 | 在 ModelScope 或 OneCode 环境中验证结构化数据、训练、推理、概率降水指标和可视化流程。 |
+| 多卡训练 | 通过 `torchrun` 验证分布式训练和 checkpoint 流程。 |
+
+# 使用说明
+
+## 1.OneCode
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/MetNet-2 --local_dir ./MetNet-2
+cd MetNet-2
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 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
+```
+
+### 训练数据介绍
+
+论文训练数据由 MRMS、HRRR、GOES、静态地理和时间信息组成，并遵循 641 通道、`512×512` 空间域协议。本仓库生成 8 条确定性窗口记录，运行时仅构造选定的 `32×32` 窗口及 halo，同时保留完整逻辑形状和通道分组。合成数据只验证工程连通性，不代表真实气象数据分布、论文训练规模或论文性能。
+
+```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
+```
+
+默认配置仅缩小样本数、网络宽度、残差块数量和训练步数，不缩小 641 通道、512 个降水类别和 12 小时时效协议。正式实验需要真实 MRMS、GOES 和 HRRR 数据及完整计算资源，训练产物保存到：
+
+```text
+result/checkpoints/metnet_2.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+论文未提供可直接加载的官方模型权重，本仓库不在 `weight/` 中内置权重。本地训练生成的 checkpoint 保存到 `result/checkpoints/metnet_2.pt`，不得描述为官方预训练权重。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载训练 checkpoint，为选定空间窗口生成 512 类降水概率及对应 CDF，并明确保存覆盖范围和完整性标记。完整数值结果保存到：
+
+```text
+result/output/predictions.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估计算离散 CRPS、Brier Score 和不同降水阈值下的 CSI，并生成目标、期望降水率和误差对比图。虚拟数据结果仅用于验证工程流程，不代表论文正式性能；结果保存到：
+
+```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 |
+
+# 引用与许可证
+
+本仓库为 MetNet-2 公开规格的独立工程复现版本。
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,42 @@
+seed: 2026
+data:
+  path: data/metnet2_fake.npz
+  format_version: metnet2_selected_windows_v1
+  logical_shape: [641, 512, 512]
+  channel_groups: {mrms_radar_history: 33, hrrr_atmosphere_history: 484, goes_satellite_history: 96, static_geography: 24, time_coordinates: 4}
+  window: 32
+  halo: 8
+  samples: 8
+model:
+  input_channels: 641
+  classes: 512
+  width: 8
+  stacks: 1
+  dilations: [1, 2, 4, 8, 16, 32, 64, 128]
+  lead_max_minutes: 720
+training:
+  batch_size: 1
+  steps: 1
+  learning_rate: 0.001
+  class_chunk: 64
+runtime:
+  device: auto
+  num_threads: 2
+inference:
+  lead_minutes: 2
+  coverage: selected 32x32 window
+  is_complete: false
+paper_model:
+  input_channels: 641
+  classes: 512
+  stacks: 3
+  dilations: [1, 2, 4, 8, 16, 32, 64, 128]
+  lead_times_minutes: {start: 2, stop: 720, step: 2}
+  spatial_upscale: 4
+  methods: [ConvLSTM, lead-time FiLM, dilated residual stacks, center crop, bilinear upsampling]
+paths:
+  checkpoint: result/checkpoints/metnet_2.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation_metrics: result/evaluation/metrics.json
+  comparison: result/evaluation/comparison.png

config.json

@@ -0,0 +1,41 @@
+{
+  "model_name": "MetNet-2",
+  "model_type": "metnet_2",
+  "architectures": ["MetNet2"],
+  "framework": "PyTorch",
+  "domain": "weather",
+  "task": "probabilistic-precipitation-forecasting",
+  "implementation": {
+    "entry_point": "model/metnet_2.py",
+    "scope": "core-method and logical full-dimension sampled-window engineering reproduction",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "logical_input_shape": ["B", 641, 512, 512],
+    "logical_output_shape": ["B", 512, 512, 512],
+    "engineering_window": [32, 32],
+    "classes": 512,
+    "lead_minutes": [2, 720, 2],
+    "core": ["ConvLSTM", "lead-time FiLM", "dilated residual blocks", "spatial and class chunking"]
+  },
+  "data": {
+    "datasets": ["MRMS", "HRRR", "GOES"],
+    "format_version": "metnet2_selected_windows_v1",
+    "input_channels": 641,
+    "precipitation_range_mm_h": [0.0, 102.4],
+    "coverage": "selected 32x32 target windows",
+    "is_complete_global": false,
+    "synthetic": true
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/metnet_2.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "probabilistic precipitation forecasting",
+  "model": "MetNet-2",
+  "input_format": "logical BCHW: [B,641,512,512] with sampled windows for engineering execution",
+  "protocol": "lead-conditioned categorical precipitation forecasting from 2 to 720 minutes",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/metnet_2.py

@@ -0,0 +1,298 @@
+"""Shape-faithful, memory-bounded MetNet-2 engineering implementation."""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset
+import yaml
+
+CHANNEL_GROUPS = (
+    ("mrms_radar_history", 33), ("hrrr_atmosphere_history", 484),
+    ("goes_satellite_history", 96), ("static_geography", 24),
+    ("time_coordinates", 4),
+)
+LOGICAL_SHAPE = (641, 512, 512)
+CLASS_RATES = np.linspace(0.0, 102.4, 512, dtype=np.float32)
+assert sum(size for _, size in CHANNEL_GROUPS) == LOGICAL_SHAPE[0]
+
+
+def load_config(path: str | Path = "conf/config.yaml") -> dict:
+    with Path(path).open(encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+class ProceduralField:
+    """Generate crops of a logical [641, 512, 512] field without materializing it."""
+
+    shape = LOGICAL_SHAPE
+
+    def __init__(self, seed: int):
+        self.seed = int(seed)
+
+    def window(self, y: int, x: int, size: int, halo: int = 0) -> Tensor:
+        if size <= 0 or halo < 0 or not (0 <= y < 512 and 0 <= x < 512):
+            raise ValueError("invalid selected-window coordinates")
+        yy = torch.arange(y - halo, y + size + halo).clamp(0, 511).float()
+        xx = torch.arange(x - halo, x + size + halo).clamp(0, 511).float()
+        channels = torch.arange(641).float()[:, None, None]
+        return (torch.sin((channels + self.seed) * .017 + yy[None, :, None] * .031)
+                + torch.cos((channels + 3 * self.seed) * .011 + xx[None, None, :] * .023)).float()
+
+    def target_window(self, y: int, x: int, size: int, lead: int) -> Tensor:
+        yy = torch.arange(y, y + size)[:, None]
+        xx = torch.arange(x, x + size)[None, :]
+        return ((yy * 7 + xx * 11 + self.seed + lead // 2) % 512).long()
+
+
+class WindowDataset(Dataset):
+    def __init__(self, data_path: str | Path, split: str = "train"):
+        with np.load(data_path) as data:
+            required = ("seed", "split", "y", "x", "size", "halo", "lead_minutes")
+            missing = set(required).difference(data.files)
+            if missing:
+                raise ValueError(f"dataset is missing fields: {sorted(missing)}")
+            indices = np.flatnonzero(data["split"].astype(str) == split)
+            self.records = [{key: data[key][i].item() for key in required} for i in indices]
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(self, index: int) -> tuple[Tensor, Tensor, Tensor]:
+        record = self.records[index]
+        field = ProceduralField(record["seed"])
+        args = record["y"], record["x"], record["size"]
+        return (field.window(*args, record["halo"]),
+                field.target_window(*args, record["lead_minutes"]),
+                torch.tensor(record["lead_minutes"], dtype=torch.long))
+
+
+def write_fake_data(path: str | Path, samples: int = 8, window: int = 32, halo: int = 8) -> Path:
+    if samples < 3 or window != 32 or window + 2 * halo > 512:
+        raise ValueError("fake data requires at least 3 selected 32x32 windows with a valid halo")
+    records = []
+    for i in range(samples):
+        records.append({
+            "id": f"sample-{i:04d}", "seed": 1000 + i,
+            "split": "train" if i < samples - 2 else "test",
+            "y": (i * 47) % (512 - window + 1), "x": (i * 83) % (512 - window + 1),
+            "size": window, "halo": halo, "lead_minutes": 2 + 2 * (i % 360),
+        })
+    output = Path(path)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, **{key: np.asarray([r[key] for r in records]) for key in records[0]})
+    return output
+
+
+class LeadFiLMConv(nn.Module):
+    def __init__(self, cin: int, cout: int, dilation: int = 1):
+        super().__init__()
+        self.conv = nn.Conv2d(cin, cout, 3, padding=dilation, dilation=dilation)
+        self.film = nn.Linear(cout, 2 * cout)
+
+    def forward(self, x: Tensor, lead: Tensor) -> Tensor:
+        result = self.conv(x)
+        add, multiply = self.film(lead).chunk(2, dim=1)
+        return result * (1.0 + torch.tanh(multiply)[:, :, None, None]) + add[:, :, None, None]
+
+
+class ConvLSTMCell(nn.Module):
+    def __init__(self, cin: int, hidden: int):
+        super().__init__()
+        self.hidden = hidden
+        self.gates = nn.Conv2d(cin + hidden, 4 * hidden, 3, padding=1)
+
+    def forward(self, x: Tensor, state: tuple[Tensor, Tensor] | None = None) -> tuple[Tensor, Tensor]:
+        if state is None:
+            shape = (x.shape[0], self.hidden, x.shape[-2], x.shape[-1])
+            state = x.new_zeros(shape), x.new_zeros(shape)
+        hidden, cell = state
+        in_gate, forget, candidate, out_gate = self.gates(torch.cat((x, hidden), 1)).chunk(4, 1)
+        cell = torch.sigmoid(forget) * cell + torch.sigmoid(in_gate) * torch.tanh(candidate)
+        return torch.sigmoid(out_gate) * torch.tanh(cell), cell
+
+
+class DilatedResidualBlock(nn.Module):
+    def __init__(self, width: int, dilation: int):
+        super().__init__()
+        self.conv1 = LeadFiLMConv(width, width, dilation)
+        self.conv2 = LeadFiLMConv(width, width, dilation)
+
+    def forward(self, x: Tensor, lead: Tensor) -> Tensor:
+        return x + self.conv2(F.relu(self.conv1(F.relu(x), lead)), lead)
+
+
+class MetNet2(nn.Module):
+    """MetNet-2 concept model retaining the 641-channel and 512-class contracts."""
+
+    def __init__(self, input_channels: int = 641, classes: int = 512, width: int = 8,
+                 stacks: int = 1, dilations: Iterable[int] = (1, 2, 4, 8, 16, 32, 64, 128),
+                 lead_max_minutes: int = 720):
+        super().__init__()
+        if input_channels != 641 or classes != 512:
+            raise ValueError("MetNet-2 requires 641 input channels and 512 output classes")
+        self.input_channels, self.classes = input_channels, classes
+        self.width, self.stacks = width, stacks
+        self.dilations = tuple(dilations)
+        self.lead_max_minutes, self.upscale = lead_max_minutes, 4
+        self.lead_embedding = nn.Sequential(nn.Linear(1, width), nn.SiLU(), nn.Linear(width, width))
+        self.input_projection = nn.Conv2d(input_channels, width, 1)
+        self.temporal = ConvLSTMCell(width, width)
+        self.blocks = nn.ModuleList(DilatedResidualBlock(width, dilation)
+                                    for _ in range(stacks) for dilation in self.dilations)
+        self.spatial = LeadFiLMConv(width, width)
+        self.head = nn.Conv2d(width, classes, 1)
+
+    def _lead(self, minutes: Tensor) -> Tensor:
+        if torch.any((minutes < 2) | (minutes > self.lead_max_minutes) | (minutes % 2 != 0)):
+            raise ValueError("lead time must be 2..720 minutes in 2-minute increments")
+        return self.lead_embedding((minutes.float() / self.lead_max_minutes).unsqueeze(1))
+
+    def _features(self, x: Tensor, lead_minutes: Tensor, output_size: int) -> Tensor:
+        if x.ndim != 4 or x.shape[1] != 641:
+            raise ValueError("x must have shape [B, 641, H, W]")
+        if output_size <= 0 or output_size % self.upscale:
+            raise ValueError("output_size must be positive and divisible by four")
+        lead = self._lead(lead_minutes.to(x.device))
+        features, _ = self.temporal(self.input_projection(x))
+        for block in self.blocks:
+            features = block(features, lead)
+        features = self.spatial(F.relu(features), lead)
+        crop = output_size // self.upscale
+        if min(features.shape[-2:]) < crop:
+            raise ValueError("input window is smaller than the requested output")
+        top, left = (features.shape[-2] - crop) // 2, (features.shape[-1] - crop) // 2
+        return F.interpolate(features[:, :, top:top + crop, left:left + crop], size=(output_size, output_size),
+                             mode="bilinear", align_corners=False)
+
+    def forward_window(self, x: Tensor, lead_minutes: Tensor, output_size: int = 32,
+                       class_slice: tuple[int, int] | None = None) -> Tensor:
+        features = self._features(x, lead_minutes, output_size)
+        start, end = class_slice or (0, self.classes)
+        if not (0 <= start < end <= self.classes):
+            raise ValueError("invalid class slice")
+        return F.conv2d(features, self.head.weight[start:end], self.head.bias[start:end])
+
+    def forward(self, x: Tensor, lead_minutes: Tensor, output_size: int = 32) -> Tensor:
+        return self.forward_window(x, lead_minutes, output_size)
+
+    @torch.no_grad()
+    def assemble_full(self, source: ProceduralField, lead_minutes: int, output_path: str | Path,
+                      tile: int = 32, halo: int = 8, class_chunk: int = 64,
+                      output: str = "probability", device: str | torch.device = "cpu") -> Path:
+        """Stream a complete [512, 512, 512] probability or CDF array to disk."""
+        if output not in {"probability", "cdf"}:
+            raise ValueError("output must be probability or cdf")
+        path = Path(output_path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        array = np.lib.format.open_memmap(path, mode="w+", dtype=np.float16, shape=(512, 512, 512))
+        self.eval().to(device)
+        lead = torch.tensor([lead_minutes], device=device)
+        for y in range(0, 512, tile):
+            for x0 in range(0, 512, tile):
+                size = min(tile, 512 - y, 512 - x0)
+                features = self._features(source.window(y, x0, size, halo).unsqueeze(0).to(device), lead, size)[0]
+                maximum = None
+                for start in range(0, 512, class_chunk):
+                    logits = F.conv2d(features.unsqueeze(0), self.head.weight[start:start + class_chunk],
+                                      self.head.bias[start:start + class_chunk])[0]
+                    value = logits.amax(0)
+                    maximum = value if maximum is None else torch.maximum(maximum, value)
+                denominator = torch.zeros_like(maximum)
+                chunks = []
+                for start in range(0, 512, class_chunk):
+                    logits = F.conv2d(features.unsqueeze(0), self.head.weight[start:start + class_chunk],
+                                      self.head.bias[start:start + class_chunk])[0]
+                    exponent = torch.exp(logits - maximum)
+                    denominator += exponent.sum(0)
+                    chunks.append(exponent)
+                cumulative = torch.zeros_like(maximum)
+                for start, exponent in zip(range(0, 512, class_chunk), chunks):
+                    values = exponent / denominator
+                    if output == "cdf":
+                        values = values.cumsum(0) + cumulative
+                        cumulative = values[-1]
+                    array[start:start + values.shape[0], y:y + size, x0:x0 + size] = values.cpu().numpy()
+        array.flush()
+        return path
+
+
+def build_model(config: dict, paper: bool = False) -> MetNet2:
+    values = dict(config["model"])
+    if paper:
+        values.update({key: value for key, value in config["paper_model"].items()
+                       if key in {"input_channels", "classes", "stacks", "dilations"}})
+    dilations = tuple(values.get("dilations", ()))
+    if dilations != (1, 2, 4, 8, 16, 32, 64, 128):
+        raise ValueError("each dilation stack must use rates 1,2,4,8,16,32,64,128")
+    if paper and values["stacks"] != 3:
+        raise ValueError("the paper model requires three dilation stacks")
+    return MetNet2(**values)
+
+
+def categorical_nll_chunked(model: MetNet2, x: Tensor, lead: Tensor, target: Tensor,
+                            output_size: int = 32, class_chunk: int = 64) -> Tensor:
+    """Compute exact categorical NLL while applying the class head in chunks."""
+    features = model._features(x, lead, output_size)
+    selected, logsumexp = torch.zeros_like(target, dtype=features.dtype), None
+    for start in range(0, model.classes, class_chunk):
+        end = min(start + class_chunk, model.classes)
+        logits = F.conv2d(features, model.head.weight[start:end], model.head.bias[start:end])
+        part = torch.logsumexp(logits, dim=1)
+        logsumexp = part if logsumexp is None else torch.logaddexp(logsumexp, part)
+        mask = (target >= start) & (target < end)
+        picked = logits.gather(1, (target - start).clamp(0, end - start - 1).unsqueeze(1)).squeeze(1)
+        selected = torch.where(mask, picked, selected)
+    return (logsumexp - selected).mean()
+
+
+def save_checkpoint(path: str | Path, model: nn.Module, model_config: dict) -> None:
+    if int(os.environ.get("RANK", "0")) != 0:
+        return
+    module = model.module if hasattr(model, "module") else model
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    temporary = Path(f"{destination}.tmp")
+    torch.save({"model": module.state_dict(), "model_config": model_config,
+                "format_version": "metnet_2_v1"}, temporary)
+    os.replace(temporary, destination)
+
+
+def load_checkpoint(path: str | Path, model: nn.Module) -> dict:
+    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
+    if set(checkpoint) != {"model", "model_config", "format_version"}:
+        raise ValueError("checkpoint must contain model, model_config, and format_version")
+    model.load_state_dict(checkpoint["model"])
+    return checkpoint
+
+
+def scores(probabilities: np.ndarray, target: np.ndarray,
+           thresholds: tuple[float, ...] = (.2, 1., 2., 4., 8.)) -> dict:
+    if probabilities.shape[0] != 512 or target.shape != probabilities.shape[1:]:
+        raise ValueError("expected probabilities [512,H,W] and target [H,W]")
+    cdf = np.cumsum(probabilities.astype(np.float32), axis=0)
+    observed_cdf = (np.arange(512)[:, None, None] >= target[None]).astype(np.float32)
+    result = {"discrete_crps": float(np.mean(np.sum((cdf - observed_cdf) ** 2, axis=0)))}
+    brier, csi = {}, {}
+    for threshold in thresholds:
+        index = min(511, int(round(threshold / .2)))
+        event_probability = 1.0 - cdf[index - 1] if index else np.ones_like(cdf[0])
+        observed, forecast = target >= index, event_probability >= .5
+        hits = np.logical_and(forecast, observed).sum()
+        denominator = hits + np.logical_and(forecast, ~observed).sum() + np.logical_and(~forecast, observed).sum()
+        brier[str(threshold)] = float(np.mean((event_probability - observed) ** 2))
+        csi[str(threshold)] = float(hits / denominator) if denominator else 1.0
+    result.update(brier=brier, csi=csi)
+    return result
+
+
+def write_json(path: str | Path, value: dict) -> None:
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    destination.write_text(json.dumps(value, indent=2) + "\n", encoding="utf-8")

scripts/fake_data.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+import sys
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.metnet_2 import load_config, write_fake_data
+
+parser = argparse.ArgumentParser(description="Create selected MetNet-2 validation windows")
+parser.add_argument("--config", default="conf/config.yaml")
+parser.add_argument("--output", default=None)
+args = parser.parse_args()
+config = load_config(args.config)
+output = args.output or config["data"]["path"]
+print(write_fake_data(output, config["data"]["samples"], config["data"]["window"], config["data"]["halo"]))

scripts/inference.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+from model.metnet_2 import CLASS_RATES, ProceduralField, build_model, load_checkpoint, load_config
+
+parser = argparse.ArgumentParser(description="Run selected-window or streamed full-domain inference")
+parser.add_argument("--config", default="conf/config.yaml")
+parser.add_argument("--lead", type=int, default=None)
+parser.add_argument("--full", action="store_true")
+parser.add_argument("--cdf", action="store_true")
+args = parser.parse_args()
+config = load_config(args.config)
+torch.set_num_threads(config["runtime"]["num_threads"])
+device = torch.device("cuda" if config["runtime"]["device"] == "auto" and torch.cuda.is_available()
+                      else "cpu" if config["runtime"]["device"] == "auto" else config["runtime"]["device"])
+model = build_model(config).to(device)
+load_checkpoint(config["paths"]["checkpoint"], model)
+field, lead = ProceduralField(2001), args.lead or config["inference"]["lead_minutes"]
+if args.full:
+    output = Path(config["paths"]["predictions"]).with_suffix(".npy")
+    print(model.assemble_full(field, lead, output, config["data"]["window"], config["data"]["halo"],
+                              config["training"]["class_chunk"], "cdf" if args.cdf else "probability", device))
+else:
+    window = config["data"]["window"]
+    model.eval()
+    with torch.no_grad():
+        logits = model(field.window(0, 0, window, config["data"]["halo"]).unsqueeze(0).to(device),
+                       torch.tensor([lead], device=device), window)[0]
+        probabilities = logits.softmax(0).cpu().numpy().astype(np.float32)
+    if not np.isfinite(probabilities).all():
+        raise FloatingPointError("inference probabilities are not finite")
+    output = Path(config["paths"]["predictions"])
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, probabilities=probabilities, cdf=np.cumsum(probabilities, axis=0),
+                        target=field.target_window(0, 0, window, lead).numpy(), rates=CLASS_RATES,
+                        lead_minutes=np.int32(lead), coverage=np.array(config["inference"]["coverage"]),
+                        is_complete=np.bool_(config["inference"]["is_complete"]))
+    print(output)

scripts/result.py

@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from model.metnet_2 import load_config, scores, write_json
+
+parser = argparse.ArgumentParser(description="Evaluate and visualize MetNet-2 predictions")
+parser.add_argument("--config", default="conf/config.yaml")
+args = parser.parse_args()
+config = load_config(args.config)
+with np.load(config["paths"]["predictions"]) as data:
+    probabilities, target, rates = data["probabilities"], data["target"], data["rates"]
+    metrics = {str(int(data["lead_minutes"])): scores(probabilities, target)}
+if not all(np.isfinite(value) for value in [metrics[next(iter(metrics))]["discrete_crps"],
+                                             *metrics[next(iter(metrics))]["brier"].values(),
+                                             *metrics[next(iter(metrics))]["csi"].values()]):
+    raise FloatingPointError("evaluation metrics are not finite")
+write_json(config["paths"]["evaluation_metrics"], metrics)
+expected, truth = (probabilities * rates[:, None, None]).sum(0), rates[target]
+figure, axes = plt.subplots(1, 3, figsize=(11, 3.5), constrained_layout=True)
+for axis, image, title in zip(axes, (truth, expected, expected - truth), ("Target", "Expected rate", "Error")):
+    plot = axis.imshow(image, cmap="viridis")
+    axis.set_title(title)
+    axis.set_axis_off()
+    figure.colorbar(plot, ax=axis, shrink=.75)
+comparison = Path(config["paths"]["comparison"])
+comparison.parent.mkdir(parents=True, exist_ok=True)
+figure.savefig(comparison, dpi=140)
+plt.close(figure)
+print(config["paths"]["evaluation_metrics"], comparison)

scripts/train.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+import argparse
+import os
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from model.metnet_2 import (WindowDataset, build_model, categorical_nll_chunked,
+                            load_config, save_checkpoint, write_json)
+
+parser = argparse.ArgumentParser(description="Train MetNet-2 on selected windows")
+parser.add_argument("--config", default="conf/config.yaml")
+parser.add_argument("--steps", type=int, default=None)
+args = parser.parse_args()
+config = load_config(args.config)
+rank = int(os.environ.get("RANK", "0"))
+world_size = int(os.environ.get("WORLD_SIZE", "1"))
+local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+distributed = world_size > 1
+requested = config["runtime"]["device"]
+use_cuda = (requested != "cpu" and torch.cuda.is_available()
+            and (not distributed or torch.cuda.device_count() >= world_size))
+if distributed:
+    torch.distributed.init_process_group(backend="nccl" if use_cuda else "gloo")
+torch.manual_seed(config["seed"] + rank)
+torch.set_num_threads(config["runtime"]["num_threads"])
+device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu" if requested == "auto" else requested)
+if use_cuda:
+    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(device)
+model = build_model(config).to(device)
+if distributed:
+    model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])
+dataset = WindowDataset(config["data"]["path"])
+sampler = DistributedSampler(dataset, shuffle=True, seed=config["seed"]) if distributed else None
+loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], shuffle=sampler is None, sampler=sampler)
+steps = args.steps if args.steps is not None else config["training"]["steps"]
+losses = []
+model.train()
+for step, (inputs, target, lead) in enumerate(loader):
+    if step >= steps:
+        break
+    optimizer.zero_grad(set_to_none=True)
+    logits = model(inputs.to(device), lead.to(device), config["data"]["window"])
+    loss = F.cross_entropy(logits, target.to(device))
+    if not torch.isfinite(loss):
+        raise FloatingPointError("training loss is not finite")
+    loss.backward()
+    optimizer.step()
+    losses.append(float(loss))
+    print(f"step={step} nll={losses[-1]:.6f}")
+if not losses:
+    raise RuntimeError("training produced no optimization steps")
+summary = torch.tensor([sum(losses), len(losses)], dtype=torch.float64, device=device)
+if distributed:
+    torch.distributed.all_reduce(summary)
+if rank == 0:
+    save_checkpoint(config["paths"]["checkpoint"], model, config["model"])
+    write_json(config["paths"]["training_metrics"],
+               {"steps": int(summary[1]), "mean_nll": float(summary[0] / summary[1]), "world_size": world_size})
+if distributed:
+    torch.distributed.destroy_process_group()