| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
|
|
| from common import load_config, resolve_output_path, save_json |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default=str(Path(__file__).resolve().parents[1] / "config" / "config.yaml")) |
| parser.add_argument("--predictions", default=None) |
| parser.add_argument("--num-samples", type=int, default=4) |
| return parser.parse_args() |
|
|
|
|
| def rel_l2(pred: np.ndarray, target: np.ndarray) -> np.ndarray: |
| diff = np.linalg.norm((pred - target).reshape(pred.shape[0], -1), axis=1) |
| denom = np.linalg.norm(target.reshape(target.shape[0], -1), axis=1) |
| return diff / np.maximum(denom, 1e-12) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| cfg = load_config(args.config) |
| results_dir = resolve_output_path(cfg, cfg["paths"]["results_dir"]) |
| pred_path = Path(args.predictions) if args.predictions else results_dir / "test_predictions.npz" |
| data = np.load(pred_path) |
| inputs = data["inputs"] |
| preds = data["predictions"] |
| targets = data["targets"] |
| errors = rel_l2(preds, targets) * 100.0 |
|
|
| sample_count = min(args.num_samples, preds.shape[0]) |
| fig, axes = plt.subplots(sample_count, 4, figsize=(13, 3.1 * sample_count), constrained_layout=True) |
| if sample_count == 1: |
| axes = axes[None, :] |
| for i in range(sample_count): |
| panels = [ |
| (inputs[i, 0], "input t=0"), |
| (targets[i, 0], "target t=5"), |
| (preds[i, 0], "prediction"), |
| (np.abs(preds[i, 0] - targets[i, 0]), f"abs error rel={errors[i]:.2f}%"), |
| ] |
| for ax, (arr, title) in zip(axes[i], panels): |
| im = ax.imshow(arr, origin="lower", cmap="RdBu_r" if title != "abs error" else "magma") |
| ax.set_title(title) |
| ax.set_xticks([]) |
| ax.set_yticks([]) |
| fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02) |
| fig_path = results_dir / "prediction_examples.png" |
| fig.savefig(fig_path, dpi=180) |
| plt.close(fig) |
|
|
| hist_fig, ax = plt.subplots(figsize=(6, 4), constrained_layout=True) |
| ax.hist(errors, bins=24, color="#386cb0", edgecolor="white") |
| ax.axvline(errors.mean(), color="#d95f02", linewidth=2, label=f"mean={errors.mean():.2f}%") |
| ax.axvline(cfg["experiment"]["paper_reference_test_percent"], color="#1b9e77", linewidth=2, linestyle="--", label="paper CROP 4.06%") |
| ax.set_xlabel("relative L2 (%)") |
| ax.set_ylabel("test samples") |
| ax.legend() |
| hist_path = results_dir / "relative_l2_histogram.png" |
| hist_fig.savefig(hist_path, dpi=180) |
| plt.close(hist_fig) |
|
|
| save_json( |
| results_dir / "result_summary.json", |
| { |
| "status": "complete", |
| "prediction_file": str(pred_path), |
| "mean_relative_l2_percent": float(errors.mean()), |
| "std_relative_l2_percent": float(errors.std(ddof=0)), |
| "min_relative_l2_percent": float(errors.min()), |
| "max_relative_l2_percent": float(errors.max()), |
| "figures": [str(fig_path), str(hist_path)], |
| "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"], |
| "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"], |
| }, |
| ) |
| print( |
| f"result_complete mean_relative_l2={errors.mean():.6f}% " |
| f"std={errors.std(ddof=0):.6f}% figures={[str(fig_path), str(hist_path)]}", |
| flush=True, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
