| """Create scientific figures and aggregate ten independent Table 3 runs.""" |
| import argparse |
| import csv |
| import json |
| import shutil |
| import sys |
| from pathlib import Path |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import h5py |
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| from common import ROOT, atomic_json, load_config, resolved |
|
|
|
|
| def aggregate(config): |
| results, weights = resolved(config, "results"), resolved(config, "weights") |
| records = [] |
| for seed in config["training"]["seeds"]: |
| directory = results/f"seed_{seed}" |
| summary = json.loads((directory/"training_summary.json").read_text()) |
| if summary["completed_epochs"] != 20000: raise ValueError(f"Incomplete seed {seed}") |
| best = json.loads((directory/"best_eval/metrics.json").read_text()) |
| final = json.loads((directory/"final_eval/metrics.json").read_text()) |
| records.append({"seed": seed, "best": best, "final": final, "training": summary}) |
| winner = min(records, key=lambda row: row["best"]["selection_score"]) |
| shutil.copy2(weights/f"seed_{winner['seed']}"/"best_model.pt", weights/"best_model.pt") |
| atomic_json(results/"best_checkpoint_selection.json", {"seed": winner["seed"], |
| "epoch": winner["best"]["epoch"], "selection_score": winner["best"]["selection_score"], |
| "metrics": winner["best"]["validation_metrics"], |
| "checkpoint": str(weights/"best_model.pt"), "phase": "all_ten_formal_runs_completed"}) |
| reference = {"u": {"mean": 25.73, "std": 2.49}, "v": {"mean": 38.00, "std": 3.76}, "p": {"mean": 42.65, "std": 8.69}} |
| statistics = {} |
| for kind in ["final", "best"]: |
| statistics[kind] = {} |
| for field in ("u", "v", "p"): |
| values = np.array([r[kind]["groups"]["cavity_internal"]["metrics"][field]["relative_l2_percent"] for r in records]) |
| statistics[kind][field] = {"mean_percent": float(values.mean()), "std_percent": float(values.std(ddof=0)), "values_percent": values.tolist()} |
| losses = np.array([r["training"]["last_loss"]["total"] for r in records]) |
| run_step_times = [] |
| for row in records: |
| with open(results/f"seed_{row['seed']}"/"history.csv") as stream: |
| times = [float(item["step_seconds"]) for item in csv.DictReader(stream)] |
| if len(times) != 20000: |
| raise ValueError(f"Expected exactly 20000 recorded updates for seed {row['seed']}") |
| run_step_times.append(float(np.mean(times))) |
| peak_rss = np.array([r["training"]["peak_rss_mb"] for r in records]) |
| report = {"status": "completed", "runs": len(records), "epochs_per_run": 20000, |
| "statistics": statistics, "std_ddof": 0, "paper_table3_percent": reference, |
| "final_loss": {"mean": float(losses.mean()), "std": float(losses.std())}, |
| "parameters": 923, |
| "time_per_update_seconds": {"mean": float(np.mean(run_step_times)), "std": float(np.std(run_step_times)), "run_means": run_step_times, "definition": "wall time for PDE/BC/IC forward, backward, gradient clip and Adam step; validation/checkpoint I/O excluded"}, |
| "peak_process_rss_mb": {"mean": float(peak_rss.mean()), "std": float(peak_rss.std()), "values": peak_rss.tolist()}, |
| "paper_table3_other": {"parameters": 923, "loss_mean": 0.10, "loss_std": 0.02, "time_seconds_mean": 0.36, "time_seconds_std": 0.01, "memory_mb_mean": 540.0, "memory_mb_std": 0.0}, |
| "selected_seed": winner["seed"], "selection_metric": "mean validation relative L2", |
| "best_model": str(weights/"best_model.pt"), |
| "limitations": config["assumptions"]} |
| atomic_json(results/"table3_comparison.json", report) |
| print(json.dumps(report, indent=2), flush=True) |
| return report |
|
|
|
|
| def plot_history(config, seed, destination): |
| path = resolved(config, "results")/f"seed_{seed}"/"history.csv" |
| with open(path) as stream: rows = list(csv.DictReader(stream)) |
| epochs = np.array([int(r["epoch"]) for r in rows]) |
| fig, axes = plt.subplots(1, 2, figsize=(12, 4.2), constrained_layout=True) |
| for key in ("total", "physics", "lid", "walls", "initial"): |
| axes[0].semilogy(epochs, [max(float(r[key]), 1e-30) for r in rows], label=key, linewidth=.8) |
| axes[0].set(xlabel="Epoch (one Adam update)", ylabel="Loss", title=f"Angle-Cascade, seed={seed}") |
| valid = [r for r in rows if r["val_u_l2"]] |
| for key in ("u", "v", "p"): |
| axes[1].plot([int(r["epoch"]) for r in valid], [100*float(r[f"val_{key}_l2"]) for r in valid], label=key) |
| axes[1].set(xlabel="Epoch", ylabel="Validation relative L2 (%)", title="Raw u/v/p, no pressure shift") |
| for ax in axes: ax.grid(alpha=.2); ax.legend() |
| target = destination/"training_curves.png" |
| fig.savefig(target, dpi=config["visualization"]["dpi"]) |
| plt.close(fig) |
| return str(target) |
|
|
|
|
| def plot_fields(config, predictions, destination): |
| images = [] |
| with h5py.File(resolved(config, "reference"), "r") as source, h5py.File(predictions, "r") as inference: |
| key = "cavity_internal" |
| raw = source[key] |
| checkpoint_meta = json.loads(inference.attrs["provenance"]) |
| actual_times = np.unique(raw[0, :]) |
| for desired in config["visualization"]["times"]: |
| t = float(actual_times[np.argmin(abs(actual_times-desired))]) |
| ids = np.flatnonzero(np.isclose(raw[0, :], t, atol=1e-8, rtol=0)) |
| rows = raw[:, ids].T |
| pred = inference[key]["prediction"][ids] |
| xs, ys = np.unique(rows[:, 1]), np.unique(rows[:, 2]) |
| if len(xs)*len(ys) != len(rows): raise ValueError("Reference snapshot is not a complete rectangular grid") |
| order = np.lexsort((rows[:, 1], rows[:, 2])) |
| truth = rows[order, 3:].reshape(len(ys), len(xs), 3) |
| predicted = pred[order].reshape(len(ys), len(xs), 3) |
| fig, axes = plt.subplots(3, 3, figsize=(11, 10), constrained_layout=True) |
| for index, field in enumerate(("u", "v", "p")): |
| low, high = min(truth[..., index].min(), predicted[..., index].min()), max(truth[..., index].max(), predicted[..., index].max()) |
| if low == high: high = low+1e-12 |
| for col, values in enumerate((truth[..., index], predicted[..., index], abs(predicted[..., index]-truth[..., index]))): |
| cmap = "coolwarm" if index == 2 and col < 2 else "viridis" if col < 2 else "magma" |
| kwargs = {"vmin": low, "vmax": high} if col < 2 else {"vmin": 0} |
| artist = axes[index, col].pcolormesh(xs, ys, values, shading="auto", cmap=cmap, **kwargs) |
| axes[index, col].set(xlabel="x (m)", ylabel="y (m)", title=f"{field}: {['Reference', 'QCPINN', 'Absolute error'][col]}") |
| axes[index, col].set_aspect("equal") |
| fig.colorbar(artist, ax=axes[index, col], shrink=.8) |
| fig.suptitle(f"Lid-driven cavity, t={t:g} s | seed={checkpoint_meta['seed']}, epoch={checkpoint_meta['epoch']} | raw fields") |
| target = destination/f"fields_t{t:g}.png" |
| fig.savefig(target, dpi=config["visualization"]["dpi"]) |
| plt.close(fig) |
| images.append(str(target)) |
| fig, axes = plt.subplots(1, 2, figsize=(10, 4.5), constrained_layout=True) |
| for ax, data, title in zip(axes, (truth, predicted), ("Reference", "QCPINN")): |
| speed = np.linalg.norm(data[..., :2], axis=-1) |
| artist = ax.streamplot(xs, ys, data[..., 0], data[..., 1], color=speed, cmap="viridis", density=1.1) |
| fig.colorbar(artist.lines, ax=ax, label="Speed (m/s)") |
| ax.set(xlabel="x (m)", ylabel="y (m)", title=f"{title}, t={t:g} s", xlim=(0, 1), ylim=(0, 1)) |
| ax.set_aspect("equal") |
| target = destination/f"streamlines_t{t:g}.png" |
| fig.savefig(target, dpi=config["visualization"]["dpi"]) |
| plt.close(fig) |
| images.append(str(target)) |
| return images |
|
|
|
|
| def make_results(config, seed=None, output_dir=None, predictions=None): |
| results = resolved(config, "results") |
| comparison = None |
| if seed is None: |
| comparison = json.loads((results/"table3_comparison.json").read_text()) |
| seed = comparison["selected_seed"] |
| evaluation = results/"best_model" |
| else: evaluation = results/f"seed_{seed}"/"best_eval" |
| output = Path(output_dir) if output_dir else evaluation/"figures" |
| output.mkdir(parents=True, exist_ok=True) |
| paths = [plot_history(config, seed, output)] |
| paths += plot_fields(config, Path(predictions) if predictions else evaluation/"predictions.h5", output) |
| metrics = json.loads((evaluation/"metrics.json").read_text()) |
| lines = ["# QCPINN Angle-Cascade cavity reproduction", "", f"Selected seed: {seed}; epoch: {metrics['epoch']}; 923 parameters.", |
| "", "Reference values were read directly from the supplied cavity.mat; no generated reference truth.", |
| "", "| Field | Internal relative L2 (%) | Paper Table 3 mean ± std (%) |", "|---|---:|---:|"] |
| for k, paper in zip(("u", "v", "p"), ("25.73 ± 2.49", "38.00 ± 3.76", "42.65 ± 8.69")): |
| value = metrics["groups"]["cavity_internal"]["metrics"][k]["relative_l2_percent"] |
| lines.append(f"| {k} | {value:.6f} | {paper} |") |
| if comparison is not None: |
| lines += ["", "## Ten independent complete runs (20000 epochs each)", "", |
| "| Field | Final checkpoints mean ± std (%) | Validation-best checkpoints mean ± std (%) | Paper Table 3 (%) |", |
| "|---|---:|---:|---:|"] |
| for field in ("u", "v", "p"): |
| final = comparison["statistics"]["final"][field] |
| best = comparison["statistics"]["best"][field] |
| paper = comparison["paper_table3_percent"][field] |
| lines.append(f"| {field} | {final['mean_percent']:.6f} ± {final['std_percent']:.6f} | {best['mean_percent']:.6f} ± {best['std_percent']:.6f} | {paper['mean']:.2f} ± {paper['std']:.2f} |") |
| loss = comparison["final_loss"] |
| timing = comparison["time_per_update_seconds"] |
| memory = comparison["peak_process_rss_mb"] |
| lines += ["", f"Final loss: {loss['mean']:.8e} ± {loss['std']:.8e}; paper: 0.10 ± 0.02.", |
| f"Time per update: {timing['mean']:.6f} ± {timing['std']:.6f} s; {timing['definition']}.", |
| f"Peak process RSS: {memory['mean']:.3f} ± {memory['std']:.3f} MB. Actual hardware is recorded in results/compute_hardware.txt; timing/memory are not hardware-identical to the paper.", |
| "Population standard deviation (ddof=0) is used; the paper does not specify this convention."] |
| lines += ["", "The selected best checkpoint is a single model; the paper entries are ten-run statistics.", |
| "See table3_comparison.json for separately aggregated final and best checkpoints.", "", "## Explicit limitations"] |
| lines += [f"- {item}" for item in config["assumptions"]] |
| lines += ["- Appendix A.2 printed sums can cancel between components. The literal printed losses were retained.", |
| "- The reference contains saved time intervals of 0.1 s, while the paper describes a solver step of 0.01 s.", |
| "- Validation uses a fixed reference subset for checkpoint selection. Holdout metrics exclude those records.", |
| "", "Source: [paper](https://arxiv.org/pdf/2503.16678), sections 4/5, Figure 2, Tables 2/3 and Appendix A.2."] |
| report = evaluation/"report.md" |
| report.write_text("\n".join(lines)+"\n") |
| atomic_json(evaluation/"visualization_manifest.json", {"figures": paths, "report": str(report), "dpi": config["visualization"]["dpi"], "reference": str(resolved(config, "reference"))}) |
| print(f"VISUALIZATION_COMPLETED figures={len(paths)} report={report}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--config", default=str(ROOT/"config/config.yaml")) |
| parser.add_argument("--aggregate", action="store_true") |
| parser.add_argument("--seed", type=int) |
| parser.add_argument("--output-dir", type=Path) |
| parser.add_argument("--predictions", type=Path) |
| args = parser.parse_args() |
| cfg = load_config(args.config) |
| if args.aggregate: aggregate(cfg) |
| else: make_results(cfg, args.seed, args.output_dir, args.predictions) |
|
|
