Update CorrDiff model package

conf/config.yaml

@@ -1,21 +1,43 @@
-# Lightweight OneScience configuration for the official CorrDiff example.
-model:
-  regression_config: conf/config_training_taiwan_regression.yaml
-  diffusion_config: conf/config_training_taiwan_diffusion.yaml
-  generation_config: conf/config_generate_taiwan.yaml
-  input_channels: 12
-  output_channels: 4
-  input_shape: [36, 36]
-  output_shape: [448, 448]
-  ensemble_size: 32
-
+seed: 42
 data:
-  type: cwb_zarr
-  path: ./data/era5_corrdiff.npz
+  path: data/corrdiff.npz
+  protocol: corrdiff_npz_v1
+  input_key: input
+  target_key: target
+  input_shape: [12, 36, 36]
+  target_shape: [4, 448, 448]
   input_variables: [tcwv, t2m, u10m, v10m, t500, z500, u500, v500, t850, z850, u850, v850]
-  output_variables: [t2m, u10m, v10m, maximum_radar_reflectivity]
-
+  target_variables: [t2m, u10m, v10m, maximum_radar_reflectivity]
+  fake_samples: 2
+model:
+  in_channels: 12
+  out_channels: 4
+  base_channels: 8
+  feature_size: 28
+  output_size: 448
+  sigma_data: 0.5
+training:
+  regression_epochs: 1
+  diffusion_epochs: 1
+  batch_size: 1
+  max_batches_per_epoch: 1
+  learning_rate: 0.0002
+  num_workers: 0
+  amp: true
+  p_mean: -1.2
+  p_std: 1.2
+sampling:
+  batch_size: 1
+  ensemble_size: 2
+  steps: 3
+  sigma_min: 0.002
+  sigma_max: 5.0
+  rho: 7.0
+  solver: heun
 runtime:
   device: auto
-  output_dir: ./result/output
-  seed: 42
+paths:
+  checkpoint: result/checkpoints/corrdiff.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation_dir: result/evaluation

config.json

@@ -2,37 +2,95 @@
   "model_name": "CorrDiff",
   "model_type": "corrdiff",
   "architectures": [
-    "CorrDiff"
+    "CorrDiff",
+    "RegressionUNet",
+    "ResidualDenoiser"
   ],
   "framework": "PyTorch",
   "domain": "atmosphere",
-  "task": "regional-weather-downscaling",
+  "task": "probabilistic-weather-downscaling",
   "implementation": {
     "entry_point": "model/corrdiff.py",
-    "scope": "compact two-stage CorrDiff reproduction"
+    "scope": "conditional-regression mean prediction followed by residual EDM ensemble generation for kilometer-scale atmospheric downscaling",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
   },
   "architecture": {
-    "family": "residual corrective diffusion model",
-    "input_grid_shape": [36, 36],
-    "output_grid_shape": [448, 448],
-    "input_channels": 12,
-    "output_channels": 4,
-    "base_channels": 32,
-    "mean_model": "RegressionUNet",
-    "residual_model": "ResidualDiffusionUNet",
-    "normalization": "group_norm",
-    "activation": "silu",
-    "sampling_steps": 4
+    "family": "conditional regression plus residual corrective diffusion",
+    "in_channels": 12,
+    "out_channels": 4,
+    "base_channels": 8,
+    "feature_size": 28,
+    "output_size": 448,
+    "sigma_data": 0.5,
+    "regression_component": "RegressionUNet",
+    "diffusion_component": "ResidualDenoiser",
+    "diffusion_preconditioning": "EDM",
+    "noise_schedule": "Karras",
+    "sampling": {
+      "ensemble_size": 2,
+      "steps": 3,
+      "sigma_min": 0.002,
+      "sigma_max": 5.0,
+      "rho": 7.0,
+      "solver": "heun"
+    }
   },
   "data": {
-    "dataset": "ERA5 and CWA-WRF",
-    "input_spatial_resolution_km": 25,
-    "output_spatial_resolution_km": 2,
-    "time_step_hours": 1,
-    "protocol": "paired-downscaling"
+    "datasets": [
+      "ERA5",
+      "CWA-WRF"
+    ],
+    "protocol": "corrdiff_npz_v1",
+    "format": "NPZ",
+    "default_file": "data/corrdiff.npz",
+    "input_key": "input",
+    "input_shape": [
+      "N",
+      12,
+      36,
+      36
+    ],
+    "target_key": "target",
+    "target_shape": [
+      "N",
+      4,
+      448,
+      448
+    ],
+    "input_variables": [
+      "tcwv",
+      "t2m",
+      "u10m",
+      "v10m",
+      "t500",
+      "z500",
+      "u500",
+      "v500",
+      "t850",
+      "z850",
+      "u850",
+      "v850"
+    ],
+    "target_variables": [
+      "t2m",
+      "u10m",
+      "v10m",
+      "maximum_radar_reflectivity"
+    ],
+    "required_metadata": [
+      "protocol",
+      "data_source"
+    ]
   },
   "configuration_sources": [
     "conf/config.yaml",
-    "model/corrdiff.py"
+    "model/corrdiff.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
   ]
 }

configuration.json

@@ -1 +1,11 @@
-{"framework":"Pytorch","task":"weather_downscaling"}
+{
+  "framework": "PyTorch",
+  "task": "probabilistic_weather_downscaling",
+  "model": "CorrDiff",
+  "license": "Apache-2.0",
+  "protocol": "corrdiff_npz_v1",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py"
+}

model/corrdiff.py

@@ -1,6 +1,7 @@
-"""Compact CorrDiff implementation for the OneScience model package."""
+"""Compact but faithful conditional-regression and residual-EDM CorrDiff."""
 
 import math
+
 import torch
 from torch import nn
 from torch.nn import functional as F
@@ -9,12 +10,13 @@ from torch.nn import functional as F
 class ConvBlock(nn.Module):
     def __init__(self, in_channels, out_channels):
         super().__init__()
+        groups = min(8, out_channels)
         self.block = nn.Sequential(
             nn.Conv2d(in_channels, out_channels, 3, padding=1),
-            nn.GroupNorm(8, out_channels),
+            nn.GroupNorm(groups, out_channels),
             nn.SiLU(),
             nn.Conv2d(out_channels, out_channels, 3, padding=1),
-            nn.GroupNorm(8, out_channels),
+            nn.GroupNorm(groups, out_channels),
             nn.SiLU(),
         )
         self.skip = nn.Conv2d(in_channels, out_channels, 1)
@@ -24,79 +26,118 @@ class ConvBlock(nn.Module):
 
 
 class RegressionUNet(nn.Module):
-    """Mean predictor: 12x36x36 coarse fields to 4x448x448 fields."""
+    """Deterministic conditional mean, computed cheaply before 448 px recovery."""
 
-    def __init__(self, in_channels=12, out_channels=4, base_channels=32):
+    def __init__(self, in_channels=12, out_channels=4, base_channels=16, feature_size=56,
+                 output_size=448):
         super().__init__()
-        self.encoder = nn.Sequential(
+        self.feature_size = feature_size
+        self.output_size = output_size
+        self.net = nn.Sequential(
             ConvBlock(in_channels, base_channels),
             ConvBlock(base_channels, base_channels * 2),
-            ConvBlock(base_channels * 2, base_channels * 4),
-        )
-        self.head = nn.Sequential(
-            ConvBlock(base_channels * 4, base_channels * 2),
-            nn.Conv2d(base_channels * 2, out_channels, 1),
+            ConvBlock(base_channels * 2, base_channels),
+            nn.Conv2d(base_channels, out_channels, 1),
         )
 
-    def forward(self, x):
-        x = self.encoder[0](x)
-        x = F.avg_pool2d(x, 2)
-        x = self.encoder[1](x)
-        x = F.avg_pool2d(x, 2)
-        x = self.encoder[2](x)
-        x = self.head[0](x)
-        x = F.interpolate(x, size=(448, 448), mode="bilinear", align_corners=False)
-        return self.head[1](x)
+    def forward(self, coarse):
+        x = F.interpolate(coarse, (self.feature_size, self.feature_size), mode="bilinear",
+                          align_corners=False)
+        x = self.net(x)
+        return F.interpolate(x, (self.output_size, self.output_size), mode="bilinear",
+                             align_corners=False)
 
 
-class ResidualDiffusionUNet(nn.Module):
-    """Conditional residual denoiser used by the compact EDM-style sampler."""
+class ResidualDenoiser(nn.Module):
+    """Noise-conditional network used inside EDM preconditioning."""
 
-    def __init__(self, condition_channels=16, out_channels=4, base_channels=32):
+    def __init__(self, condition_channels=16, out_channels=4, base_channels=16,
+                 feature_size=56):
         super().__init__()
-        self.time = nn.Sequential(
+        self.feature_size = feature_size
+        self.noise_mlp = nn.Sequential(
             nn.Linear(1, base_channels), nn.SiLU(), nn.Linear(base_channels, base_channels)
         )
-        self.net = nn.Sequential(
-            ConvBlock(condition_channels, base_channels),
+        self.input = ConvBlock(condition_channels, base_channels)
+        self.body = nn.Sequential(
             ConvBlock(base_channels, base_channels * 2),
             ConvBlock(base_channels * 2, base_channels),
             nn.Conv2d(base_channels, out_channels, 1),
         )
 
-    def forward(self, noisy_residual, condition, sigma):
-        x = torch.cat((noisy_residual, condition), dim=1)
-        embedding = self.time(torch.log(sigma).reshape(-1, 1).float())
-        x = self.net[0](x)
-        x = x + embedding[:, :, None, None]
-        for layer in self.net[1:]:
-            x = layer(x)
-        return x
+    def forward(self, noisy, condition, c_noise):
+        size = noisy.shape[-2:]
+        x = torch.cat((noisy, condition), dim=1)
+        x = F.interpolate(x, (self.feature_size, self.feature_size), mode="bilinear",
+                          align_corners=False)
+        x = self.input(x)
+        x = x + self.noise_mlp(c_noise[:, None].float())[:, :, None, None].to(x.dtype)
+        x = self.body(x)
+        return F.interpolate(x, size, mode="bilinear", align_corners=False)
 
 
 class CorrDiff(nn.Module):
-    """Two-stage mean plus stochastic residual atmospheric downscaler."""
-
-    def __init__(self, in_channels=12, out_channels=4, base_channels=32):
+    def __init__(self, in_channels=12, out_channels=4, base_channels=16, feature_size=56,
+                 output_size=448, sigma_data=0.5):
         super().__init__()
-        self.regression = RegressionUNet(in_channels, out_channels, base_channels)
-        self.diffusion = ResidualDiffusionUNet(in_channels + 2 * out_channels, out_channels, base_channels)
+        self.sigma_data = sigma_data
+        self.output_size = output_size
+        self.regression = RegressionUNet(in_channels, out_channels, base_channels,
+                                         feature_size, output_size)
+        self.diffusion = ResidualDenoiser(in_channels + 2 * out_channels, out_channels,
+                                          base_channels, feature_size)
 
     def mean(self, coarse):
         return self.regression(coarse)
 
-    def denoise(self, noisy_residual, coarse, mean, sigma):
-        coarse = F.interpolate(coarse, size=mean.shape[-2:], mode="bilinear", align_corners=False)
-        return self.diffusion(noisy_residual, torch.cat((coarse, mean), dim=1), sigma)
+    def condition(self, coarse, mean):
+        coarse = F.interpolate(coarse, mean.shape[-2:], mode="bilinear", align_corners=False)
+        return torch.cat((coarse, mean), dim=1)
 
-    def forward(self, coarse, steps=4):
+    def denoise(self, noisy_residual, coarse, mean, sigma):
+        sigma = sigma.reshape(-1, 1, 1, 1).to(noisy_residual.dtype)
+        sigma_data = self.sigma_data
+        c_skip = sigma_data**2 / (sigma.square() + sigma_data**2)
+        c_out = sigma * sigma_data / (sigma.square() + sigma_data**2).sqrt()
+        c_in = (sigma.square() + sigma_data**2).rsqrt()
+        c_noise = sigma.flatten().log() / 4
+        network = self.diffusion(c_in * noisy_residual, self.condition(coarse, mean), c_noise)
+        return c_skip * noisy_residual + c_out * network
+
+    @staticmethod
+    def karras_schedule(steps, sigma_min, sigma_max, rho, device):
+        ramp = torch.linspace(0, 1, steps, device=device)
+        maximum = sigma_max ** (1 / rho)
+        minimum = sigma_min ** (1 / rho)
+        sigmas = (maximum + ramp * (minimum - maximum)) ** rho
+        return torch.cat((sigmas, sigmas.new_zeros(1)))
+
+    def sample(self, coarse, steps=4, sigma_min=0.002, sigma_max=5.0, rho=7.0,
+               solver="heun"):
         mean = self.mean(coarse)
-        residual = torch.randn_like(mean)
-        schedule = torch.linspace(1.0, 0.01, steps, device=coarse.device)
-        for sigma in schedule:
-            sigma_batch = sigma.expand(coarse.shape[0])
-            residual = residual - 0.2 * self.denoise(residual, coarse, mean, sigma_batch)
-        return mean + residual
-
-
-__all__ = ["CorrDiff", "RegressionUNet", "ResidualDiffusionUNet"]
+        sigmas = self.karras_schedule(steps, sigma_min, sigma_max, rho, coarse.device)
+        x = torch.randn_like(mean) * sigmas[0]
+        for index, (current, following) in enumerate(zip(sigmas[:-1], sigmas[1:])):
+            sigma = current.expand(coarse.shape[0])
+            denoised = self.denoise(x, coarse, mean, sigma)
+            derivative = (x - denoised) / current
+            proposal = x + (following - current) * derivative
+            if solver == "heun" and index < len(sigmas) - 2:
+                next_sigma = following.expand(coarse.shape[0])
+                next_denoised = self.denoise(proposal, coarse, mean, next_sigma)
+                next_derivative = (proposal - next_denoised) / following
+                x = x + (following - current) * (derivative + next_derivative) / 2
+            else:
+                x = proposal
+        return mean + x
+
+    def forward(self, coarse, mode="sample", mean=None, noisy=None, sigma=None,
+                **sample_options):
+        if mode == "mean":
+            return self.mean(coarse)
+        if mode == "denoise":
+            return self.denoise(noisy, coarse, mean, sigma)
+        return self.sample(coarse, **sample_options)
+
+
+__all__ = ["CorrDiff", "RegressionUNet", "ResidualDenoiser"]

scripts/fake_data.py

@@ -1,22 +1,41 @@
-from pathlib import Path
+"""Generate a tiny, structurally realistic CorrDiff NPZ dataset."""
+
 import argparse
+from pathlib import Path
+
 import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--output", default="data/fake_corrdiff.npz")
-    parser.add_argument("--samples", type=int, default=1)
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
+    parser.add_argument("--output")
     args = parser.parse_args()
-    path = Path(args.output)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    rng = np.random.default_rng(42)
-    inputs = rng.standard_normal((args.samples, 12, 36, 36)).astype("float32")
-    targets = rng.standard_normal((args.samples, 4, 448, 448)).astype("float32")
-    np.savez_compressed(path, input=inputs, target=targets)
-    print(f"input: {inputs.shape}")
-    print(f"target: {targets.shape}")
-    print(f"saved: {path}")
+    config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
+    data = config["data"]
+    count = data["fake_samples"]
+    rng = np.random.default_rng(config["seed"])
+    coarse = rng.normal(size=(count, *data["input_shape"])).astype("float32")
+    # Correlated targets make the regression/backward smoke test meaningful.
+    y = np.linspace(-1, 1, data["target_shape"][1], dtype="float32")
+    x = np.linspace(-1, 1, data["target_shape"][2], dtype="float32")
+    yy, xx = np.meshgrid(y, x, indexing="ij")
+    target = np.empty((count, *data["target_shape"]), dtype="float32")
+    coarse_signal = coarse.mean(axis=(2, 3))
+    for sample in range(count):
+        for channel in range(data["target_shape"][0]):
+            target[sample, channel] = coarse_signal[sample, channel] + 0.3 * np.sin(
+                (channel + 1) * np.pi * xx
+            ) + 0.2 * np.cos((channel + 1) * np.pi * yy)
+    target += rng.normal(0, 0.05, target.shape).astype("float32")
+    output = Path(args.output) if args.output else ROOT / data["path"]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, input=coarse, target=target,
+                        protocol=np.asarray(data["protocol"]), data_source=np.asarray("synthetic"))
+    print(f"saved={output} input={coarse.shape} target={target.shape}")
 
 
 if __name__ == "__main__":

scripts/inference.py

@@ -1,60 +1,66 @@
+"""Generate a checkpoint-backed CorrDiff ensemble in device-sized batches."""
+
 import argparse
-from pathlib import Path
 import sys
+from pathlib import Path
 
 import numpy as np
 import torch
+import yaml
 
-sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
 from model.corrdiff import CorrDiff
 
 
+def scalar(archive, key):
+    if key not in archive or archive[key].ndim != 0:
+        raise ValueError(f"NPZ metadata {key} must be present as a scalar")
+    return str(archive[key].item())
+
+
 def main():
-    parser = argparse.ArgumentParser(description="Generate a CorrDiff ensemble")
-    parser.add_argument("--data", default="data/era5_corrdiff.npz")
-    parser.add_argument("--regression-checkpoint", default="data/checkpoints/regression_model.pth")
-    parser.add_argument("--diffusion-checkpoint", default="data/checkpoints/diffusion_model.pth")
-    parser.add_argument("--ensemble-size", type=int, default=32)
-    parser.add_argument("--output", default="result/output/predictions.npz")
-    parser.add_argument("--seed", type=int, default=42)
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
+    parser.add_argument("--checkpoint")
     args = parser.parse_args()
-
-    data = np.load(args.data)
-    coarse = torch.from_numpy(data["input"])
-    target = data["target"].astype("float32") if "target" in data else None
-    model = CorrDiff()
-    regression_checkpoint = Path(args.regression_checkpoint)
-    diffusion_checkpoint = Path(args.diffusion_checkpoint)
-    if regression_checkpoint.exists():
-        state = torch.load(regression_checkpoint, map_location="cpu", weights_only=False)
-        model.regression.load_state_dict(state["model"])
-    if diffusion_checkpoint.exists():
-        state = torch.load(diffusion_checkpoint, map_location="cpu", weights_only=False)
-        model.diffusion.load_state_dict(state["model"])
+    config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
+    checkpoint_path = Path(args.checkpoint) if args.checkpoint else ROOT / config["paths"]["checkpoint"]
+    if not checkpoint_path.is_file():
+        raise FileNotFoundError(f"Checkpoint is required: {checkpoint_path}")
+    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
+    state = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model = CorrDiff(**config["model"]).to(device)
+    model.load_state_dict(state["model"])
     model.eval()
-
-    samples = []
-    with torch.no_grad():
-        for member in range(args.ensemble_size):
-            torch.manual_seed(args.seed + member)
-            samples.append(model(coarse).numpy().astype("float32"))
-
-    ensemble = np.stack(samples)
-    output = Path(args.output)
+    archive = np.load(ROOT / config["data"]["path"])
+    protocol, data_source = scalar(archive, "protocol"), scalar(archive, "data_source")
+    if protocol != config["data"]["protocol"] or state.get("protocol") != protocol:
+        raise ValueError("Data, checkpoint, and configured protocols must match")
+    if not data_source or state.get("data_source") != data_source:
+        raise ValueError("Data and checkpoint data_source metadata must match")
+    coarse = archive[config["data"]["input_key"]]
+    target = archive[config["data"]["target_key"]]
+    if coarse.ndim != 4 or target.ndim != 4 or tuple(coarse.shape[1:]) != tuple(config["data"]["input_shape"]) or tuple(target.shape[1:]) != tuple(config["data"]["target_shape"]) or len(coarse) != len(target):
+        raise ValueError("Invalid CorrDiff NPZ tensor contract")
+    member_count = config["sampling"]["ensemble_size"]
+    ensemble = np.empty((member_count, len(coarse), *config["data"]["target_shape"]), dtype="float32")
+    options = {key: config["sampling"][key] for key in ("steps", "sigma_min", "sigma_max", "rho", "solver")}
+    batch_size = config["sampling"]["batch_size"]
+    with torch.inference_mode():
+        for member in range(member_count):
+            torch.manual_seed(config["seed"] + member)
+            for start in range(0, len(coarse), batch_size):
+                stop = min(start + batch_size, len(coarse))
+                batch = torch.from_numpy(coarse[start:stop]).to(device)
+                ensemble[member, start:stop] = model.sample(batch, **options).cpu().numpy()
+                del batch
+    output = ROOT / config["paths"]["predictions"]
     output.parent.mkdir(parents=True, exist_ok=True)
-    payload = {
-        "ensemble": ensemble,
-        "ensemble_mean": ensemble.mean(axis=0),
-        "ensemble_std": ensemble.std(axis=0),
-        "input": data["input"].astype("float32"),
-    }
-    if target is not None:
-        payload["target"] = target
-    np.savez_compressed(output, **payload)
-    print(f"ensemble: {ensemble.shape}")
-    print(f"ensemble mean: {payload['ensemble_mean'].shape}")
-    print(f"ensemble std: {payload['ensemble_std'].shape}")
-    print(f"saved: {output}")
+    np.savez_compressed(output, ensemble=ensemble, ensemble_mean=ensemble.mean(0),
+                        ensemble_std=ensemble.std(0), target=target, protocol=np.asarray(protocol),
+                        data_source=np.asarray(data_source), checkpoint=np.asarray(str(checkpoint_path)))
+    print(f"saved={output} ensemble={ensemble.shape} batch_size={batch_size}")
 
 
 if __name__ == "__main__":

scripts/result.py

@@ -1,3 +1,5 @@
+"""Compute deterministic and probabilistic CorrDiff metrics and plots."""
+
 import argparse
 import json
 from pathlib import Path
@@ -6,60 +8,56 @@ import matplotlib
 matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 import numpy as np
+import yaml
 
-
-VARIABLES = ["t2m", "u10m", "v10m", "radar_reflectivity"]
+ROOT = Path(__file__).resolve().parents[1]
 
 
 def crps_ensemble(ensemble, target):
-    first = np.abs(ensemble - target[None]).mean(axis=0)
-    sorted_ensemble = np.sort(ensemble, axis=0)
-    members = ensemble.shape[0]
-    weights = 2 * np.arange(1, members + 1) - members - 1
-    second = (sorted_ensemble * weights[:, None, None, None, None]).sum(axis=0)
-    return first - second / members**2
+    first = np.abs(ensemble - target[None]).mean(0)
+    sorted_members = np.sort(ensemble, axis=0)
+    m = ensemble.shape[0]
+    weights = (2 * np.arange(1, m + 1) - m - 1).reshape(m, 1, 1, 1, 1)
+    return first - (sorted_members * weights).sum(0) / m**2
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Evaluate CorrDiff predictions")
-    parser.add_argument("--prediction", default="result/output/predictions.npz")
-    parser.add_argument("--target", default="data/era5_corrdiff.npz")
-    parser.add_argument("--output-dir", default="result/output")
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
     args = parser.parse_args()
-
-    predictions = np.load(args.prediction)
-    ensemble = predictions["ensemble"]
-    mean = predictions["ensemble_mean"]
-    std = predictions["ensemble_std"]
-    target = predictions["target"] if "target" in predictions else np.load(args.target)["target"]
-
-    output = Path(args.output_dir)
+    config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
+    archive = np.load(ROOT / config["paths"]["predictions"])
+    if "protocol" not in archive or archive["protocol"].ndim != 0 or str(archive["protocol"].item()) != config["data"]["protocol"]:
+        raise ValueError("Prediction protocol does not match the configured protocol")
+    if "data_source" not in archive or archive["data_source"].ndim != 0 or not str(archive["data_source"].item()):
+        raise ValueError("Prediction data_source must be a non-empty scalar")
+    protocol = str(archive["protocol"].item())
+    data_source = str(archive["data_source"].item())
+    ensemble, target = archive["ensemble"], archive["target"]
+    mean, spread = ensemble.mean(0), ensemble.std(0)
+    axes = (0, 2, 3)
+    mae = np.abs(mean - target).mean(axis=axes)
+    rmse = np.sqrt(((mean - target) ** 2).mean(axis=axes))
+    crps = crps_ensemble(ensemble, target).mean(axis=axes)
+    spread_value = spread.mean(axis=axes)
+    names = config["data"]["target_variables"]
+    metrics = {name: {"mae": float(mae[i]), "rmse": float(rmse[i]), "crps": float(crps[i]),
+                      "ensemble_spread": float(spread_value[i])} for i, name in enumerate(names)}
+    metrics["aggregate"] = {key: float(np.mean([metrics[n][key] for n in names]))
+                            for key in ("mae", "rmse", "crps", "ensemble_spread")}
+    output = ROOT / config["paths"]["evaluation_dir"]
     output.mkdir(parents=True, exist_ok=True)
-    mae = np.abs(mean - target).mean(axis=(0, 2, 3))
-    crps = crps_ensemble(ensemble, target).mean(axis=(0, 2, 3))
-    metrics = {
-        name: {"mae": float(mae[channel]), "crps": float(crps[channel])}
-        for channel, name in enumerate(VARIABLES)
-    }
-    (output / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
-
-    figure, axes = plt.subplots(4, 3, figsize=(12, 14))
-    for channel, name in enumerate(VARIABLES):
-        axes[channel, 0].imshow(mean[0, channel], cmap="viridis")
-        axes[channel, 0].set_title(f"{name}: ensemble mean")
-        axes[channel, 1].imshow(std[0, channel], cmap="magma")
-        axes[channel, 1].set_title(f"{name}: ensemble std")
-        axes[channel, 2].imshow(target[0, channel], cmap="viridis")
-        axes[channel, 2].set_title(f"{name}: target")
-        for axis in axes[channel]:
-            axis.set_xticks([])
-            axis.set_yticks([])
-    figure.tight_layout()
-    figure.savefig(output / "prediction_comparison.png", dpi=150)
-    plt.close(figure)
-
-    print(json.dumps(metrics, indent=2))
-    print(f"saved evaluation outputs to: {output}")
+    payload = {"metrics": metrics, "protocol": protocol, "data_source": data_source}
+    (output / "metrics.json").write_text(json.dumps(payload, indent=2) + "\n")
+    figure, plot_axes = plt.subplots(len(names), 4, figsize=(13, 3 * len(names)))
+    for channel, name in enumerate(names):
+        fields = (target[0, channel], mean[0, channel], spread[0, channel], mean[0, channel] - target[0, channel])
+        titles = ("target", "ensemble mean", "ensemble spread", "mean error")
+        for axis, field, title in zip(plot_axes[channel], fields, titles):
+            axis.imshow(field, cmap="coolwarm" if title == "mean error" else "viridis")
+            axis.set_title(f"{name}: {title}"); axis.axis("off")
+    figure.tight_layout(); figure.savefig(output / "ensemble_diagnostics.png", dpi=120); plt.close(figure)
+    print(json.dumps(payload, indent=2)); print(f"evaluation={output}")
 
 
 if __name__ == "__main__":

scripts/train.py

@@ -1,74 +1,164 @@
+"""Train CorrDiff's conditional mean, then its frozen-mean residual EDM."""
+
 import argparse
+import json
+import os
+import random
+import sys
+from contextlib import nullcontext
 from pathlib import Path
+
 import numpy as np
 import torch
+import yaml
+from torch import distributed as dist
 from torch.nn import functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
 
-import sys
-sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
 from model.corrdiff import CorrDiff
 
 
+def scalar(archive, key, default=None):
+    if key not in archive:
+        if default is not None:
+            return default
+        raise ValueError(f"NPZ is missing required metadata: {key}")
+    value = archive[key]
+    if value.ndim != 0:
+        raise ValueError(f"NPZ metadata {key} must be a scalar")
+    return str(value.item())
+
+
+def reduced_average(total, count, device, distributed):
+    values = torch.tensor([total, count], dtype=torch.float64, device=device)
+    if distributed:
+        dist.all_reduce(values, op=dist.ReduceOp.SUM)
+    if values[1].item() == 0:
+        raise RuntimeError("Training stage processed no batches")
+    return (values[0] / values[1]).item()
+
+
 def main():
-    parser = argparse.ArgumentParser(description="Train CorrDiff regression and diffusion stages")
-    parser.add_argument("--data", default="data/era5_corrdiff.npz")
-    parser.add_argument("--steps", type=int, default=2)
-    parser.add_argument("--checkpoint-dir", default="data/checkpoints")
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
     args = parser.parse_args()
-    data = np.load(args.data)
-    coarse = torch.from_numpy(data["input"])
-    target = torch.from_numpy(data["target"])
-    model = CorrDiff()
-    regression_optimizer = torch.optim.Adam(model.regression.parameters(), lr=2e-4)
-    diffusion_optimizer = torch.optim.Adam(model.diffusion.parameters(), lr=2e-4)
-    model.train()
-    regression_losses = []
-    diffusion_losses = []
+    config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8"))
+    global_rank = int(os.getenv("RANK", 0))
+    local_rank = int(os.getenv("LOCAL_RANK", 0))
+    world = int(os.getenv("WORLD_SIZE", 1))
+    distributed = world > 1
+    if distributed:
+        dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    use_cuda = torch.cuda.is_available() and config["runtime"]["device"] != "cpu"
+    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
+    if use_cuda:
+        torch.cuda.set_device(local_rank)
+    seed = config["seed"] + global_rank
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    archive = np.load(ROOT / config["data"]["path"])
+    protocol = scalar(archive, "protocol")
+    data_source = scalar(archive, "data_source")
+    if protocol != config["data"]["protocol"]:
+        raise ValueError(f"Expected protocol {config['data']['protocol']}, got {protocol}")
+    if not data_source:
+        raise ValueError("data_source must be a non-empty scalar")
+    coarse = torch.from_numpy(archive[config["data"]["input_key"]])
+    target = torch.from_numpy(archive[config["data"]["target_key"]])
+    if coarse.ndim != 4 or target.ndim != 4:
+        raise ValueError("CorrDiff input and target must be NCHW tensors")
+    if len(coarse) != len(target) or tuple(coarse.shape[1:]) != tuple(config["data"]["input_shape"]) or tuple(target.shape[1:]) != tuple(config["data"]["target_shape"]):
+        raise ValueError("NPZ tensor shapes do not match config")
+    dataset = TensorDataset(coarse, target)
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler,
+                        shuffle=sampler is None, num_workers=config["training"]["num_workers"])
+    model = CorrDiff(**config["model"]).to(device)
+    if distributed:
+        model = DDP(model, device_ids=[local_rank] if use_cuda else None,
+                    find_unused_parameters=True)
+    base = model.module if distributed else model
+    reg_opt = torch.optim.AdamW(base.regression.parameters(), lr=config["training"]["learning_rate"])
+    diff_opt = torch.optim.AdamW(base.diffusion.parameters(), lr=config["training"]["learning_rate"])
+    amp = bool(config["training"]["amp"] and use_cuda)
+    scaler = torch.amp.GradScaler("cuda", enabled=amp)
+    autocast = (lambda: torch.amp.autocast("cuda", enabled=True)) if amp else nullcontext
+    history = []
 
-    # Stage 1: learn the conditional mean.
-    for step in range(args.steps):
-        mean = model.mean(coarse)
-        regression_loss = F.mse_loss(mean, target)
-        regression_optimizer.zero_grad()
-        regression_loss.backward()
-        regression_optimizer.step()
-        regression_losses.append(regression_loss.item())
-        print(f"regression step={step + 1} loss={regression_loss.item():.6f}")
+    # Stage 1 is completed in full before any residual-EDM update occurs.
+    for epoch in range(config["training"]["regression_epochs"]):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        total = count = 0
+        for batch_index, (coarse_batch, target_batch) in enumerate(loader):
+            coarse_batch, target_batch = coarse_batch.to(device), target_batch.to(device)
+            reg_opt.zero_grad(set_to_none=True)
+            with autocast():
+                loss = F.mse_loss(model(coarse_batch, mode="mean"), target_batch)
+            scaler.scale(loss).backward()
+            scaler.step(reg_opt)
+            scaler.update()
+            total += loss.item()
+            count += 1
+            if batch_index + 1 >= config["training"]["max_batches_per_epoch"]:
+                break
+        value = reduced_average(total, count, device, distributed)
+        record = {"stage": "regression", "epoch": epoch + 1, "regression_mse": value}
+        history.append(record)
+        if global_rank == 0:
+            print(json.dumps(record))
 
-    # Stage 2: freeze the mean model and learn its stochastic residual.
-    model.regression.eval()
-    for parameter in model.regression.parameters():
+    base.regression.eval()
+    for parameter in base.regression.parameters():
         parameter.requires_grad_(False)
-    for step in range(args.steps):
-        with torch.no_grad():
-            mean = model.mean(coarse)
-        residual = target - mean
-        sigma = torch.rand(coarse.shape[0]).clamp_min(0.01)
-        noisy = residual + sigma[:, None, None, None] * torch.randn_like(residual)
-        predicted = model.denoise(noisy, coarse, mean, sigma)
-        diffusion_loss = F.mse_loss(predicted, residual)
-        diffusion_optimizer.zero_grad()
-        diffusion_loss.backward()
-        diffusion_optimizer.step()
-        diffusion_losses.append(diffusion_loss.item())
-        print(f"diffusion step={step + 1} loss={diffusion_loss.item():.6f}")
+    for epoch in range(config["training"]["diffusion_epochs"]):
+        if sampler is not None:
+            sampler.set_epoch(config["training"]["regression_epochs"] + epoch)
+        base.diffusion.train()
+        total = count = 0
+        for batch_index, (coarse_batch, target_batch) in enumerate(loader):
+            coarse_batch, target_batch = coarse_batch.to(device), target_batch.to(device)
+            with torch.no_grad():
+                mean = base.mean(coarse_batch)
+            residual = target_batch - mean
+            sigma = (torch.randn(len(coarse_batch), device=device) * config["training"]["p_std"] + config["training"]["p_mean"]).exp()
+            noisy = residual + sigma[:, None, None, None] * torch.randn_like(residual)
+            diff_opt.zero_grad(set_to_none=True)
+            with autocast():
+                denoised = model(coarse_batch, mode="denoise", mean=mean, noisy=noisy, sigma=sigma)
+                weight = (sigma.square() + base.sigma_data**2) / (sigma * base.sigma_data).square()
+                loss = (weight[:, None, None, None] * (denoised - residual).square()).mean()
+            scaler.scale(loss).backward()
+            scaler.step(diff_opt)
+            scaler.update()
+            total += loss.item()
+            count += 1
+            if batch_index + 1 >= config["training"]["max_batches_per_epoch"]:
+                break
+        value = reduced_average(total, count, device, distributed)
+        record = {"stage": "diffusion", "epoch": epoch + 1, "edm_loss": value}
+        history.append(record)
+        if global_rank == 0:
+            print(json.dumps(record))
 
-    checkpoint_dir = Path(args.checkpoint_dir)
-    checkpoint_dir.mkdir(parents=True, exist_ok=True)
-    torch.save(
-        {"model": model.regression.state_dict(), "format": "corrdiff-regression-v1"},
-        checkpoint_dir / "regression_model.pth",
-    )
-    torch.save(
-        {"model": model.diffusion.state_dict(), "format": "corrdiff-diffusion-v1"},
-        checkpoint_dir / "diffusion_model.pth",
-    )
-    np.savez(
-        checkpoint_dir / "training_history.npz",
-        regression_loss=np.asarray(regression_losses, dtype="float32"),
-        diffusion_loss=np.asarray(diffusion_losses, dtype="float32"),
-    )
-    print(f"saved checkpoints to: {checkpoint_dir}")
+    if global_rank == 0:
+        checkpoint = ROOT / config["paths"]["checkpoint"]
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"model": base.state_dict(), "config": config, "format": "corrdiff-edm-v3",
+                    "protocol": protocol, "data_source": data_source}, checkpoint)
+        metrics = ROOT / config["paths"]["training_metrics"]
+        metrics.parent.mkdir(parents=True, exist_ok=True)
+        metrics.write_text(json.dumps({"history": history, "protocol": protocol,
+                                       "data_source": data_source}, indent=2) + "\n")
+        print(f"checkpoint={checkpoint}")
+    if distributed:
+        dist.destroy_process_group()
 
 
 if __name__ == "__main__":