| |
| 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() |
|
|
