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