File size: 3,357 Bytes
eac1b4e | 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 | 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()
|