File size: 2,159 Bytes
6c3f19f | 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 | import argparse
import json
import sys
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from scripts.common import load_config, resolve_path
def parse_args():
parser = argparse.ArgumentParser(description="Summarize GP_for_TO inference outputs.")
parser.add_argument("--output-dir", default=None)
parser.add_argument("--no-plot", action="store_true")
return parser.parse_args()
def save_plot(data, output_dir):
x = data["x"]
fields = [("u", data["u"]), ("v", data["v"]), ("p", data["p"]), ("ro", data["ro"])]
fig, axes = plt.subplots(2, 2, figsize=(9, 7))
for ax, (name, values) in zip(axes.reshape(-1), fields):
im = ax.tricontourf(x[:, 0], x[:, 1], values, levels=32, cmap="viridis")
ax.set_title(name)
ax.set_xlabel("x")
ax.set_ylabel("y")
fig.colorbar(im, ax=ax)
fig.tight_layout()
path = output_dir / "field_summary.png"
fig.savefig(path, dpi=160)
plt.close(fig)
return path
def main():
args = parse_args()
cfg = load_config()
output_dir = resolve_path(args.output_dir or cfg["inference"]["output_dir"])
pred_path = output_dir / "predictions.npz"
summary_path = output_dir / "inference_summary.json"
if not pred_path.is_file():
raise FileNotFoundError(f"Missing inference output: {pred_path}")
data = np.load(pred_path)
print(f"Prediction file: {pred_path}")
for name in ("x", "u", "v", "p", "ro"):
arr = data[name]
print(
f"{name}: shape={arr.shape}, dtype={arr.dtype}, "
f"min={float(arr.min()):.6e}, max={float(arr.max()):.6e}, mean={float(arr.mean()):.6e}"
)
if summary_path.is_file():
summary = json.loads(summary_path.read_text(encoding="utf-8"))
print(f"Problem: {summary.get('problem')}, checkpoint source: {summary.get('checkpoint_metadata', {}).get('problem')}")
if not args.no_plot:
plot_path = save_plot(data, output_dir)
print(f"Plot: {plot_path}")
if __name__ == "__main__":
main()
|