File size: 5,045 Bytes
bdccc62 | 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 117 118 119 120 121 122 123 124 125 126 127 128 129 | from __future__ import annotations
import argparse
import csv
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
import sys
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from scripts.common import load_config, write_json # noqa: E402
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Create CFO Burgers result visualizations")
parser.add_argument("--config", default=str(PROJECT_ROOT / "config" / "config.yaml"))
return parser.parse_args()
def plot_history(history_path: Path, out_path: Path) -> None:
steps, losses, val_steps, vals = [], [], [], []
with history_path.open("r", encoding="utf-8") as f:
for row in csv.DictReader(f):
steps.append(int(row["step"]))
losses.append(float(row["loss"]))
if row["val_rel_l2"]:
val_steps.append(int(row["step"]))
vals.append(float(row["val_rel_l2"]))
fig, ax1 = plt.subplots(figsize=(8, 4.5))
ax1.plot(steps, losses, color="#1f77b4", linewidth=1.0, label="train loss")
ax1.set_yscale("log")
ax1.set_xlabel("training step")
ax1.set_ylabel("flow matching MSE")
if vals:
ax2 = ax1.twinx()
ax2.plot(val_steps, vals, color="#d62728", marker="o", linewidth=1.2, label="val rel L2")
ax2.set_ylabel("validation trajectory relative L2")
fig.tight_layout()
fig.savefig(out_path, dpi=180)
plt.close(fig)
def plot_worst_case(npz_path: Path, metrics_path: Path, out_path: Path) -> None:
data = np.load(npz_path)
pred = data["prediction"]
target = data["target"]
metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
idx = int(metrics["worst_sample_index"])
times = [25, 50, 75, 100]
fig, axes = plt.subplots(3, len(times), figsize=(13, 6), sharex=True, sharey=True)
for col, t in enumerate(times):
axes[0, col].plot(target[idx, t - 1], color="#222222", linewidth=1.2)
axes[0, col].set_title(f"target t={t}")
axes[1, col].plot(pred[idx, t - 1], color="#0072B2", linewidth=1.2)
axes[1, col].set_title(f"prediction t={t}")
axes[2, col].plot(pred[idx, t - 1] - target[idx, t - 1], color="#D55E00", linewidth=1.0)
axes[2, col].set_title("error")
for ax in axes.flat:
ax.grid(True, linewidth=0.3, alpha=0.4)
fig.suptitle(f"Worst test trajectory index {idx}, rel L2={metrics['per_sample_relative_l2_max']:.6f}")
fig.tight_layout()
fig.savefig(out_path, dpi=180)
plt.close(fig)
def plot_spacetime(npz_path: Path, metrics_path: Path, out_path: Path) -> None:
data = np.load(npz_path)
pred = data["prediction"]
target = data["target"]
metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
idx = int(metrics["worst_sample_index"])
err = pred[idx] - target[idx]
vmax = max(abs(target[idx]).max(), abs(pred[idx]).max())
emax = max(abs(err).max(), 1e-8)
fig, axes = plt.subplots(1, 3, figsize=(12, 4), constrained_layout=True)
ims = [
axes[0].imshow(target[idx], aspect="auto", origin="lower", cmap="viridis", vmin=-vmax, vmax=vmax),
axes[1].imshow(pred[idx], aspect="auto", origin="lower", cmap="viridis", vmin=-vmax, vmax=vmax),
axes[2].imshow(err, aspect="auto", origin="lower", cmap="coolwarm", vmin=-emax, vmax=emax),
]
for ax, title in zip(axes, ["target", "prediction", "prediction - target"]):
ax.set_title(title)
ax.set_xlabel("space index")
ax.set_ylabel("time index")
for ax, im in zip(axes, ims):
fig.colorbar(im, ax=ax, shrink=0.8)
fig.savefig(out_path, dpi=180)
plt.close(fig)
def main() -> None:
args = parse_args()
cfg = load_config(args.config)
results = Path(cfg["paths"]["results_dir"])
figures = results / "figures"
figures.mkdir(parents=True, exist_ok=True)
history = results / "train_history.csv"
pred_npz = results / "test_predictions.npz"
metrics_json = results / "test_metrics.json"
outputs = {}
if history.exists():
outputs["training_curve"] = str(figures / "training_curve.png")
plot_history(history, Path(outputs["training_curve"]))
if pred_npz.exists() and metrics_json.exists():
outputs["worst_case_profiles"] = str(figures / "worst_case_profiles.png")
outputs["worst_case_spacetime"] = str(figures / "worst_case_spacetime.png")
plot_worst_case(pred_npz, metrics_json, Path(outputs["worst_case_profiles"]))
plot_spacetime(pred_npz, metrics_json, Path(outputs["worst_case_spacetime"]))
report = {
"status": "success",
"figures": outputs,
"metrics": json.loads(metrics_json.read_text(encoding="utf-8")) if metrics_json.exists() else None,
}
write_json(results / "result_report.json", report)
print(report, flush=True)
if __name__ == "__main__":
main()
|