bundle/scripts/make_figs.py

#!/usr/bin/env python
"""Generate plotly HTML figures + raw CSVs for the logbook figure cells."""
import json, os, csv
import plotly.graph_objects as go

OUT = "outputs/figs"; os.makedirs(OUT, exist_ok=True)
COL = {"mfa": "#5B4FCF", "base": "#F2A93B", "sae": "#F2A93B", "dm": "#F2A93B"}

def save(fig, name, rows, header):
    fig.write_html(f"{OUT}/{name}.html", include_plotlyjs="cdn", full_html=True)
    with open(f"{OUT}/{name}.csv", "w", newline="") as f:
        w = csv.writer(f); w.writerow(header); w.writerows(rows)
    print("wrote", name)

def claim1():
    d = json.load(open("outputs/claim1_gemma_l18.json"))
    # panel A: relative reconstruction error
    fig = go.Figure()
    xs = ["centroid only<br>(region)", "centroid + rank-10<br>offset (Eq 11-13)"]
    ys = [d["err_centroid_only"], d["err_centroid_plus_offset"]]
    fig.add_bar(x=xs, y=ys, marker_color=["#B9B3E8", "#5B4FCF"],
                text=[f"{v:.3f}" for v in ys], textposition="outside")
    fig.update_layout(title="Claim 1 — MFA reconstruction of Gemma-2-2B L18 activations "
                            "(relative error, lower=better)", yaxis_title="mean ||x-x̂||/||x||",
                      template="plotly_white", height=430, showlegend=False)
    save(fig, "claim1_recon", [[a, b] for a, b in zip(xs, ys)], ["segment", "rel_error"])
    # panel B: variance explained by rank-10 subspace vs random
    fig2 = go.Figure()
    xs2 = ["MFA local rank-10<br>offset", "random rank-10<br>subspace"]
    ys2 = [d["r2_two_segment"], d["r2_random_rank10_subspace"]]
    fig2.add_bar(x=xs2, y=ys2, marker_color=["#5B4FCF", "#F2A93B"],
                 text=[f"{v:.3f}" for v in ys2], textposition="outside")
    fig2.update_layout(title="Claim 1 — centered variance explained: local subspaces capture "
                             f"{d['r2_two_segment']/max(d['r2_random_rank10_subspace'],1e-9):.0f}× a random 10-dim subspace",
                       yaxis_title="R² (centered)", template="plotly_white", height=430, showlegend=False)
    save(fig2, "claim1_variance", [[a, b] for a, b in zip(xs2, ys2)], ["subspace", "r2"])

def claim2():
    d = json.load(open("outputs/claim2_gemma_l18.json"))
    fig = go.Figure()
    xs = ["nearest neighbor<br>(centroid kNN)", "random component"]
    ys = [d["knn_topTokenOverlap_nearest_mean"], d["knn_topTokenOverlap_random_mean"]]
    fig.add_bar(x=xs, y=ys, marker_color=["#5B4FCF", "#F2A93B"],
                text=[f"{v:.3f}" for v in ys], textposition="outside")
    r = ys[0]/max(ys[1], 1e-9)
    fig.update_layout(title=f"Claim 2 — neighboring Gaussians share {r:.1f}× more top tokens "
                            "than random (semantic neighborhoods, Fig 1)",
                      yaxis_title="mean top-token overlap", template="plotly_white",
                      height=430, showlegend=False)
    save(fig, "claim2_knn", [[a, b] for a, b in zip(xs, ys)], ["pair", "overlap"])

def claim3():
    d = json.load(open("outputs/claim3_gemma_l18.json"))
    fig = go.Figure()
    xs = ["MFA<br>(centroid+offset)", "Gemmascope SAE"]
    ys = [d["mfa_active_components_meanR>1e-4"], d["sae_L0_mean"]]
    fig.add_bar(x=xs, y=ys, marker_color=["#5B4FCF", "#F2A93B"],
                text=[f"{v:.0f} feats" for v in ys], textposition="outside")
    fig.update_layout(title="Claim 3 — #active features per activation (MFA concentrates mass in ~2 objects)",
                      yaxis_title="active features", template="plotly_white", height=430,
                      yaxis_type="log", showlegend=False)
    save(fig, "claim3_features", [[a, b] for a, b in zip(xs, ys)], ["method", "active_features"])

def claim5():
    p = "outputs/claim5_scores.json"
    if not os.path.exists(p):
        print("claim5 scores not ready"); return
    d = json.load(open(p)); s = d["summary"]
    fig = go.Figure()
    xs = ["MFA centroid<br>(Eq 14)", "DiffMeans<br>(supervised)"]
    ys = [s["mfa_median"], s["diffmeans_median"]]
    fig.add_bar(x=xs, y=ys, marker_color=["#5B4FCF", "#F2A93B"],
                text=[f"{v:.3f}" for v in ys], textposition="outside")
    fig.update_layout(title=f"Claim 5 — median steering score on Gemma-2-2B L18 "
                            f"(MFA {s['mfa_over_diffmeans_ratio_median']}× DiffMeans)",
                      yaxis_title="steering score (harmonic mean)", template="plotly_white",
                      height=430, showlegend=False)
    save(fig, "claim5_steer", [[a, b] for a, b in zip(xs, ys)], ["method", "median_score"])

if __name__ == "__main__":
    import sys
    for name in (sys.argv[1:] or ["claim1", "claim2", "claim3", "claim5"]):
        try:
            globals()[name]()
        except FileNotFoundError as e:
            print("skip", name, e)