#!/usr/bin/env python3 """Visualize CNO fields and ID/OOD relative-L1 distributions.""" from __future__ import annotations import argparse import json from pathlib import Path import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np PROJECT_ROOT = Path(__file__).resolve().parents[1] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--results-dir", default=str(PROJECT_ROOT / "results")) parser.add_argument( "--sample-index", type=int, default=None, help="array index to plot; default is the sample nearest the split median error", ) parser.add_argument("--dpi", type=int, default=180) return parser.parse_args() def load_predictions(path: Path) -> dict[str, np.ndarray]: if not path.is_file(): raise FileNotFoundError(f"prediction artifact not found: {path}") with np.load(path) as payload: required = {"sample_ids", "inputs", "targets", "predictions", "relative_l1"} missing = sorted(required.difference(payload.files)) if missing: raise KeyError(f"{path} is missing arrays: {missing}") arrays = {key: np.asarray(payload[key]) for key in required} count = arrays["sample_ids"].shape[0] for key in ("inputs", "targets", "predictions", "relative_l1"): if arrays[key].shape[0] != count: raise ValueError(f"sample count mismatch for {key} in {path}") if not np.isfinite(arrays[key]).all(): raise ValueError(f"nonfinite values in {key} from {path}") if arrays["inputs"].ndim != 4 or arrays["inputs"].shape[1] != 1: raise ValueError(f"expected N1HW fields in {path}, got {arrays['inputs'].shape}") return arrays def representative_index(errors: np.ndarray, requested: int | None) -> int: if requested is not None: if requested < 0 or requested >= errors.size: raise IndexError(f"sample-index {requested} outside [0,{errors.size})") return requested median = np.median(errors) return int(np.argmin(np.abs(errors - median))) def plot_fields( split: str, arrays: dict[str, np.ndarray], output_path: Path, requested_index: int | None, dpi: int, ) -> None: errors_percent = arrays["relative_l1"] * 100.0 index = representative_index(errors_percent, requested_index) input_field = arrays["inputs"][index, 0] target = arrays["targets"][index, 0] prediction = arrays["predictions"][index, 0] absolute_error = np.abs(prediction - target) field_min = float(min(target.min(), prediction.min())) field_max = float(max(target.max(), prediction.max())) fig, axes = plt.subplots(1, 4, figsize=(16, 3.8), constrained_layout=True) input_image = axes[0].imshow(input_field, origin="lower", cmap="RdBu_r") fig.colorbar(input_image, ax=axes[0], shrink=0.78) target_image = axes[1].imshow( target, origin="lower", cmap="RdBu_r", vmin=field_min, vmax=field_max ) prediction_image = axes[2].imshow( prediction, origin="lower", cmap="RdBu_r", vmin=field_min, vmax=field_max ) error_image = axes[3].imshow(absolute_error, origin="lower", cmap="magma") fig.colorbar(target_image, ax=[axes[1], axes[2]], shrink=0.78) fig.colorbar(error_image, ax=axes[3], shrink=0.78) titles = ("Initial velocity", "Target at T=1", "CNO prediction", "Absolute error") for axis, title in zip(axes, titles): axis.set_title(title) axis.set_xlabel("x index") axis.set_ylabel("y index") sample_id = int(arrays["sample_ids"][index]) fig.suptitle( f"{split.upper()} Sample_{sample_id} — relative L1={errors_percent[index]:.3f}%" ) fig.savefig(output_path, dpi=dpi) plt.close(fig) def plot_distribution( split_arrays: dict[str, dict[str, np.ndarray]], paper_reference: dict, output_path: Path, dpi: int, ) -> None: fig, axes = plt.subplots(1, 2, figsize=(11, 4), constrained_layout=True) colors = {"id": "#2166ac", "ood": "#b2182b"} for split in ("id", "ood"): values = split_arrays[split]["relative_l1"] * 100.0 axes[0].hist(values, bins=20, alpha=0.55, label=split.upper(), color=colors[split]) axes[0].axvline( np.median(values), color=colors[split], linewidth=2, linestyle="-" ) reference = paper_reference.get(split) if reference is not None: axes[0].axvline( float(reference), color=colors[split], linewidth=1.5, linestyle="--" ) axes[0].set_title("Per-sample relative L1") axes[0].set_xlabel("Relative L1 (%)") axes[0].set_ylabel("Count") axes[0].legend(title="solid=reproduction\ndashed=paper") values = [ split_arrays["id"]["relative_l1"] * 100.0, split_arrays["ood"]["relative_l1"] * 100.0, ] box = axes[1].boxplot(values, tick_labels=["ID", "OOD"], patch_artist=True) for patch, color in zip(box["boxes"], (colors["id"], colors["ood"])): patch.set_facecolor(color) patch.set_alpha(0.55) axes[1].set_title("Error distribution summary") axes[1].set_ylabel("Relative L1 (%)") fig.savefig(output_path, dpi=dpi) plt.close(fig) def main() -> None: args = parse_args() results_dir = Path(args.results_dir).expanduser().resolve() metrics_path = results_dir / "metrics.json" if not metrics_path.is_file(): raise FileNotFoundError(f"metrics file not found: {metrics_path}") with metrics_path.open("r", encoding="utf-8") as handle: metrics = json.load(handle) split_arrays = { split: load_predictions(results_dir / f"{split}_predictions.npz") for split in ("id", "ood") } for split, arrays in split_arrays.items(): output_path = results_dir / f"{split}_fields.png" plot_fields(split, arrays, output_path, args.sample_index, args.dpi) values = arrays["relative_l1"] * 100.0 print( f"visualization split={split} median={np.median(values):.6f}% " f"mean={np.mean(values):.6f}% saved={output_path}", flush=True, ) distribution_path = results_dir / "error_distribution.png" plot_distribution( split_arrays, metrics.get("paper_reference", {}), distribution_path, args.dpi, ) print(f"visualization saved={distribution_path}", flush=True) if __name__ == "__main__": main()