File size: 4,371 Bytes
81ae663 | 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 | """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()
|