scripts/plot_20layer_3way.py

"""20-layer 3-way: single-state normalized (P=1) vs hard patch (P=8) vs slotroute,
best ckpt. Grouped bars over the 4 higher=better axes + a fidelity train/ood
split panel (the interesting reversal). Three separately-trained 20-layer models."""
import csv, numpy as np, matplotlib
matplotlib.use("Agg"); import matplotlib.pyplot as plt

rows = list(csv.DictReader(open("workspace/markdown/arl2_20layer_three_way_quality.csv")))
best = [r for r in rows if r["checkpoint"] == "best"]
variants = ["frame", "patch", "patch_overlap"]
LABEL = {"frame": "single-state\n(P=1)", "patch": "hard patch\n(P=8)", "patch_overlap": "slotroute"}
COL = {"frame": "#2F6FBF", "patch": "#E8710A", "patch_overlap": "#7B4FA3"}
METR = ["clip_fidelity_teacher", "clip_text_align", "clip_temporal", "flow_smoothness"]
TITLE = {"clip_fidelity_teacher": "Fidelity vs teacher ↑", "clip_text_align": "Prompt fidelity ↑",
         "clip_temporal": "Temporal coherence ↑", "flow_smoothness": "Flow smoothness ↑"}

def wavg(v, m, split=None):
    vr = [r for r in best if r["variant"] == v and (split is None or r["split"] == split)]
    xs = [float(r[m]) for r in vr if r[m] not in ("", "nan")]
    return float(np.mean(xs)) if xs else float("nan")  # per-video already; mean over videos

# per-video mean over all 8 (n weighted naturally since we average raw videos)
def vids(v, m, split=None):
    out = []
    for r in best:
        if r["variant"] == v and (split is None or r["split"] == split) and r[m] not in ("", "nan"):
            out.append(float(r[m]))
    return out
def mean(v, m, split=None):
    x = vids(v, m, split); return float(np.mean(x)) if x else float("nan")

fig = plt.figure(figsize=(12, 8), dpi=130)
fig.suptitle("20-layer GDN (v2): single-state vs hard-patch(P=8) vs slotroute  —  best ckpt, 8 videos/variant\n"
             "480p/21f/seed0 · reference-free CLIP + optical-flow + fidelity-vs-teacher", fontsize=10.5, y=0.995)
gs = fig.add_gridspec(2, 3)
axmap = [fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[0, 1]), fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[1, 1])]
for ax, m in zip(axmap, METR):
    ys = [mean(v, m) for v in variants]; x = np.arange(3)
    ax.bar(x, ys, width=0.62, color=[COL[v] for v in variants], edgecolor="white", lw=1.4, zorder=3)
    lo, hi = min(ys), max(ys); pad = (hi - lo) * 0.6 + 1e-4; ax.set_ylim(lo - pad, hi + pad * 1.5)
    for xi, y in zip(x, ys):
        ax.text(xi, y, f"{y:.4f}", ha="center", va="bottom" if y >= 0 else "top", fontsize=8.5)
    wi = int(np.argmax(ys)); ax.text(x[wi], ax.get_ylim()[1], "best", ha="center", va="top",
                                     fontsize=8.5, color=COL[variants[wi]], fontweight="bold")
    ax.set_title(TITLE[m], fontsize=9.5, loc="left")
    ax.set_xticks(x); ax.set_xticklabels([LABEL[v] for v in variants], fontsize=8)
    ax.tick_params(colors="#6b6b6b", labelsize=8); ax.grid(True, axis="y", color="#e6e6e6", lw=0.8, zorder=0)
    for s in ax.spines.values(): s.set_color("#e6e6e6")

# fidelity train vs ood split (the reversal)
axs = fig.add_subplot(gs[:, 2])
w = 0.38; x = np.arange(3)
tr = [mean(v, "clip_fidelity_teacher", "train") for v in variants]
oo = [mean(v, "clip_fidelity_teacher", "ood") for v in variants]
axs.bar(x - w/2, tr, w, color="#9DBBE0", edgecolor="white", lw=1.2, label="train (5)", zorder=3)
axs.bar(x + w/2, oo, w, color="#2F6FBF", edgecolor="white", lw=1.2, label="ood (3)", zorder=3)
for xi, y in zip(x - w/2, tr): axs.text(xi, y, f"{y:.3f}", ha="center", va="bottom", fontsize=7.5)
for xi, y in zip(x + w/2, oo): axs.text(xi, y, f"{y:.3f}", ha="center", va="bottom", fontsize=7.5)
axs.set_title("Fidelity vs teacher — train vs OOD split\n(patch WINS ood, loses train)", fontsize=9.5, loc="left")
axs.set_xticks(x); axs.set_xticklabels([LABEL[v] for v in variants], fontsize=8)
axs.set_ylim(min(tr+oo) - 0.01, max(tr+oo) + 0.012); axs.legend(fontsize=8, frameon=False)
axs.tick_params(colors="#6b6b6b", labelsize=8); axs.grid(True, axis="y", color="#e6e6e6", lw=0.8, zorder=0)
for s in axs.spines.values(): s.set_color("#e6e6e6")

fig.tight_layout(rect=[0, 0, 1, 0.95])
out = "workspace/markdown/arl2_20layer_3way_bars.png"; fig.savefig(out, bbox_inches="tight"); print("wrote", out)
print("\n=== 20-layer best-ckpt, mean over 8 videos ===")
for m in METR + ["flow_magnitude"]:
    print(f"{m:24s}", {LABEL[v].replace(chr(10),' '): round(mean(v, m), 4) for v in variants})
print("\n=== fidelity train/ood split ===")
for v in variants:
    print(f"{LABEL[v].replace(chr(10),' '):18s} train={mean(v,'clip_fidelity_teacher','train'):.4f}  ood={mean(v,'clip_fidelity_teacher','ood'):.4f}")