| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import yaml |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def load_config(path: str | Path) -> dict[str, Any]: |
| with open(path, "r", encoding="utf-8") as handle: |
| return yaml.safe_load(handle) |
|
|
|
|
| def resolve(path: str | Path) -> Path: |
| path = Path(path) |
| return path if path.is_absolute() else PROJECT_ROOT / path |
|
|
|
|
| def load_json(path: Path) -> dict[str, Any]: |
| with open(path, "r", encoding="utf-8") as handle: |
| return json.load(handle) |
|
|
|
|
| def save_json(path: Path, payload: dict[str, Any]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with open(path, "w", encoding="utf-8") as handle: |
| json.dump(payload, handle, indent=2) |
|
|
|
|
| def plot_sample(sample_id: int, inputs: np.ndarray, targets: np.ndarray, preds: np.ndarray, channel_names: list[str], out_dir: Path) -> Path: |
| fig, axes = plt.subplots(4, len(channel_names), figsize=(3.2 * len(channel_names), 10.0), constrained_layout=True) |
| rows = [("Boussinesq input", inputs), ("Compressible target", targets), ("FNO prediction", preds), ("Abs error", np.abs(preds - targets))] |
| for row, (title, array) in enumerate(rows): |
| for col, channel in enumerate(channel_names): |
| ax = axes[row, col] |
| image = array[sample_id, col] |
| vmax = 1.0 if row < 3 else max(float(image.max()), 1.0e-8) |
| im = ax.imshow(image, origin="lower", cmap="viridis" if row < 3 else "magma", vmin=0.0, vmax=vmax) |
| ax.set_xticks([]) |
| ax.set_yticks([]) |
| if row == 0: |
| ax.set_title(channel) |
| if col == 0: |
| ax.set_ylabel(title) |
| if row == 3: |
| mse = float(np.mean((preds[sample_id, col] - targets[sample_id, col]) ** 2)) |
| ax.set_xlabel(f"MSE={mse:.2e}") |
| fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| path = out_dir / f"sample_{sample_id:03d}.png" |
| fig.savefig(path, dpi=160) |
| plt.close(fig) |
| return path |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Summarize metrics and render NeuralConvection FNO figures.") |
| parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml")) |
| args = parser.parse_args() |
|
|
| config = load_config(args.config) |
| viz_cfg = config["visualization"] |
| predictions = np.load(resolve(viz_cfg["predictions_path"])) |
| preds = predictions["pred_norm"] |
| targets = predictions["target_norm"] |
| inputs = predictions["input_norm"] |
| channel_names = [str(x) for x in predictions["channel_names"].tolist()] |
| metrics = load_json(resolve(viz_cfg["metrics_path"])) |
| out_dir = resolve(viz_cfg["output_dir"]) |
| figures = [] |
| for sample_id in range(min(int(viz_cfg["max_samples"]), preds.shape[0])): |
| figures.append(str(plot_sample(sample_id, inputs, targets, preds, channel_names, out_dir))) |
|
|
| comparison = {} |
| paper_channels = config["paper_table5_compact_fno_2d"]["channels"] |
| for name in channel_names: |
| current = metrics["per_channel"][name] |
| paper = paper_channels[name] |
| comparison[name] = { |
| "mse": current["mse"], |
| "paper_mse": paper["mse"], |
| "mse_ratio_to_paper": current["mse"] / paper["mse"], |
| "ssim": current["ssim"], |
| "paper_ssim": paper["ssim"], |
| "ssim_delta": current["ssim"] - paper["ssim"], |
| } |
| summary = { |
| "status": "success", |
| "predictions_path": str(resolve(viz_cfg["predictions_path"])), |
| "metrics_path": str(resolve(viz_cfg["metrics_path"])), |
| "figures": figures, |
| "table5_comparison": comparison, |
| } |
| save_json(resolve(viz_cfg["summary_path"]), summary) |
| print(f"result_summary={resolve(viz_cfg['summary_path'])}", flush=True) |
| for channel, values in comparison.items(): |
| print( |
| f"{channel}: mse={values['mse']:.8e} paper={values['paper_mse']:.8e} " |
| f"ratio={values['mse_ratio_to_paper']:.3f} ssim={values['ssim']:.6f} paper_ssim={values['paper_ssim']:.6f}", |
| flush=True, |
| ) |
| print(f"figures={len(figures)} saved_to={out_dir}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
