Upload folder using huggingface_hub

.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 binary

README.md

@@ -0,0 +1,143 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Physics Parameterization
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">StableNN-Phys</span></strong>
+</p>
+
+# Model Introduction
+
+StableNN-Phys is an engineering reproduction of a neural-network unified physics parameterization for an atmospheric single-column model, continuously predicting thermodynamic-state evolution at three-hour intervals from column states and surface fluxes.
+
+Paper: Prognostic Validation of a Neural Network Unified Physics Parameterization  
+https://doi.org/10.1029/2018GL078510
+
+# Model Description
+
+The method reproduced by StableNN-Phys was proposed by researchers from the Department of Atmospheric Sciences at the University of Washington. The model takes 34-level liquid-water static energy and total-water states with surface sensible heat, latent heat, and incoming solar radiation, maps 71 input features to 68 physics tendencies at each step, and integrates multiple steps with advective forcing and no teacher forcing. The model is suitable for neural physics-parameterization training, three-hourly single-column prediction, 64-step long-term stability validation, and column-water-budget diagnosis.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Single-column prediction | Continuously predict three-hourly state evolution from 34-level thermodynamic states and surface fluxes. |
+| Physics-parameterization training | Learn the mapping from 71 input features to 68 physics tendencies with `T=20` multistep windows. |
+| Long-term stability validation | Run a fixed 64-step, eight-day single-column integration without teacher forcing. |
+| Water-budget evaluation | Diagnose precipitation from column-water storage, latent heat flux, and advective moisture convergence. |
+| Hugging Face/OneCode execution | Validate training, inference, evaluation, visualization, and checkpoint workflows in Hugging Face or OneCode environments. |
+
+# 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/StableNN-Phys --local-dir ./StableNN-Phys
+cd StableNN-Phys
+```
+
+### Environment Dependencies
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for connectivity validation with the default small-sample configuration.
+- DCU users must install DTK first. DTK 25.04.2 or later, or the OneScience-recommended version matching the cluster, is recommended.
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+Training data consist of atmospheric-column thermodynamic states and external physical forcing. Inputs contain 34-level static energy, 34-level total water, surface fluxes, and insolation for 71 values. Targets are 34-level temperature and moisture physics tendencies for the same column, giving 68 values. Samples form continuous `T=20` windows at three-hour intervals. This repository uses a small amount of synthetic data to validate training, inference, and evaluation and does not represent the paper data distribution, training scale, or formal performance.
+
+```bash
+python scripts/fake_data.py
+```
+
+### Training
+
+For single-device training, use:
+
+```bash
+python scripts/train.py
+```
+
+For multi-GPU training, use:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training uses Adam to fit `T=20` multistep state sequences, with layer-mass-weighted MAD as the default `paper` loss. The recorded paper configuration uses a learning rate of 0.01, batch size 200, five epochs, and hidden width 128, while the default engineering configuration reduces the hidden width and sample count for quick workflow validation; results are saved to:
+
+```text
+result/checkpoints/stablenn_phys.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not include pretrained weights under `weight/`. Checkpoints generated with synthetic data validate the engineering workflow only and are not official paper weights.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference performs a fixed 64-step, eight-day single-column rollout without teacher forcing and saves the complete state, tendency, and forcing sequence at every three-hour time. Complete numerical results are saved to `result/output/rollout.npz`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation reports length-65 per-lead mass-weighted MAD and bias and computes R2 by data source. Precipitation is diagnosed from column-water storage, latent heat flux, and advective moisture convergence while the budget residual is checked; synthetic-data results validate the engineering workflow only and do not represent formal paper performance. Structured results and the figure are saved to:
+
+```text
+result/evaluation/metrics.json
+result/evaluation/state_precipitation_timeseries.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 StableNN-Phys specifications.
+
+Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.

conf/config.yaml

@@ -0,0 +1,33 @@
+seed: 42
+data:
+  file: data/stablenn_phys_virtual.npz
+  format_version: stablenn_phys_v1
+  levels: 34
+  window_steps: 20
+  columns: 6
+  windows_per_column: 4
+  step_hours: 3
+model:
+  name: StableNN-Phys
+  hidden_size: 32
+  paper_hidden_size: 128
+  input_features: 71
+  output_features: 68
+train:
+  epochs: 3
+  batch_size: 4
+  learning_rate: 0.003
+  loss_mode: paper
+  paper_model:
+    hidden_size: 128
+    learning_rate: 0.01
+    batch_size: 200
+    epochs: 5
+runtime:
+  device: auto
+  rollout_steps: 64
+paths:
+  checkpoint: result/checkpoints/stablenn_phys.pt
+  training_metrics: result/training/metrics.json
+  inference: result/output/rollout.npz
+  evaluation: result/evaluation/metrics.json

config.json

@@ -0,0 +1,42 @@
+{
+  "model_name": "StableNN-Phys",
+  "model_type": "stablenn_phys",
+  "architectures": ["StableNNPhys"],
+  "framework": "PyTorch",
+  "domain": "atmospheric physics",
+  "task": "single-column prognostic rollout",
+  "implementation": {
+    "entry_point": "model/stablenn_phys.py",
+    "scope": "Core-method, full-window and full-vertical-dimension 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": {
+    "levels": 34,
+    "input_features": 71,
+    "output_features": 68,
+    "engineering_hidden_size": 32,
+    "paper_hidden_size": 128,
+    "activation": "ReLU",
+    "linear_bypass": true
+  },
+  "integration": {
+    "step_hours": 3,
+    "training_window_steps": 20,
+    "rollout_steps": 64,
+    "forcing": "trapezoidal horizontal-advection update before neural-network Euler physics update",
+    "teacher_forcing": false
+  },
+  "checkpoint": {
+    "path": "result/checkpoints/stablenn_phys.pt",
+    "required_fields": ["model", "model_config", "format_version"],
+    "format_version": "stablenn_phys_checkpoint_v1"
+  },
+  "loss_modes": {
+    "default": "paper",
+    "paper": "layer-mass-weighted MAD of prognostic state errors over all rollout steps",
+    "official_v0_3": "equal-level normalized MAD of prognostic state errors over all rollout steps"
+  }
+}

configuration.json

@@ -0,0 +1,17 @@
+{
+  "framework": "PyTorch",
+  "task": "atmospheric_single_column_rollout",
+  "model": "StableNN-Phys",
+  "config": "config.json",
+  "input_format": "sL [B,34] + qT [B,34] + SHF/LHF/SOLIN [B,3] = [B,71]",
+  "output_format": "sL and qT tendencies [B,68]",
+  "training_window_steps": 20,
+  "rollout_steps": 64,
+  "per_step_length": 65,
+  "step_hours": 3,
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/stablenn_phys.py

@@ -0,0 +1,51 @@
+import torch
+from torch import nn
+
+
+class StableNNPhys(nn.Module):
+    def __init__(self, hidden_size=32, input_size=71, output_size=68):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.input_size = input_size
+        self.output_size = output_size
+        self.hidden = nn.Linear(input_size, hidden_size)
+        self.output = nn.Linear(hidden_size, output_size)
+        self.bypass = nn.Linear(input_size, output_size)
+        nn.init.zeros_(self.output.weight)
+        nn.init.zeros_(self.output.bias)
+        nn.init.zeros_(self.bypass.weight)
+        nn.init.zeros_(self.bypass.bias)
+
+    def forward(self, x):
+        return self.output(torch.relu(self.hidden(x))) + self.bypass(x)
+
+
+def rollout(model, initial_state, surface, advection, state_mean, state_std,
+            tendency_mean, tendency_std, dt_seconds=10800.0):
+    """Integrate advection trapezoidally, then neural physics with Euler."""
+    states = [initial_state]
+    physics = []
+    state = initial_state
+    for step in range(surface.shape[1]):
+        adv_now = advection[:, step]
+        adv_next = advection[:, min(step + 1, advection.shape[1] - 1)]
+        forced = state + 0.5 * dt_seconds * (adv_now + adv_next)
+        surface_scaled = surface[:, step] / surface.new_tensor([100.0, 100.0, 1000.0])
+        features = torch.cat(((forced - state_mean) / state_std, surface_scaled), dim=-1)
+        tendency = model(features) * tendency_std + tendency_mean
+        state = forced + dt_seconds * tendency
+        physics.append(tendency)
+        states.append(state)
+    return torch.stack(states, dim=1), torch.stack(physics, dim=1)
+
+
+def rollout_loss(prediction, target, layer_mass, mode="paper"):
+    error = torch.abs(prediction[:, 1:] - target[:, 1:])
+    if mode == "paper":
+        weights = torch.cat((layer_mass, layer_mass), dim=-1)
+        weights = weights / weights.mean(dim=-1, keepdim=True)
+        return (error * weights[:, None]).mean()
+    if mode == "official_v0_3":
+        scale = target[:, 1:].std(dim=(0, 1), unbiased=False).clamp_min(1e-6)
+        return (error / scale).mean()
+    raise ValueError(f"Unknown loss mode: {mode}")

scripts/fake_data.py

@@ -0,0 +1,70 @@
+from pathlib import Path
+import math
+import yaml
+import numpy as np
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def main():
+    cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    rng = np.random.default_rng(cfg["seed"])
+    ncol = cfg["data"]["columns"]
+    nwin = cfg["data"]["windows_per_column"]
+    steps = cfg["data"]["window_steps"]
+    levels = cfg["data"]["levels"]
+    dt = cfg["data"]["step_hours"] * 3600.0
+    total_steps = nwin * steps + 64
+    sigma = np.linspace(0.02, 0.995, levels, dtype=np.float32)
+    pressure = 100000.0 * sigma
+    interfaces = np.linspace(0.0, 100000.0, levels + 1, dtype=np.float32)
+    layer_mass = np.diff(interfaces) / 9.80665
+    time = np.arange(total_steps + 1, dtype=np.float32)
+    states = np.empty((ncol, total_steps + 1, 68), dtype=np.float32)
+    surface = np.empty((ncol, total_steps, 3), dtype=np.float32)
+    advection = np.empty((ncol, total_steps + 1, 68), dtype=np.float32)
+    for c in range(ncol):
+        phase = 2.0 * math.pi * c / ncol
+        sl0 = 2.82e5 + 1.25e5 * (1.0 - sigma) + 1400.0 * np.sin(math.pi * sigma + phase)
+        qt0 = 0.0175 * sigma ** 2.6 + 0.00012 + 0.0004 * np.sin(2 * math.pi * sigma + phase)
+        states[c, 0] = np.concatenate((sl0, np.maximum(qt0, 2e-5)))
+        solar = np.maximum(0.0, np.sin(2.0 * math.pi * (time[:-1] % 8.0) / 8.0 + phase))
+        surface[c, :, 0] = 30.0 + 18.0 * np.sin(2 * math.pi * time[:-1] / 8.0 + phase)
+        surface[c, :, 1] = 92.0 + 34.0 * np.sin(2 * math.pi * time[:-1] / 8.0 + phase - 0.5)
+        surface[c, :, 2] = 760.0 * solar
+        vertical = np.sin(math.pi * sigma)
+        for t in range(total_steps + 1):
+            wave = math.sin(2 * math.pi * t / 16.0 + phase)
+            advection[c, t, :levels] = (0.010 * wave * vertical + rng.normal(0, 0.0003, levels))
+            advection[c, t, levels:] = (1.4e-9 * wave * vertical + rng.normal(0, 4e-11, levels))
+        for t in range(total_steps):
+            adv = 0.5 * (advection[c, t] + advection[c, t + 1])
+            sl, qt = np.split(states[c, t], 2)
+            shf, lhf, solin = surface[c, t]
+            heat = 2.0e-7 * (300000.0 - sl) + 1.0e-5 * shf * sigma + 1.0e-7 * solin * vertical
+            moist = 1.4e-7 * (0.012 * sigma ** 2.4 - qt) + 2.5e-11 * lhf * sigma
+            next_state = states[c, t] + dt * (adv + np.concatenate((heat, moist)))
+            next_state[levels:] = np.maximum(next_state[levels:], 1e-6)
+            states[c, t + 1] = next_state
+    initial, targets, surf_windows, adv_windows, source = [], [], [], [], []
+    for c in range(ncol):
+        for w in range(nwin):
+            start = w * steps
+            initial.append(states[c, start])
+            targets.append(states[c, start:start + steps + 1])
+            surf_windows.append(surface[c, start:start + steps])
+            adv_windows.append(advection[c, start:start + steps + 1])
+            source.append(c)
+    out = ROOT / cfg["data"]["file"]
+    out.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(out, initial=np.array(initial), target=np.array(targets),
+                        surface=np.array(surf_windows), advection=np.array(adv_windows),
+                        layer_mass=layer_mass, pressure=pressure, source=np.array(source),
+                        long_state=states, long_surface=surface, long_advection=advection,
+                        dt_seconds=np.float32(dt), format_version=cfg["data"]["format_version"])
+    print(f"saved {out}: windows={len(initial)}, T={steps}, levels={levels}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from model.stablenn_phys import StableNNPhys, rollout
+
+
+CHECKPOINT_FORMAT_VERSION = "stablenn_phys_checkpoint_v1"
+REQUIRED_CHECKPOINT_FIELDS = {"model", "model_config", "format_version"}
+
+
+def main():
+    cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    raw = np.load(ROOT / cfg["data"]["file"])
+    checkpoint = torch.load(ROOT / cfg["paths"]["checkpoint"], map_location="cpu", weights_only=False)
+    missing = REQUIRED_CHECKPOINT_FIELDS.difference(checkpoint)
+    if missing:
+        raise ValueError(f"Checkpoint is missing required fields: {', '.join(sorted(missing))}")
+    if checkpoint["format_version"] != CHECKPOINT_FORMAT_VERSION:
+        raise ValueError(
+            f"Unsupported checkpoint format_version {checkpoint['format_version']!r}; "
+            f"expected {CHECKPOINT_FORMAT_VERSION!r}"
+        )
+    model = StableNNPhys(checkpoint["model_config"]["hidden_size"])
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    steps = cfg["runtime"]["rollout_steps"]
+    if steps != 64:
+        raise ValueError(f"runtime.rollout_steps must be 64, got {steps}")
+    initial = torch.from_numpy(raw["long_state"][:, 0].astype(np.float32))
+    surface = torch.from_numpy(raw["long_surface"][:, :steps].astype(np.float32))
+    advection = torch.from_numpy(raw["long_advection"][:, :steps + 1].astype(np.float32))
+    norm = checkpoint["normalization"]
+    with torch.no_grad():
+        prediction, tendency = rollout(model, initial, surface, advection, norm["state_mean"], norm["state_std"],
+                                       norm["tendency_mean"], norm["tendency_std"], checkpoint["dt_seconds"])
+    target = raw["long_state"][:, :steps + 1]
+    out = ROOT / cfg["paths"]["inference"]
+    out.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(out, prediction=prediction.numpy(), target=target, tendency=tendency.numpy(),
+                        surface=surface.numpy(), advection=advection.numpy(), layer_mass=raw["layer_mass"],
+                         source=np.arange(initial.shape[0]),
+                         lead_hours=np.arange(steps + 1) * cfg["data"]["step_hours"],
+                        dt_seconds=np.float32(checkpoint["dt_seconds"]))
+    print(f"saved {out}: {steps} steps, {steps * 3 / 24:.1f} days")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,67 @@
+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 main():
+    cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = np.load(ROOT / cfg["paths"]["inference"])
+    pred, target = data["prediction"], data["target"]
+    if pred.shape != target.shape or pred.shape[1:] != (65, 68):
+        raise ValueError(f"Expected prediction and target shaped [N,65,68], got {pred.shape} and {target.shape}")
+    mass = np.concatenate((data["layer_mass"], data["layer_mass"]))
+    mass = mass / mass.mean()
+    error = pred - target
+    mad = (np.abs(error) * mass[None, None]).mean(axis=(0, 2))
+    bias = error.mean(axis=(0, 2))
+    flat_target = target[:, 1:].astype(np.float64).reshape(target.shape[0], -1)
+    flat_pred = pred[:, 1:].astype(np.float64).reshape(pred.shape[0], -1)
+    source_r2 = []
+    for truth, estimate in zip(flat_target, flat_pred):
+        denom = np.sum((truth - truth.mean()) ** 2)
+        source_r2.append(float(1.0 - np.sum((truth - estimate) ** 2) / max(denom, 1e-12)))
+    dt = float(data["dt_seconds"])
+    qt_pred, qt_target = pred[..., 34:], target[..., 34:]
+    column_water_pred = (qt_pred * data["layer_mass"][None, None]).sum(-1)
+    column_water_target = (qt_target * data["layer_mass"][None, None]).sum(-1)
+    lhf_evap = data["surface"][:, :, 1] / 2.5e6
+    adv_q = (0.5 * (data["advection"][:, :-1, 34:] + data["advection"][:, 1:, 34:]) *
+             data["layer_mass"][None, None]).sum(-1)
+    storage = np.diff(column_water_pred, axis=1) / dt
+    precipitation = np.maximum(0.0, lhf_evap + adv_q - storage)
+    residual = storage + precipitation - lhf_evap - adv_q
+    metrics = {"mass_weighted_mad_per_step": mad.tolist(), "bias_per_step": bias.tolist(),
+               "source_r2": source_r2, "mean_source_r2": float(np.mean(source_r2)),
+               "water_budget": {"mean_precipitation_kg_m2_s": float(precipitation.mean()),
+                                "mean_abs_residual_kg_m2_s": float(np.abs(residual).mean())},
+                "per_step_length": int(len(mad)), "step_hours": cfg["data"]["step_hours"]}
+    out = ROOT / cfg["paths"]["evaluation"]
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(json.dumps(metrics, indent=2))
+    hours = data["lead_hours"]
+    fig, axes = plt.subplots(2, 1, figsize=(9, 7), constrained_layout=True)
+    axes[0].plot(hours, target[:, :, :34].mean((0, 2)), label="target sL")
+    axes[0].plot(hours, pred[:, :, :34].mean((0, 2)), "--", label="predicted sL")
+    ax2 = axes[0].twinx()
+    ax2.plot(hours, target[:, :, 34:].mean((0, 2)), color="tab:green", label="target qT")
+    ax2.plot(hours, pred[:, :, 34:].mean((0, 2)), "--", color="tab:red", label="predicted qT")
+    axes[0].set(xlabel="lead time (h)", title="SCM state rollout")
+    axes[0].legend(loc="upper left")
+    ax2.legend(loc="upper right")
+    axes[1].plot(hours[1:], precipitation.mean(0) * 86400.0, color="navy")
+    axes[1].set(xlabel="lead time (h)", ylabel="mm day-1", title="Diagnosed precipitation")
+    fig.savefig(out.parent / "state_precipitation_timeseries.png", dpi=150)
+    plt.close(fig)
+    print(f"saved {out}: per_step={len(mad)}, source_R2={metrics['mean_source_r2']:.4f}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,95 @@
+import json
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from model.stablenn_phys import StableNNPhys, rollout, rollout_loss
+
+
+CHECKPOINT_FORMAT_VERSION = "stablenn_phys_checkpoint_v1"
+
+
+def main():
+    cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    torch.manual_seed(cfg["seed"])
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    use_accelerator = torch.cuda.is_available() and torch.cuda.device_count() >= world_size
+    if distributed:
+        dist.init_process_group("nccl" if use_accelerator else "gloo")
+    rank = dist.get_rank() if distributed else 0
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    device = torch.device(f"cuda:{local_rank}" if use_accelerator else "cpu")
+    if use_accelerator:
+        torch.cuda.set_device(device)
+    raw = np.load(ROOT / cfg["data"]["file"])
+    tensors = [torch.from_numpy(raw[name].astype(np.float32)) for name in ("initial", "target", "surface", "advection")]
+    source = torch.from_numpy(raw["source"].astype(np.int64))
+    dataset = TensorDataset(*tensors, source)
+    sampler = DistributedSampler(dataset, shuffle=True, seed=cfg["seed"]) if distributed else None
+    loader = DataLoader(dataset, batch_size=cfg["train"]["batch_size"], sampler=sampler,
+                        shuffle=sampler is None)
+    all_state = torch.from_numpy(raw["target"].astype(np.float32))
+    state_mean = all_state.mean((0, 1)).to(device)
+    state_std = all_state.std((0, 1), unbiased=False).clamp_min(1e-5).to(device)
+    tendency_mean = torch.zeros(68, device=device)
+    tendency_std = torch.cat((torch.full((34,), 0.02), torch.full((34,), 2e-7))).to(device)
+    layer_mass = torch.from_numpy(raw["layer_mass"].astype(np.float32)).to(device)
+    model = StableNNPhys(cfg["model"]["hidden_size"]).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg["train"]["learning_rate"])
+    history = []
+    for epoch in range(cfg["train"]["epochs"]):
+        if sampler:
+            sampler.set_epoch(epoch)
+        total = 0.0
+        for initial, target, surface, advection, _ in loader:
+            initial, target, surface, advection = [x.to(device) for x in (initial, target, surface, advection)]
+            prediction, _ = rollout(model, initial, surface, advection, state_mean, state_std,
+                                    tendency_mean, tendency_std, float(raw["dt_seconds"]))
+            loss = rollout_loss(prediction, target, layer_mass.expand(initial.shape[0], -1), cfg["train"]["loss_mode"])
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            total += loss.item() * initial.shape[0]
+        value = total / len(dataset)
+        history.append(value)
+        if rank == 0:
+            print(f"epoch={epoch + 1} loss={value:.6g}")
+    if rank == 0:
+        module = model.module if distributed else model
+        checkpoint = ROOT / cfg["paths"]["checkpoint"]
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        payload = {"format_version": CHECKPOINT_FORMAT_VERSION, "model": module.state_dict(),
+                   "optimizer_state_dict": optimizer.state_dict(), "epoch": cfg["train"]["epochs"],
+                   "model_config": dict(cfg["model"]), "training_config": dict(cfg["train"]),
+                   "normalization": {"state_mean": state_mean.cpu(), "state_std": state_std.cpu(),
+                                     "tendency_mean": tendency_mean.cpu(), "tendency_std": tendency_std.cpu()},
+                   "variables": {"input": ["sL[34]", "qT[34]", "SHF", "LHF", "SOLIN"],
+                                 "output": ["dsL_dt[34]", "dqT_dt[34]"]}, "dt_seconds": float(raw["dt_seconds"])}
+        torch.save(payload, checkpoint)
+        metrics = ROOT / cfg["paths"]["training_metrics"]
+        metrics.parent.mkdir(parents=True, exist_ok=True)
+        metrics.write_text(json.dumps({"loss": history, "world_size": dist.get_world_size() if distributed else 1,
+                                       "paper_model": cfg["train"]["paper_model"]}, indent=2))
+        print(f"saved {checkpoint}")
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()