Publish NNCAM reproduction

.gitattributes

@@ -1,35 +1,7 @@
-*.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
+*.py text eol=lf
+*.md text eol=lf
+*.json text eol=lf
+*.yaml text eol=lf
+*.npz binary
+*.pt binary
+*.png binary

README.md

@@ -0,0 +1,148 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Climate Parameterization
+- Subgrid Processes
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">NNCAM</span></strong>
+</p>
+
+# Model Introduction
+
+NNCAM predicts physical tendencies and fluxes produced by cloud, convection, and radiation subgrid processes from atmospheric-column states for data-driven climate-model parameterization research.
+
+Paper: Deep learning to represent subgrid processes in climate models  
+https://gmd.copernicus.org/articles/11/3999/2018/
+
+# Model Description
+
+The method was proposed by research teams from Ludwig Maximilian University of Munich, the University of California Irvine, and Columbia University. The paper trains on approximately 140 million atmospheric-column samples from one year of SPCAM aquaplanet simulation. The model predicts 65 heating, moistening, radiative-flux, and precipitation outputs from a 94-dimensional atmospheric-column state.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Subgrid-process parameterization | Predict physical tendencies and fluxes from temperature, humidity, wind, and surface forcing. |
+| Atmospheric-column diagnostics | Validate heating, moistening, radiation, and precipitation relationships over 30 levels. |
+| Local engineering validation | Validate training, inference, conservation diagnostics, and visualization with structured synthetic samples. |
+| ModelScope/OneCode execution | Validate structured data, training, inference, parameterization 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/NNCAM --local-dir ./NNCAM
+cd NNCAM
+```
+
+### 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 uses SPCAM aquaplanet simulations with a 30-minute timestep and 30 vertical levels. Inputs are `[B,94]` temperature, humidity, wind, and surface-forcing columns, and targets are `[B,65]` heating, moistening, four radiative fluxes, and precipitation. This repository uses a small structured synthetic dataset for engineering validation only and does not represent the real SPCAM distribution, training scale, or paper performance.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+For single-GPU training, use:
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+The default reduces the paper's nine 256-node layers to four 32-node layers and shortens training without reducing the 94 inputs, 65 outputs, or 30-level vertical protocol. Formal experiments require real SPCAM data and the paper-scale model, with artifacts saved to:
+
+```text
+result/checkpoints/nncam.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/nncam.pt` and must not be represented as an official pretrained weight.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the training checkpoint and generates subgrid tendencies, radiative fluxes, and precipitation from complete atmospheric-column states. Complete numerical outputs are saved to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation computes grouped RMSE and R² and generates grouped-error and precipitation-prediction 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 NNCAM 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
+- 地球科学
+- 气候参数化
+- 次网格过程
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">NNCAM</span></strong>
+</p>
+
+# 模型介绍
+
+NNCAM 根据大气柱状态预测云、对流和辐射等次网格过程产生的物理倾向与通量，主要用于数据驱动气候模式参数化研究。
+
+论文：Deep learning to represent subgrid processes in climate models  
+https://gmd.copernicus.org/articles/11/3999/2018/
+
+# 模型描述
+
+该方法由慕尼黑大学、加州大学欧文分校和哥伦比亚大学的研究团队提出。论文使用 SPCAM 水行星模拟一年的约 1.4 亿个大气柱样本训练网络。模型适用于从 94 维大气柱状态预测 65 维加热、湿化、辐射通量和降水输出。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 次网格过程参数化 | 根据温度、湿度、风和地表强迫预测物理倾向与通量。 |
+| 大气柱诊断 | 验证 30 层大气柱的加热、湿化、辐射和降水关系。 |
+| 本地工程验证 | 使用结构化虚拟样本验证训练、推理、守恒诊断和可视化流程。 |
+| 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/NNCAM --local_dir ./NNCAM
+cd NNCAM
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 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
+```
+
+### 训练数据介绍
+
+论文训练数据来自 SPCAM 水行星模拟，时间步为 30 分钟并包含 30 个垂直层。输入为 `[B,94]` 的温度、湿度、风和地表强迫，目标为 `[B,65]` 的加热、湿化、四个辐射通量及降水。本仓库使用少量具有垂直和物理关联的虚拟样本验证工程流程，不代表 SPCAM 的真实数据分布、训练规模或论文性能。
+
+```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
+```
+
+默认配置将网络由论文的 9 层、每层 256 个节点缩小为 4 层、每层 32 个节点，并减少训练轮数，但不缩小 94 维输入、65 维输出和 30 层垂直协议。正式实验需要真实 SPCAM 数据和论文规模模型，训练产物保存到：
+
+```text
+result/checkpoints/nncam.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+论文未提供可直接加载的官方模型权重，本仓库不在 `weight/` 中内置权重。本地训练生成的 checkpoint 保存到 `result/checkpoints/nncam.pt`，不得描述为官方预训练权重。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载训练 checkpoint，根据完整大气柱状态生成次网格倾向、辐射通量和降水预测。完整数值结果保存到：
+
+```text
+result/output/predictions.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估按输出组计算 RMSE 和 R²，并生成分组误差及降水预测对比图。虚拟数据结果仅用于验证工程流程，不代表论文正式性能；结果保存到：
+
+```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 |
+
+# 引用与许可证
+
+本仓库为 NNCAM 公开规格的独立工程复现版本。
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,39 @@
+project:
+  name: NNCAM
+  format_version: "1.0"
+  seed: 42
+
+dimensions:
+  input: 94
+  output: 65
+  levels: 30
+
+paths:
+  data: data/nncam_fake.npz
+  checkpoint: result/checkpoints/nncam.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation_metrics: result/evaluation/metrics.json
+  figure: result/evaluation/comparison.png
+
+data:
+  samples: 192
+  seed: 42
+
+model:
+  depth: 4
+  width: 32
+  negative_slope: 0.3
+
+training:
+  epochs: 3
+  batch_size: 64
+  learning_rate: 0.001
+
+paper_model:
+  depth: 9
+  width: 256
+  epochs: 18
+  batch_size: 1024
+  parameters: 567361
+  note: Recorded paper-scale architecture; enabled only by the explicit --paper-model flag.

config.json

@@ -0,0 +1,42 @@
+{
+  "model_name": "NNCAM",
+  "model_type": "nncam",
+  "architectures": ["NNCAM"],
+  "framework": "PyTorch",
+  "domain": "atmospheric-physics",
+  "task": "climate-model-subgrid-parameterization",
+  "implementation": {
+    "entry_point": "model/nncam.py",
+    "scope": "core-method and full-column reduced-sample engineering reproduction",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "input_shape": ["B", 94],
+    "output_shape": ["B", 65],
+    "engineering_depth": 4,
+    "engineering_width": 32,
+    "paper_depth": 9,
+    "paper_width": 256,
+    "activation": "LeakyReLU"
+  },
+  "data": {
+    "dataset": "SPCAM aquaplanet simulation",
+    "format_version": "1.0",
+    "vertical_levels": 30,
+    "input_variables": ["T", "Q", "V", "Ps", "Sin", "H", "E"],
+    "output_variables": ["dT", "dQ", "SWtoa", "SWsfc", "LWtoa", "LWsfc", "P"],
+    "input_layout": "NF",
+    "synthetic": true
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/nncam.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "atmospheric subgrid-process parameterization",
+  "model": "NNCAM",
+  "input_format": "NF: atmospheric column state [N,94]",
+  "protocol": "SPCAM-supervised offline tendency prediction and engineering column diagnostics",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/nncam.py

@@ -0,0 +1,88 @@
+import math
+
+import numpy as np
+import torch
+from torch import nn
+
+
+INPUT_GROUPS = {"T": slice(0, 30), "Q": slice(30, 60), "V": slice(60, 90), "Ps": slice(90, 91), "Sin": slice(91, 92), "H": slice(92, 93), "E": slice(93, 94)}
+OUTPUT_GROUPS = {"dT": slice(0, 30), "dQ": slice(30, 60), "SWtoa": slice(60, 61), "SWsfc": slice(61, 62), "LWtoa": slice(62, 63), "LWsfc": slice(63, 64), "P": slice(64, 65)}
+CP = 1004.0
+LV = 2.5e6
+OUTPUT_SCALE = np.r_[np.full(30, CP), np.full(30, LV), np.full(4, 1e-3), 2e-2].astype(np.float32)
+
+
+class NNCAM(nn.Module):
+    """Fully connected 94-to-65 atmospheric-column parameterization."""
+
+    def __init__(self, input_dim=94, output_dim=65, width=32, depth=4, negative_slope=0.3):
+        super().__init__()
+        if input_dim != 94 or output_dim != 65:
+            raise ValueError("NNCAM requires 94 inputs and 65 outputs")
+        if width < 1 or depth < 1:
+            raise ValueError("width and depth must be positive")
+        layers = []
+        in_features = input_dim
+        for _ in range(depth):
+            layers.extend((nn.Linear(in_features, width), nn.LeakyReLU(negative_slope)))
+            in_features = width
+        layers.append(nn.Linear(in_features, output_dim))
+        self.network = nn.Sequential(*layers)
+        self.model_config = {"input_dim": input_dim, "output_dim": output_dim, "width": width, "depth": depth, "negative_slope": negative_slope}
+
+    def forward(self, x):
+        if x.ndim != 2 or x.shape[1] != 94:
+            raise ValueError(f"expected input [B, 94], got {tuple(x.shape)}")
+        output = self.network(x)
+        if output.shape != (x.shape[0], 65):
+            raise RuntimeError(f"expected output [B, 65], got {tuple(output.shape)}")
+        return output
+
+
+def build_model(width=32, depth=4, negative_slope=0.3):
+    return NNCAM(width=width, depth=depth, negative_slope=negative_slope)
+
+
+def generate_fake_data(n_samples=192, seed=42):
+    rng = np.random.default_rng(seed)
+    sigma = np.linspace(0.02, 1.0, 30, dtype=np.float32)[None, :]
+    lat = rng.uniform(-math.pi / 2, math.pi / 2, (n_samples, 1)).astype(np.float32)
+    time = rng.uniform(0, 2 * math.pi, (n_samples, 1)).astype(np.float32)
+    ps = rng.normal(101000.0, 1800.0, (n_samples, 1)).astype(np.float32)
+    insolation = np.maximum(0.0, 950.0 * np.cos(lat) * (0.65 + 0.35 * np.sin(time))).astype(np.float32)
+    height = rng.uniform(0.0, 2500.0, (n_samples, 1)).astype(np.float32)
+    evaporation = (35.0 + 80.0 * np.maximum(np.cos(lat), 0.0) + rng.normal(0, 4, (n_samples, 1))).astype(np.float32)
+    temperature = 205.0 + 83.0 * sigma**0.24 - 0.006 * height + 4.0 * np.cos(lat) * sigma
+    temperature += rng.normal(0, 1.2, temperature.shape)
+    humidity = (0.00005 + 0.017 * sigma**3 * np.maximum(np.cos(lat), 0.15)) * rng.lognormal(0, 0.12, temperature.shape)
+    wind = 18.0 * np.sin(lat) * (1.0 - sigma) + 5.0 * np.sin(time + 3.0 * sigma) + rng.normal(0, 2, temperature.shape)
+    instability = np.maximum(temperature[:, -1:] - temperature[:, 18:19] - 25.0, 0.0)
+    moisture = humidity[:, -8:].mean(1, keepdims=True)
+    precipitation = np.maximum(0.0, 4.0e-5 * instability * moisture / 0.012 + rng.normal(0, 1.5e-6, (n_samples, 1))).astype(np.float32)
+    heating_shape = np.exp(-((sigma - 0.55) / 0.23) ** 2)
+    drying_shape = np.exp(-((sigma - 0.78) / 0.18) ** 2)
+    dt = (precipitation * LV / CP / 86400.0 * heating_shape - 8e-6 * (temperature - temperature.mean(1, keepdims=True))).astype(np.float32)
+    dq = (-precipitation / 86400.0 * drying_shape + evaporation / LV / 30.0 / 86400.0).astype(np.float32)
+    cloud = np.clip(moisture / 0.014, 0.0, 1.0)
+    targets = (dt, dq, 0.30 * insolation, insolation * (0.72 - 0.18 * cloud), 215.0 + 0.65 * (temperature[:, -1:] - 273.0) - 24.0 * cloud, 330.0 + 1.1 * (temperature[:, -1:] - 285.0) + 18.0 * cloud, precipitation)
+    x = np.concatenate((temperature, humidity, wind, ps, insolation, height, evaporation), axis=1).astype(np.float32)
+    y = np.concatenate(targets, axis=1).astype(np.float32)
+    return x, y, lat[:, 0], time[:, 0]
+
+
+def fit_normalizer(values):
+    mean = values.mean(0).astype(np.float32)
+    scale = np.maximum(np.ptp(values, axis=0), values.std(0))
+    return mean, np.where(scale > 1e-12, scale, 1.0).astype(np.float32)
+
+
+def normalize_input(values, mean, scale):
+    return (values - mean) / scale
+
+
+def scale_output(values):
+    return values * OUTPUT_SCALE
+
+
+def unscale_output(values):
+    return values / OUTPUT_SCALE

scripts/fake_data.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[1]
+import sys
+sys.path.insert(0, str(ROOT))
+from model.nncam import generate_fake_data
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Create deterministic synthetic NNCAM columns.")
+    parser.add_argument("--samples", type=int, default=192)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--output", type=Path, default=ROOT / "data/nncam_fake.npz")
+    args = parser.parse_args()
+    x, y, lat, time = generate_fake_data(args.samples, args.seed)
+    assert x.shape == (args.samples, 94) and y.shape == (args.samples, 65)
+    assert np.isfinite(x).all() and np.isfinite(y).all() and (y[:, 64] >= 0).all()
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(args.output, x=x, y=y, lat=lat, time=time)
+    print(f"saved {args.output}: x={x.shape}, y={y.shape}, finite=true")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from model.nncam import NNCAM, unscale_output
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run offline NNCAM inference.")
+    parser.add_argument("--data", type=Path, default=ROOT / "data/nncam_fake.npz")
+    parser.add_argument("--checkpoint", type=Path, default=ROOT / "result/checkpoints/nncam.pt")
+    parser.add_argument("--output", type=Path, default=ROOT / "result/output/predictions.npz")
+    args = parser.parse_args()
+    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    required = {"model", "model_config", "format_version", "normalization"}
+    if not required.issubset(checkpoint):
+        raise ValueError(f"checkpoint missing {sorted(required - checkpoint.keys())}")
+    model = NNCAM(**checkpoint["model_config"])
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    with np.load(args.data) as data:
+        x, truth, lat, time = (data[name] for name in ("x", "y", "lat", "time"))
+    norm = checkpoint["normalization"]
+    normalized = (torch.from_numpy(x) - norm["input_mean"]) / norm["input_scale"]
+    with torch.no_grad():
+        scaled = model(normalized) * norm["target_scale"] + norm["target_mean"]
+    prediction = unscale_output(scaled.numpy()).astype(np.float32)
+    if prediction.shape != truth.shape or not np.isfinite(prediction).all():
+        raise RuntimeError(f"invalid prediction shape or values: {prediction.shape}")
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(args.output, input=x, truth=truth, prediction=prediction, lat=lat, time=time)
+    print(f"saved {args.output}: prediction={prediction.shape}, finite=true")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def score(truth, prediction):
+    rmse = float(np.sqrt(np.mean((truth - prediction) ** 2)))
+    denominator = float(np.sum((truth - truth.mean()) ** 2))
+    return {"rmse": rmse, "r2": float(1.0 - np.sum((truth - prediction) ** 2) / denominator) if denominator else None}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate NNCAM predictions.")
+    parser.add_argument("--input", type=Path, default=ROOT / "result/output/predictions.npz")
+    parser.add_argument("--metrics", type=Path, default=ROOT / "result/evaluation/metrics.json")
+    parser.add_argument("--figure", type=Path, default=ROOT / "result/evaluation/comparison.png")
+    args = parser.parse_args()
+    with np.load(args.input) as data:
+        truth, prediction = data["truth"], data["prediction"]
+    if truth.shape != prediction.shape or truth.ndim != 2 or truth.shape[1] != 65:
+        raise ValueError(f"expected matching [N,65] arrays, got {truth.shape}, {prediction.shape}")
+    groups = {"dT": slice(0, 30), "dQ": slice(30, 60), "SW": slice(60, 62), "LW": slice(62, 64), "P": slice(64, 65)}
+    metrics = {name: score(truth[:, indices], prediction[:, indices]) for name, indices in groups.items()}
+    metrics["overall"] = score(truth, prediction)
+    values = np.array([value for group in metrics.values() for value in group.values() if value is not None])
+    if not np.isfinite(values).all():
+        raise RuntimeError("evaluation metrics contain non-finite values")
+    args.metrics.parent.mkdir(parents=True, exist_ok=True)
+    args.metrics.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+    names = list(groups)
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4), constrained_layout=True)
+    axes[0].bar(names, [metrics[name]["rmse"] for name in names])
+    axes[0].set_title("Grouped RMSE")
+    axes[1].scatter(truth[:, 64], prediction[:, 64], s=12, alpha=0.7)
+    axes[1].set(xlabel="True precipitation", ylabel="Predicted precipitation", title="Precipitation comparison")
+    fig.savefig(args.figure, dpi=150)
+    plt.close(fig)
+    print(f"saved {args.metrics} and {args.figure}; shape={prediction.shape}, finite=true")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from model.nncam import build_model, fit_normalizer, normalize_input, scale_output
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train NNCAM on the prepared NPZ dataset.")
+    parser.add_argument("--data", type=Path, default=ROOT / "data/nncam_fake.npz")
+    parser.add_argument("--checkpoint", type=Path, default=ROOT / "result/checkpoints/nncam.pt")
+    parser.add_argument("--metrics", type=Path, default=ROOT / "result/training/metrics.json")
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--batch-size", type=int)
+    parser.add_argument("--width", type=int)
+    parser.add_argument("--depth", type=int)
+    parser.add_argument("--paper-model", action="store_true", help="Explicitly use depth=9, width=256, epochs=18, batch_size=1024 (567361 parameters).")
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+    if not args.data.is_file():
+        raise FileNotFoundError(f"missing dataset {args.data}; run python scripts/fake_data.py first")
+    defaults = {"depth": 9, "width": 256, "epochs": 18, "batch_size": 1024} if args.paper_model else {"depth": 4, "width": 32, "epochs": 3, "batch_size": 64}
+    depth, width = args.depth or defaults["depth"], args.width or defaults["width"]
+    epochs, batch_size = args.epochs or defaults["epochs"], args.batch_size or defaults["batch_size"]
+    torch.manual_seed(args.seed)
+    with np.load(args.data) as data:
+        x, y = data["x"].astype(np.float32), data["y"].astype(np.float32)
+    if x.ndim != 2 or x.shape[1] != 94 or y.shape != (x.shape[0], 65):
+        raise ValueError(f"expected x=[N,94], y=[N,65], got {x.shape}, {y.shape}")
+    input_mean, input_scale = fit_normalizer(x)
+    scaled_y = scale_output(y)
+    target_mean, target_scale = fit_normalizer(scaled_y)
+    dataset = TensorDataset(torch.from_numpy(normalize_input(x, input_mean, input_scale)), torch.from_numpy(normalize_input(scaled_y, target_mean, target_scale)))
+    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+    model = build_model(width=width, depth=depth)
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.2)
+    history = []
+    for epoch in range(epochs):
+        total = 0.0
+        for xb, yb in loader:
+            optimizer.zero_grad(set_to_none=True)
+            loss = torch.nn.functional.mse_loss(model(xb), yb)
+            loss.backward()
+            optimizer.step()
+            total += loss.item() * len(xb)
+        history.append(total / len(dataset))
+        scheduler.step()
+        print(f"epoch={epoch + 1:02d} loss={history[-1]:.6f}")
+    parameter_count = sum(parameter.numel() for parameter in model.parameters())
+    checkpoint = {
+        "format_version": 1,
+        "model": model.state_dict(),
+        "model_config": model.model_config,
+        "normalization": {"input_mean": torch.from_numpy(input_mean), "input_scale": torch.from_numpy(input_scale), "target_mean": torch.from_numpy(target_mean), "target_scale": torch.from_numpy(target_scale)},
+        "training": {"epochs": epochs, "batch_size": batch_size, "learning_rate": args.lr, "paper_model": args.paper_model, "parameters": parameter_count},
+    }
+    args.checkpoint.parent.mkdir(parents=True, exist_ok=True)
+    args.metrics.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(checkpoint, args.checkpoint)
+    args.metrics.write_text(json.dumps({"loss": history, "final_loss": history[-1], "parameters": parameter_count, "model_config": model.model_config}, indent=2), encoding="utf-8")
+    print(f"saved {args.checkpoint}; parameters={parameter_count}; paper_model={args.paper_model}")
+
+
+if __name__ == "__main__":
+    main()