| 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 ROOT, load_config, resolve, save_json |
|
|
|
|
| def plot_case(xy: np.ndarray, pred: np.ndarray, target: np.ndarray, out_path: Path, title: str) -> None: |
| x = xy[..., 0] |
| y = xy[..., 1] |
| p = pred[..., 0] |
| t = target[..., 0] |
| err = p - t |
| vmin = float(min(p.min(), t.min())) |
| vmax = float(max(p.max(), t.max())) |
| fig, axes = plt.subplots(1, 3, figsize=(13, 3.6), constrained_layout=True) |
| for ax, field, name, lo, hi, cmap in [ |
| (axes[0], t, "target", vmin, vmax, "coolwarm"), |
| (axes[1], p, "prediction", vmin, vmax, "coolwarm"), |
| (axes[2], err, "error", float(-np.max(np.abs(err))), float(np.max(np.abs(err))), "coolwarm"), |
| ]: |
| mesh = ax.pcolormesh(x, y, field, shading="auto", cmap=cmap, vmin=lo, vmax=hi) |
| ax.set_aspect("equal", adjustable="box") |
| ax.set_title(name) |
| ax.set_xticks([]) |
| ax.set_yticks([]) |
| fig.colorbar(mesh, ax=ax, fraction=0.046, pad=0.02) |
| fig.suptitle(title) |
| fig.savefig(out_path, dpi=180) |
| plt.close(fig) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Create HO-FNO Pipe visualizations from inference arrays.") |
| parser.add_argument("--config", default=str(ROOT / "config" / "config.yaml")) |
| parser.add_argument("--num-cases", type=int, default=8) |
| args = parser.parse_args() |
|
|
| cfg = load_config(args.config) |
| output_dir = resolve(cfg["paths"]["output_dir"]) |
| results_dir = resolve(cfg["paths"]["results_dir"], output_dir) |
| pred_path = results_dir / "predictions.npy" |
| target_path = results_dir / "targets.npy" |
| input_path = results_dir / "inputs.npy" |
| rel_path = results_dir / "relative_l2_per_sample.npy" |
| for path in [pred_path, target_path, input_path, rel_path]: |
| if not path.exists(): |
| raise FileNotFoundError(f"missing required inference artifact: {path}") |
|
|
| pred = np.load(pred_path) |
| target = np.load(target_path) |
| inputs = np.load(input_path) |
| rel = np.load(rel_path) |
| fig_dir = results_dir / "figures" |
| fig_dir.mkdir(parents=True, exist_ok=True) |
| order = np.argsort(rel) |
| chosen = [] |
| if len(order) > 0: |
| chosen.extend([int(order[0]), int(order[len(order) // 2]), int(order[-1])]) |
| for idx in range(min(args.num_cases, len(rel))): |
| if idx not in chosen: |
| chosen.append(idx) |
| chosen = chosen[: max(1, args.num_cases)] |
| figures = [] |
| for idx in chosen: |
| out_path = fig_dir / f"pipe_case_{idx:03d}.png" |
| plot_case(inputs[idx], pred[idx], target[idx], out_path, f"Pipe case {idx} | rel L2={rel[idx]:.6f}") |
| figures.append(str(out_path)) |
|
|
| summary = { |
| "figure_count": len(figures), |
| "figures": figures, |
| "mean_relative_l2": float(rel.mean()), |
| "best_case": int(order[0]) if len(order) else None, |
| "median_case": int(order[len(order) // 2]) if len(order) else None, |
| "worst_case": int(order[-1]) if len(order) else None, |
| } |
| save_json(results_dir / "result_summary.json", summary) |
| print(f"result_complete figures={len(figures)} summary={results_dir / 'result_summary.json'}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
