Upload folder using huggingface_hub

README.md

@@ -13,9 +13,8 @@ tags:
 - arxiv:2309.15214
 tasks: []
 datasets:
-- OneScience/ERA5
+- OneScience-Group/ERA5
 ---
-
 <p align="center">
   <strong>
     <span style="font-size: 30px;">CorrDiff</span>
@@ -32,17 +31,17 @@ https://arxiv.org/abs/2309.15214
 
 # Model Description
 
-CorrDiff was proposed by NVIDIA and its collaborators for kilometer-scale regional weather downscaling and was trained with coarse-resolution ERA5 reanalysis and high-resolution WRF regional model data from Taiwan's Central Weather Administration.
+CorrDiff was proposed by NVIDIA and its collaborators and was trained with coarse-resolution ERA5 reanalysis and high-resolution WRF regional model data from Taiwan's Central Weather Administration.
 The model is suitable for converting coarse-resolution global weather fields into high-resolution regional weather fields and producing probabilistic weather predictions.
 
 # Use Cases
 
 | Scenario | Description |
 | :---: | :--- |
-| Weather forecast training | Train CorrDiff with ERA5 HDF5 data. |
-| Local quick validation | Use synthetic data to check data loading, model training, inference, and result visualization. |
-| ModelScope / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
-| Multi-GPU training | Launch multi-process training with `torchrun`. |
+| Regional weather downscaling training | Train CorrDiff with time-aligned coarse-resolution ERA5 inputs and high-resolution CWA-WRF targets. |
+| Local quick validation | Use synthetic paired data to check data loading, two-stage training, ensemble inference, and result visualization. |
+| Hugging Face / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Launch multi-process training with `torchrun` after adapting the real-data pipeline for distributed training. |
 
 # Usage Guide
 
@@ -91,18 +90,28 @@ pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simp
 
 ### Training Data Introduction
 
-The OneScience community provides an ERA5 data slice that can be downloaded as follows:
+This repository uses synthetic paired data by default to validate training, ensemble inference, and evaluation:
 
 ```bash
-hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data
+python scripts/fake_data.py --output data/era5_corrdiff.npz
 ```
 
-When real data is unavailable, generate synthetic data for pipeline validation:
+The synthetic file contains `input` with shape `[N, 12, 36, 36]` and `target` with shape `[N, 4, 448, 448]`. It is intended for pipeline validation only and does not represent real weather predictions.
+
+ERA5 inputs can be downloaded from the OneScience community:
 
 ```bash
-python scripts/fake_data.py --output data/era5_corrdiff.npz
+hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data/era5
+```
+
+The paired CWA-WRF resource can be downloaded from NVIDIA NGC:
+
+```bash
+ngc registry resource download-version "nvidia/modulus/modulus_datasets_cwa:v1"
 ```
 
+ERA5 alone is insufficient for supervised CorrDiff training. Real data must be time-aligned and preprocessed into the paired `input` and `target` interface expected by the scripts. Verify the data path in `conf/config.yaml` before training.
+
 ### Training
 
 Single GPU:
@@ -120,14 +129,16 @@ torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_
 Training outputs:
 
 ```text
-data/checkpoints/model_bak.pth
-data/checkpoints/trloss.npy
-data/checkpoints/valoss.npy
+data/checkpoints/regression_model.pth
+data/checkpoints/diffusion_model.pth
+data/checkpoints/training_history.npz
 ```
 
+Training first fits the conditional-mean regression model and then freezes it while fitting the residual diffusion model. `training_history.npz` stores the regression and diffusion denoising losses.
+
 ### Training Weights
 
-This repository provides weights trained on 39 years of ERA5 reanalysis data in the `weight/` folder. The weight files will be uploaded soon and are expected to be available in the near future.
+The training command generates `regression_model.pth` and `diffusion_model.pth` from `data/era5_corrdiff.npz`. Weights produced with synthetic data validate the training and inference pipeline only and do not provide real weather forecasting skill. This repository does not present them as pretrained weights reproducing the paper.
 
 ### Inference
 
@@ -135,7 +146,13 @@ This repository provides weights trained on 39 years of ERA5 reanalysis data in
 python scripts/inference.py
 ```
 
-Prediction results are written to `result/output/` by default.
+Prediction output:
+
+```text
+result/output/predictions.npz
+```
+
+The file stores the ensemble predictions, ensemble mean, ensemble standard deviation, input, and target generated by the current trained weights.
 
 ### Evaluation and Visualization
 
@@ -143,6 +160,15 @@ Prediction results are written to `result/output/` by default.
 python scripts/result.py
 ```
 
+Evaluation outputs:
+
+```text
+result/output/metrics.json
+result/output/prediction_comparison.png
+```
+
+`metrics.json` stores MAE and CRPS for the four output variables. `prediction_comparison.png` compares the ensemble mean, ensemble standard deviation, and target. Metrics and figures generated from synthetic data are pipeline checks, not paper results.
+
 # Official OneScience Resources
 
 | Platform | OneScience Main Repository | Skills Repository |

config.json

@@ -9,42 +9,27 @@
   "task": "regional-weather-downscaling",
   "implementation": {
     "entry_point": "model/corrdiff.py",
-    "scope": "compact two-stage reproduction of the CorrDiff regression and residual diffusion architecture"
+    "scope": "compact two-stage CorrDiff reproduction"
   },
   "architecture": {
     "family": "residual corrective diffusion model",
-    "input_grid_shape": [
-      36,
-      36
-    ],
-    "output_grid_shape": [
-      448,
-      448
-    ],
+    "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",
-    "conditioning": [
-      "bilinearly upsampled coarse input",
-      "high-resolution regression mean",
-      "noise level sigma"
-    ],
     "normalization": "group_norm",
     "activation": "silu",
-    "sampling_steps": 4,
-    "sigma_max": 1.0,
-    "sigma_min": 0.01
+    "sampling_steps": 4
   },
   "data": {
     "dataset": "ERA5 and CWA-WRF",
     "input_spatial_resolution_km": 25,
     "output_spatial_resolution_km": 2,
     "time_step_hours": 1,
-    "input_steps": 1,
-    "output_steps": 1,
-    "protocol": "synthetic_era5_corrdiff"
+    "protocol": "paired-downscaling"
   },
   "configuration_sources": [
     "conf/config.yaml",

scripts/inference.py

@@ -1,35 +1,59 @@
 import argparse
 from pathlib import Path
+import sys
+
 import numpy as np
 import torch
 
-import sys
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 from model.corrdiff import CorrDiff
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Run CorrDiff inference")
+    parser = argparse.ArgumentParser(description="Generate a CorrDiff ensemble")
     parser.add_argument("--data", default="data/era5_corrdiff.npz")
-    parser.add_argument("--checkpoint", default="data/checkpoints/model_bak.pth")
-    parser.add_argument("--ensemble-size", type=int, default=1)
+    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)
     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()
-    checkpoint = Path(args.checkpoint)
-    if checkpoint.exists():
-        model.load_state_dict(torch.load(checkpoint, map_location="cpu")["model"])
+    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"])
     model.eval()
+
     samples = []
     with torch.no_grad():
-        for _ in range(args.ensemble_size):
+        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)
     output.parent.mkdir(parents=True, exist_ok=True)
-    np.savez_compressed(output, prediction=np.stack(samples), input=data["input"])
-    print(f"prediction: {np.stack(samples).shape}")
+    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}")
 
 

scripts/result.py

@@ -1,20 +1,65 @@
 import argparse
+import json
 from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
 import numpy as np
 
 
+VARIABLES = ["t2m", "u10m", "v10m", "radar_reflectivity"]
+
+
+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
+
+
 def main():
-    parser = argparse.ArgumentParser(description="Evaluate CorrDiff NPZ predictions")
+    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")
     args = parser.parse_args()
-    prediction = np.load(args.prediction)["prediction"]
-    target = np.load(args.target)["target"]
-    mae = np.abs(prediction.mean(axis=0) - target).mean()
-    output = Path(args.prediction).with_name("metrics.npz")
-    np.savez(output, mae=np.array(mae, dtype="float32"))
-    print(f"ensemble={prediction.shape[0]} mae={mae:.6f}")
-    print(f"saved: {output}")
+
+    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)
+    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}")
 
 
 if __name__ == "__main__":

scripts/train.py

@@ -10,34 +10,65 @@ from model.corrdiff import CorrDiff
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Train compact CorrDiff on NPZ data")
+    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("--output", default="data/checkpoints/model_bak.pth")
+    parser.add_argument("--checkpoint-dir", default="data/checkpoints")
     args = parser.parse_args()
     data = np.load(args.data)
     coarse = torch.from_numpy(data["input"])
     target = torch.from_numpy(data["target"])
     model = CorrDiff()
-    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)
+    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 = []
+
+    # 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 2: freeze the mean model and learn its stochastic residual.
+    model.regression.eval()
+    for parameter in model.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.detach(), sigma)
-        loss = F.mse_loss(mean, target) + F.mse_loss(predicted, residual)
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-        print(f"step={step + 1} loss={loss.item():.6f}")
-    output = Path(args.output)
-    output.parent.mkdir(parents=True, exist_ok=True)
-    torch.save({"model": model.state_dict(), "format": "corrdiff-compact-v1"}, output)
-    np.save(output.parent / "trloss.npy", np.asarray([loss.item()], dtype="float32"))
-    np.save(output.parent / "valoss.npy", np.asarray([loss.item()], dtype="float32"))
-    print(f"saved: {output}")
+        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}")
+
+    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 __name__ == "__main__":