File size: 4,428 Bytes
5c14105 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | 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()
|