| """Create vector plots and a LaTeX table from recorded benchmark measurements.""" |
|
|
| import argparse |
| from pathlib import Path |
| import json |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| from matplotlib import pyplot as plt, font_manager |
| import pandas as pd |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| COLORS = ["#4E95C0", "#245778", "#A7D3EE", "#718594", "#7EB4D5", "#364C5B", "#D7EAF6"] |
| LABELS = { |
| "dooable": "DooABLe", |
| "tb_uniform": "TB, uniform backward", |
| "tb_exact": "TB, exact backward", |
| "uniform": "Uniform executable", |
| "reference_tilt": "Reference tilt", |
| "exact": "Exact joint law", |
| "zero_cost": "Zero route cost", |
| "duplicate_endpoints": "Duplicate endpoints", |
| "unnormalized": "Unnormalized backward", |
| } |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--run", required=True) |
| parser.add_argument("--output") |
| args = parser.parse_args() |
| run = Path(args.run) |
| out = Path(args.output) if args.output else run / "plots" |
| out.mkdir(parents=True, exist_ok=True) |
| font = ROOT / "assets/fonts/Ubuntu-Regular.ttf" |
| if not font.exists(): |
| font = ROOT / "paper/figures/fonts/Ubuntu-Regular.ttf" |
| font_manager.fontManager.addfont(str(font)) |
| plt.rcParams.update( |
| { |
| "font.family": "Ubuntu", |
| "font.size": 9, |
| "mathtext.fontset": "cm", |
| "pdf.fonttype": 42, |
| "axes.spines.top": False, |
| "axes.spines.right": False, |
| } |
| ) |
| df = pd.read_csv(run / "metrics.csv") |
| methods = list(dict.fromkeys(df.method)) |
| fig, axes = plt.subplots(1, 3, figsize=(10, 3.2)) |
| specifications = [ |
| ("endpoint_tv", "Endpoint TV"), |
| ("conditional_free_energy_gap", "Conditional gap (cost units)"), |
| ("mean_cost", "Mean execution cost"), |
| ] |
| for ax, (metric, label) in zip(axes, specifications): |
| summary = df.groupby("method")[metric].agg(["mean", "sem"]).reindex(methods) |
| ax.bar( |
| range(len(methods)), |
| summary["mean"], |
| yerr=summary["sem"].fillna(0), |
| color=COLORS[: len(methods)], |
| capsize=2, |
| ) |
| ax.set_ylabel(label) |
| ax.set_xticks( |
| range(len(methods)), [LABELS[x] for x in methods], rotation=40, ha="right" |
| ) |
| fig.tight_layout() |
| fig.savefig(out / "benchmark.pdf") |
| fig.savefig(out / "benchmark.png", dpi=220) |
| plt.close(fig) |
| histories = [] |
| for path in sorted(run.glob("*_seed*/training.json")): |
| frame = pd.read_json(path) |
| frame["method"] = path.parent.name.rsplit("_seed", 1)[0] |
| frame["seed"] = int(path.parent.name.rsplit("_seed", 1)[1]) |
| histories.append(frame) |
| if histories: |
| history = pd.concat(histories) |
| fig, ax = plt.subplots(figsize=(5.3, 3.1)) |
| for color, (method, group) in zip(COLORS, history.groupby("method")): |
| stats = group.groupby("step").endpoint_tv.agg(["mean", "sem"]) |
| x = stats.index.to_numpy() |
| mean = stats["mean"].to_numpy() |
| sem = stats["sem"].fillna(0).to_numpy() |
| ax.plot(x, mean, label=LABELS[method], color=color) |
| ax.fill_between(x, mean - sem, mean + sem, color=color, alpha=0.15) |
| ax.set_xlabel("Training updates") |
| ax.set_ylabel("Endpoint TV") |
| ax.legend(frameon=False) |
| fig.tight_layout() |
| fig.savefig(out / "training.pdf") |
| fig.savefig(out / "training.png", dpi=220) |
| plt.close(fig) |
| summary = df.groupby("method")[[x[0] for x in specifications]].agg(["mean", "sem"]) |
| lines = [ |
| "\\begin{tabular}{lccc}", |
| "\\toprule", |
| "Method & TV & Conditional gap & Mean cost \\\\", |
| "\\midrule", |
| ] |
| for method in methods: |
| cells = [] |
| for metric, _ in specifications: |
| value = summary.loc[method, (metric, "mean")] |
| sem = summary.loc[method, (metric, "sem")] |
| cells.append( |
| f"${value:.4g}$" if pd.isna(sem) else f"${value:.4g} \\pm {sem:.2g}$" |
| ) |
| lines.append(LABELS[method] + " & " + " & ".join(cells) + " \\\\") |
| lines.extend(["\\bottomrule", "\\end{tabular}"]) |
| (out / "measurements.tex").write_text("\n".join(lines) + "\n") |
| print(json.dumps({"rows": len(df), "plots": str(out)})) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
