Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,2 @@
-*.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

README.md

@@ -0,0 +1,147 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Data Assimilation
+- Mass Conservation
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">MassConservingCNN</span></strong>
+</p>
+
+# Model Introduction
+
+MassConservingCNN corrects mass-conservation violations introduced by localization in ensemble Kalman filter data assimilation. Given an unconstrained analysis and radar-observation locations, it generates an analysis field with nonnegative rain and improved mass conservation for research on physically constrained data assimilation and analysis postprocessing.
+
+Paper: Training a convolutional neural network to conserve mass in data assimilation  
+https://doi.org/10.5194/npg-28-111-2021
+
+# Model Description
+
+MassConservingCNN was proposed by researchers from the Meteorological Institute of Ludwig-Maximilians-Universität München and ClimateAi. The paper trains and validates the model with EnKF unconstrained analyses, QPEns constrained analyses, and radar-observation locations generated by twin experiments with a one-dimensional modified shallow-water model. The model is suitable for mass-conserving data-assimilation correction, rain non-negativity constraints, and physically consistent analysis generation.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Data-assimilation correction | Predict a QPEns-style analysis from `X^a` and a radar-location indicator. |
+| Mass-aware training | Train with the paper Equation 6 error and Equation 7 mass penalty. |
+| ModelScope/OneCode execution | Validate training, inference, evaluation, visualization, and checkpoint workflows in ModelScope or OneCode. |
+| Multi-GPU training | Launch distributed data-parallel training with `torchrun`. |
+
+# Usage Instructions
+
+## 1.OneCode
+
+Experience intelligent, one-click AI4S programming through the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.Download and Installation
+
+```bash
+hf download OneScience-Group/MassConservingCNN --local-dir ./MassConservingCNN
+cd MassConservingCNN
+```
+
+### 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 default generator creates 24 training samples and 12 validation samples while preserving `[B,4,250]` inputs and `[B,3,250]` targets. The data combine periodic waves, smooth convective cells, nonnegative rain related to velocity convergence, rainy-region radar masks, and smooth EnKF-style errors. They validate the engineering workflow only and are not equivalent to the paper's 48,000-sample QPEns datasets.
+
+```bash
+python scripts/fake_data.py --force
+```
+
+### 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 checks the data version, shapes, dtypes, finite values, binary radar masks, and rain non-negativity. The checkpoint stores model parameters, optimizer state, model configuration, normalization statistics, variable order, data version, `eta`, epoch, and seed. Outputs are written to:
+
+```text
+result/checkpoints/massconservingcnn.pt
+result/training/metrics.json
+```
+
+### Trained Weights
+
+This repository does not include weights under `weight/`. The paper does not provide a confirmed official checkpoint, and the current engineering checkpoint is not claimed to be compatible with external weights.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference restores the model parameters produced by training and generates mass-corrected analysis fields from unconstrained analyses and radar-location indicators. Results include input analyses, target analyses, model predictions, radar-observation locations, and the corresponding physical and normalization information, and are saved to:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation results include core prediction-error, mass-conservation-error, and relative-improvement results, together with an input, target, and prediction comparison figure. Structured metrics and the auxiliary figure are saved to the paths below; synthetic-data results validate the engineering workflow only and do not represent paper performance.
+
+```text
+result/evaluation/metrics.json
+result/evaluation/input_target_prediction.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 MassConservingCNN 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,47 @@
+seed: 42
+data:
+  root: data
+  format_version: mass_conserving_msw_v1
+  grid_points: 250
+  domain_km: 125
+  variables: [u, h, r]
+  input_channels: 4
+  train_samples: 24
+  validation_samples: 12
+model:
+  input_channels: 4
+  hidden_channels: 32
+  hidden_layers: 4
+  kernel_size: 3
+train:
+  epochs: 2
+  batch_size: 8
+  learning_rate: 0.001
+  optimizer: Adam
+  eta: 2.0
+  num_workers: 0
+runtime:
+  device: auto
+paths:
+  checkpoint: result/checkpoints/massconservingcnn.pt
+  training_metrics: result/training/metrics.json
+  inference: result/output/predictions.npz
+  evaluation_dir: result/evaluation
+paper_model:
+  input_channels: 4
+  output_channels: 3
+  grid_points: 250
+  hidden_channels: 32
+  hidden_layers: 4
+  kernel_size: 3
+  hidden_activation: SELU
+  rain_activation: ReLU
+  padding: circular
+  influence_radius: 5
+  train_samples: 48000
+  validation_samples: 48000
+  batch_size: 96
+  epochs: 100
+  optimizer: Adam
+  eta: 2.0
+  experiment: dT10_eta2

config.json

@@ -0,0 +1,45 @@
+{
+  "model_name": "MassConservingCNN",
+  "model_type": "massconservingcnn",
+  "architectures": ["MassConservingCNN"],
+  "framework": "PyTorch",
+  "domain": "earth-science",
+  "task": "mass-aware-data-assimilation-analysis",
+  "implementation": {
+    "entry_point": "model/massconservingcnn.py",
+    "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", 4, 250],
+    "output_shape": ["B", 3, 250],
+    "variable_order": ["u", "h", "r"],
+    "hidden_layers": 4,
+    "filters_per_layer": 32,
+    "kernel_size": 3,
+    "padding": "circular",
+    "influence_radius": 5
+  },
+  "paper_model": {
+    "input_channels": 4,
+    "output_channels": 3,
+    "grid_points": 250,
+    "hidden_channels": 32,
+    "hidden_layers": 4,
+    "kernel_size": 3,
+    "hidden_activation": "SELU",
+    "rain_activation": "ReLU",
+    "padding": "circular",
+    "influence_radius": 5,
+    "train_samples": 48000,
+    "validation_samples": 48000,
+    "batch_size": 96,
+    "epochs": 100,
+    "optimizer": "Adam",
+    "eta": 2.0,
+    "experiment": "dT10_eta2"
+  },
+  "configuration_sources": ["conf/config.yaml", "model/massconservingcnn.py", "scripts/fake_data.py", "scripts/train.py", "scripts/inference.py", "scripts/result.py"]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "mass-aware-data-assimilation-analysis",
+  "model": "MassConservingCNN",
+  "input_format": "BCX",
+  "protocol": "1D modified shallow-water X^a plus radar indicator to QPEns analysis",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/massconservingcnn.py

@@ -0,0 +1,50 @@
+"""Periodic one-dimensional CNN for mass-aware data-assimilation correction."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class PeriodicConv1d(nn.Module):
+    """Conv1d with explicit circular padding and unchanged spatial length."""
+
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int):
+        super().__init__()
+        if kernel_size % 2 != 1:
+            raise ValueError("kernel_size must be odd")
+        self.pad = kernel_size // 2
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=0)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.conv(F.pad(inputs, (self.pad, self.pad), mode="circular"))
+
+
+class MassConservingCNN(nn.Module):
+    """Four hidden SELU convolutions followed by the u/h/r output layer."""
+
+    def __init__(self, input_channels: int = 4, hidden_channels: int = 32,
+                 hidden_layers: int = 4, kernel_size: int = 3):
+        super().__init__()
+        if input_channels != 4 or hidden_layers != 4 or kernel_size != 3:
+            raise ValueError("paper architecture requires 4 inputs, 4 hidden layers, kernel size 3")
+        layers = []
+        channels = input_channels
+        for _ in range(hidden_layers):
+            layers.extend((PeriodicConv1d(channels, hidden_channels, kernel_size), nn.SELU()))
+            channels = hidden_channels
+        self.hidden = nn.Sequential(*layers)
+        self.output = PeriodicConv1d(hidden_channels, 3, kernel_size)
+
+    @property
+    def influence_radius(self) -> int:
+        return 5
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        if inputs.ndim != 3 or inputs.shape[1] != 4 or inputs.shape[2] != 250:
+            raise ValueError(f"expected float tensor [B,4,250], got {tuple(inputs.shape)}")
+        if not inputs.is_floating_point():
+            raise TypeError("inputs must have a floating-point dtype")
+        raw = self.output(self.hidden(inputs))
+        return torch.cat((raw[:, :2], F.relu(raw[:, 2:3])), dim=1)

scripts/fake_data.py

@@ -0,0 +1,84 @@
+"""Generate small, structured modified-shallow-water analysis pairs."""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def periodic_gaussian(x, center, width):
+    distance = np.minimum(np.abs(x - center), 1.0 - np.abs(x - center))
+    return np.exp(-0.5 * (distance / width) ** 2)
+
+
+def make_split(path, count, config, seed):
+    rng = np.random.default_rng(seed)
+    n = int(config["data"]["grid_points"])
+    x = np.arange(n, dtype=np.float32) / n
+    xa = np.empty((count, 3, n), dtype=np.float32)
+    target = np.empty_like(xa)
+    radar = np.empty((count, 1, n), dtype=np.float32)
+    for sample in range(count):
+        phase = rng.uniform(0.0, 1.0)
+        wave = np.sin(2 * np.pi * (x - phase))
+        harmonic = np.sin(4 * np.pi * (x - 0.6 * phase))
+        convective = periodic_gaussian(x, (phase + 0.23) % 1.0, 0.045)
+        secondary = periodic_gaussian(x, (phase + 0.66) % 1.0, 0.07)
+        u_true = 0.75 * wave + 0.22 * harmonic - 0.28 * np.gradient(convective)
+        h_true = 10.0 + 0.35 * np.cos(2 * np.pi * (x - phase)) + 0.5 * convective
+        convergence = np.maximum(-np.gradient(u_true), 0.0)
+        r_true = np.maximum(0.0, 0.7 * convective + 0.28 * convergence - 0.09)
+        rain_mask = (r_true > 0.08).astype(np.float32)
+
+        # Smooth EnKF-like errors are tied to convection and dry-region mass drift.
+        dry = 1.0 - rain_mask
+        u_error = 0.11 * secondary - 0.07 * convective + 0.025 * harmonic
+        h_error = 0.16 * dry + 0.08 * secondary - 0.05 * convective
+        r_error = 0.13 * secondary * dry - 0.06 * convective
+        xa[sample, 0] = u_true + u_error
+        xa[sample, 1] = h_true + h_error
+        xa[sample, 2] = np.maximum(0.0, r_true + r_error)
+        target[sample] = np.stack((u_true, h_true, r_true))
+        radar[sample, 0] = rain_mask
+
+    # Shared synthetic climatology keeps train and validation normalization identical.
+    means = np.asarray([0.0, 10.0], dtype=np.float32)
+    stds = np.asarray([0.6, 0.4, 0.3], dtype=np.float32)
+    normalized_x = xa.copy()
+    normalized_y = target.copy()
+    normalized_x[:, :2] = (xa[:, :2] - means[None, :, None]) / stds[None, :2, None]
+    normalized_y[:, :2] = (target[:, :2] - means[None, :, None]) / stds[None, :2, None]
+    normalized_x[:, 2] = xa[:, 2] / stds[2]
+    normalized_y[:, 2] = target[:, 2] / stds[2]
+    inputs = np.concatenate((normalized_x, radar), axis=1).astype(np.float32)
+    np.savez_compressed(
+        path, inputs=inputs, targets=normalized_y.astype(np.float32), xa=xa,
+        targets_physical=target, radar=radar, climate_mean_uh=means,
+        climate_std_uhr=stds, format_version=np.asarray(config["data"]["format_version"]),
+        variable_order=np.asarray(["u", "h", "r"]), input_layout=np.asarray("BCX"),
+        data_source=np.asarray("structured_synthetic_msw"),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--force", action="store_true")
+    args = parser.parse_args()
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    output = ROOT / config["data"]["root"]
+    output.mkdir(parents=True, exist_ok=True)
+    splits = (("train.npz", int(config["data"]["train_samples"])),
+              ("validation.npz", int(config["data"]["validation_samples"])))
+    for offset, (name, count) in enumerate(splits):
+        path = output / name
+        if args.force or not path.exists():
+            make_split(path, count, config, int(config["seed"]) + offset)
+        print(f"generated={path.relative_to(ROOT)} samples={count} shape=({count},4,250)")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,54 @@
+"""Run validated inference and save normalized and physical fields."""
+
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.massconservingcnn import MassConservingCNN
+from train import MSWDataset, device_from_config
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    device = device_from_config(config)
+    checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
+    required = {"model", "optimizer_state_dict", "model_config", "epoch", "eta",
+                "format_version", "variable_order", "normalization", "climate_mean_uh", "climate_std_uhr", "seed"}
+    if not required.issubset(checkpoint):
+        raise ValueError(f"incomplete checkpoint, missing {sorted(required - set(checkpoint))}")
+    if checkpoint["format_version"] != config["data"]["format_version"] or checkpoint["variable_order"] != ["u", "h", "r"]:
+        raise ValueError("checkpoint protocol mismatch")
+    dataset = MSWDataset(ROOT / config["data"]["root"] / "validation.npz", config)
+    loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), shuffle=False)
+    model = MassConservingCNN(**checkpoint["model_config"]).to(device)
+    model.load_state_dict(checkpoint["model"]); model.eval()
+    outputs = []
+    with torch.no_grad():
+        for inputs, _ in loader:
+            outputs.append(model(inputs.to(device)).cpu().numpy())
+    predictions = np.concatenate(outputs).astype(np.float32)
+    if predictions.shape != dataset.data["targets"].shape or predictions.dtype != np.float32 or not np.isfinite(predictions).all():
+        raise ValueError("invalid inference output")
+    means = np.asarray(checkpoint["climate_mean_uh"], dtype=np.float32)
+    stds = np.asarray(checkpoint["climate_std_uhr"], dtype=np.float32)
+    physical = predictions.copy()
+    physical[:, :2] = predictions[:, :2] * stds[None, :2, None] + means[None, :, None]
+    physical[:, 2] = predictions[:, 2] * stds[2]
+    output = ROOT / config["paths"]["inference"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, predictions=predictions, predictions_physical=physical,
+                        inputs=dataset.data["inputs"], xa=dataset.data["xa"], targets=dataset.data["targets"],
+                        targets_physical=dataset.data["targets_physical"], radar=dataset.data["radar"],
+                        format_version=np.asarray(config["data"]["format_version"]), variable_order=np.asarray(["u", "h", "r"]))
+    print(f"predictions={output.relative_to(ROOT)} shape={predictions.shape} dtype={predictions.dtype}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,72 @@
+"""Compute paper-aligned metrics and plot input, target, and prediction."""
+
+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 metrics(candidate_n, candidate_p, target_n, target_p):
+    rmse = np.sqrt(np.mean((candidate_n - target_n) ** 2, axis=(0, 2)))
+    sample_variable_rmse = np.sqrt(np.mean((candidate_n - target_n) ** 2, axis=2))
+    mass_h = np.mean(np.abs(candidate_p[:, 1].sum(1) - target_p[:, 1].sum(1)) / candidate_p.shape[2])
+    mass_r = np.mean(np.abs(candidate_p[:, 2].sum(1) - target_p[:, 2].sum(1)) / candidate_p.shape[2])
+    h_bias = np.mean(candidate_p[:, 1] - target_p[:, 1])
+    return {"J": float(sample_variable_rmse.mean()), "rmse_u": float(rmse[0]),
+            "rmse_h": float(rmse[1]), "rmse_r": float(rmse[2]),
+            "mass_error_h_per_point": float(mass_h), "mass_error_r_per_point": float(mass_r),
+            "h_bias": float(h_bias)}
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    data = np.load(ROOT / config["paths"]["inference"])
+    if str(data["format_version"]) != config["data"]["format_version"]:
+        raise ValueError("prediction format mismatch")
+    input_n, target_n, prediction_n = data["inputs"][:, :3], data["targets"], data["predictions"]
+    input_p, target_p, prediction_p = data["xa"], data["targets_physical"], data["predictions_physical"]
+    expected = (len(target_n), 3, 250)
+    if any(array.shape != expected for array in (input_n, target_n, prediction_n, input_p, target_p, prediction_p)):
+        raise ValueError("evaluation arrays must have shape [B,3,250]")
+    baseline = metrics(input_n, input_p, target_n, target_p)
+    prediction = metrics(prediction_n, prediction_p, target_n, target_p)
+    improvement = {key: float(100 * (baseline[key] - prediction[key]) / baseline[key])
+                   for key in ("J", "rmse_u", "rmse_h", "rmse_r", "mass_error_h_per_point", "mass_error_r_per_point")
+                   if baseline[key] != 0}
+    if baseline["h_bias"] != 0:
+        improvement["absolute_h_bias"] = float(
+            100 * (abs(baseline["h_bias"]) - abs(prediction["h_bias"])) / abs(baseline["h_bias"])
+        )
+    report = {"samples": len(target_n), "baseline_input": baseline, "prediction": prediction,
+              "relative_improvement_percent": improvement,
+              "note": "Structured synthetic engineering validation; not paper performance."}
+    values = list(baseline.values()) + list(prediction.values()) + list(improvement.values())
+    if not np.isfinite(values).all():
+        raise FloatingPointError("non-finite evaluation metric")
+    output = ROOT / config["paths"]["evaluation_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
+    x = np.arange(250) * float(config["data"]["domain_km"]) / 250
+    figure, axes = plt.subplots(3, 1, figsize=(11, 8), sharex=True)
+    for index, (axis, variable) in enumerate(zip(axes, ("u", "h", "r"))):
+        axis.plot(x, input_p[0, index], color="steelblue", label="input X^a", linewidth=1.4)
+        axis.plot(x, target_p[0, index], color="black", label="QPEns target", linewidth=1.5)
+        axis.plot(x, prediction_p[0, index], color="firebrick", label="CNN prediction", linewidth=1.3)
+        if variable == "r":
+            axis.fill_between(x, 0, data["radar"][0, 0] * max(target_p[0, 2].max(), 1e-6), color="gold", alpha=0.2, label="radar mask")
+        axis.set_ylabel(variable); axis.grid(alpha=0.2)
+    axes[0].legend(ncol=3); axes[-1].set_xlabel("distance (km)")
+    figure.suptitle("MassConservingCNN structured synthetic validation")
+    figure.tight_layout(); figure.savefig(output / "input_target_prediction.png", dpi=150); plt.close(figure)
+    print(f"evaluation={output.relative_to(ROOT)} J={prediction['J']:.6f} h_mass={prediction['mass_error_h_per_point']:.6f}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,128 @@
+"""Train MassConservingCNN with optional torchrun DDP."""
+
+import json
+import os
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+from model.massconservingcnn import MassConservingCNN
+
+
+class MSWDataset(Dataset):
+    def __init__(self, path, config):
+        self.data = np.load(path)
+        if str(self.data["format_version"]) != config["data"]["format_version"]:
+            raise ValueError("incompatible data format version")
+        count = len(self.data["inputs"])
+        if self.data["inputs"].shape != (count, 4, 250):
+            raise ValueError("inputs must have shape [B,4,250]")
+        if self.data["targets"].shape != (count, 3, 250):
+            raise ValueError("targets must have shape [B,3,250]")
+        if self.data["inputs"].dtype != np.float32 or self.data["targets"].dtype != np.float32:
+            raise TypeError("inputs and targets must be float32")
+        if not np.isfinite(self.data["inputs"]).all() or not np.isfinite(self.data["targets"]).all():
+            raise ValueError("data must be finite")
+        if not np.isin(self.data["radar"], (0.0, 1.0)).all():
+            raise ValueError("radar indicator must be binary")
+        if (self.data["inputs"][:, 2] < 0).any() or (self.data["targets"][:, 2] < 0).any():
+            raise ValueError("normalized rain must remain non-negative")
+
+    def __len__(self):
+        return len(self.data["inputs"])
+
+    def __getitem__(self, index):
+        return torch.from_numpy(self.data["inputs"][index]), torch.from_numpy(self.data["targets"][index])
+
+
+def paper_j(prediction, target):
+    return torch.sqrt(torch.mean((prediction - target) ** 2, dim=2) + 1e-12).mean(dim=1)
+
+
+def mass_aware_loss(prediction, target, eta):
+    base = paper_j(prediction, target)
+    mass = eta / prediction.shape[2] * torch.abs(prediction[:, 1].sum(1) - target[:, 1].sum(1))
+    return (base + mass).mean(), base.mean(), mass.mean()
+
+
+def device_from_config(config, local_rank=0):
+    requested = config["runtime"]["device"]
+    if requested == "auto":
+        return torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu")
+    return torch.device(requested)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    seed = int(config["seed"])
+    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if distributed:
+        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    rank = torch.distributed.get_rank() if distributed else 0
+    device = device_from_config(config, local_rank)
+    if device.type == "cuda":
+        torch.cuda.set_device(device); torch.cuda.manual_seed_all(seed)
+    train_set = MSWDataset(ROOT / config["data"]["root"] / "train.npz", config)
+    valid_set = MSWDataset(ROOT / config["data"]["root"] / "validation.npz", config)
+    sampler = DistributedSampler(train_set, shuffle=True, seed=seed) if distributed else None
+    loader = DataLoader(train_set, batch_size=int(config["train"]["batch_size"]),
+                        shuffle=sampler is None, sampler=sampler,
+                        num_workers=int(config["train"]["num_workers"]))
+    valid_loader = DataLoader(valid_set, batch_size=int(config["train"]["batch_size"]), shuffle=False)
+    model = MassConservingCNN(**config["model"]).to(device)
+    wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model
+    bare = wrapped.module if distributed else wrapped
+    optimizer = torch.optim.Adam(wrapped.parameters(), lr=float(config["train"]["learning_rate"]))
+    history = []
+    for epoch in range(int(config["train"]["epochs"])):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        wrapped.train(); total = 0.0; seen = 0
+        for inputs, targets in loader:
+            prediction = wrapped(inputs.to(device)); loss, _, _ = mass_aware_loss(prediction, targets.to(device), float(config["train"]["eta"]))
+            optimizer.zero_grad(set_to_none=True); loss.backward(); optimizer.step()
+            total += float(loss.detach()) * len(inputs); seen += len(inputs)
+        totals = torch.tensor([total, seen], dtype=torch.float64, device=device)
+        if distributed:
+            torch.distributed.all_reduce(totals)
+        wrapped.eval(); valid_total = valid_j = valid_mass = 0.0; valid_seen = 0
+        if rank == 0:
+            with torch.no_grad():
+                for inputs, targets in valid_loader:
+                    loss, base, mass = mass_aware_loss(bare(inputs.to(device)), targets.to(device), float(config["train"]["eta"]))
+                    valid_total += float(loss) * len(inputs); valid_j += float(base) * len(inputs)
+                    valid_mass += float(mass) * len(inputs); valid_seen += len(inputs)
+            history.append({"epoch": epoch + 1, "train_loss": float(totals[0] / totals[1]),
+                            "validation_loss": valid_total / valid_seen, "validation_J": valid_j / valid_seen,
+                            "validation_mass_penalty": valid_mass / valid_seen})
+    if rank == 0:
+        checkpoint_path = ROOT / config["paths"]["checkpoint"]
+        metrics_path = ROOT / config["paths"]["training_metrics"]
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True); metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        model_state = bare.state_dict()
+        torch.save({"model": model_state, "model_state_dict": model_state,
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "model_config": config["model"], "epoch": int(config["train"]["epochs"]),
+                    "eta": float(config["train"]["eta"]), "format_version": config["data"]["format_version"],
+                    "variable_order": ["u", "h", "r"], "normalization": "u,h: center/scale; r: scale only",
+                    "climate_mean_uh": train_set.data["climate_mean_uh"],
+                    "climate_std_uhr": train_set.data["climate_std_uhr"], "seed": seed}, checkpoint_path)
+        metrics_path.write_text(json.dumps({"history": history}, indent=2) + "\n")
+        print(f"checkpoint={checkpoint_path.relative_to(ROOT)} validation_loss={history[-1]['validation_loss']:.6f}")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()