OneScience-Group/NNCAM

1#!/usr/bin/env python32import argparse3import json4from pathlib import Path5 6import matplotlib7matplotlib.use("Agg")8import matplotlib.pyplot as plt9import numpy as np10 11 12ROOT = Path(__file__).resolve().parents[1]13 14 15def score(truth, prediction):16    rmse = float(np.sqrt(np.mean((truth - prediction) ** 2)))17    denominator = float(np.sum((truth - truth.mean()) ** 2))18    return {"rmse": rmse, "r2": float(1.0 - np.sum((truth - prediction) ** 2) / denominator) if denominator else None}19 20 21def main():22    parser = argparse.ArgumentParser(description="Evaluate NNCAM predictions.")23    parser.add_argument("--input", type=Path, default=ROOT / "result/output/predictions.npz")24    parser.add_argument("--metrics", type=Path, default=ROOT / "result/evaluation/metrics.json")25    parser.add_argument("--figure", type=Path, default=ROOT / "result/evaluation/comparison.png")26    args = parser.parse_args()27    with np.load(args.input) as data:28        truth, prediction = data["truth"], data["prediction"]29    if truth.shape != prediction.shape or truth.ndim != 2 or truth.shape[1] != 65:30        raise ValueError(f"expected matching [N,65] arrays, got {truth.shape}, {prediction.shape}")31    groups = {"dT": slice(0, 30), "dQ": slice(30, 60), "SW": slice(60, 62), "LW": slice(62, 64), "P": slice(64, 65)}32    metrics = {name: score(truth[:, indices], prediction[:, indices]) for name, indices in groups.items()}33    metrics["overall"] = score(truth, prediction)34    values = np.array([value for group in metrics.values() for value in group.values() if value is not None])35    if not np.isfinite(values).all():36        raise RuntimeError("evaluation metrics contain non-finite values")37    args.metrics.parent.mkdir(parents=True, exist_ok=True)38    args.metrics.write_text(json.dumps(metrics, indent=2), encoding="utf-8")39    names = list(groups)40    fig, axes = plt.subplots(1, 2, figsize=(10, 4), constrained_layout=True)41    axes[0].bar(names, [metrics[name]["rmse"] for name in names])42    axes[0].set_title("Grouped RMSE")43    axes[1].scatter(truth[:, 64], prediction[:, 64], s=12, alpha=0.7)44    axes[1].set(xlabel="True precipitation", ylabel="Predicted precipitation", title="Precipitation comparison")45    fig.savefig(args.figure, dpi=150)46    plt.close(fig)47    print(f"saved {args.metrics} and {args.figure}; shape={prediction.shape}, finite=true")48 49 50if __name__ == "__main__":51    main()52