File size: 3,634 Bytes
8bf8ea0 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | 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()
|