spb/wpai0001_uqo

1"""Figure recipes for the Tier-0 transplant run (graft_tier0) and the Tier-1 task-vector run.2Every plotted number comes from results.json of a completed run; no computation happens here."""3from __future__ import annotations45from pathlib import Path67import matplotlib.pyplot as plt8import numpy as np910from wp_ai_0001.visualization.style import CATEGORICAL, DOUBLE_WIDTH_IN, ENTITY_COLORS, SEQUENTIAL_CMAP, TEXT_SECONDARY, apply_style, heatmap1112MAIN_CONDS = [("subspace_k8", "subspace graft k=8"), ("subspace_random_k8", "random subspace k=8"), ("act_identity_all", "identity (same width)"), ("act_procrustes_all", "Procrustes"), ("act_unrelated_procrustes_all", "Procrustes, task-free fit"),13              ("act_ridge_all", "ridge"), ("act_cca_all", "CCA"), ("ctl_unrelated_data_map", "ridge, task-free fit"), ("ctl_sham_map", "sham map (random orth.)"),14              ("ctl_norm_matched_random", "norm-matched random"), ("ctl_shuffled_token", "shuffled-token donor"), ("ctl_shuffled_layer", "shuffled-layer donor"),15              ("ctl_unrelated_donor", "unrelated-skill donor"), ("steer_a1.0", "steering vector α=1"), ("adapter_graft", "adapter graft (donor-free)"),16              ("weight_L0_mlp", "weight graft L0 MLP"), ("weight_L0_block_ln", "weight graft L0 block"), ("baseline_supervised_adapter", "supervised adapter (labels)"),17              ("baseline_distillation", "distillation (labels-free, donor outputs)")]181920LABELLED = {"act_identity_all", "act_procrustes_all", "act_unrelated_procrustes_all", "ctl_unrelated_data_map", "ctl_sham_map", "weight_L0_block_ln", "baseline_supervised_adapter", "act_ridge_all", "subspace_k8", "adapter_graft"}21LABEL_OFFSETS = {"act_unrelated_procrustes_all": (-4, -2), "baseline_supervised_adapter": (-4, 6), "act_identity_all": (-4, 2), "act_procrustes_all": (-4, -6), "ctl_unrelated_data_map": (-4, 4), "ctl_sham_map": (4, 4), "weight_L0_block_ln": (4, 2), "act_ridge_all": (-4, 2), "subspace_k8": (-4, -8), "adapter_graft": (-4, 2)}222324def _pairs(results, same_init: bool | None = None):25    pm = results.get("pair_meta")26    if not pm:  # older runs: infer from condition keys (same init seed <=> same seed suffix)27        pm = {p: {"same_init_seed": p.split("->")[0].rsplit("_s", 1)[1] == p.split("->")[1].rsplit("_s", 1)[1]} for p in results["conditions"].get("act_identity_all", {})}28    return [p for p, m in pm.items() if same_init is None or m["same_init_seed"] == same_init]293031def fig_pareto(results: dict, out: Path) -> None:32    """Target gain vs control retention for every condition (mean over pairs; same-init and different-init separately)."""33    apply_style()34    fig, axes = plt.subplots(1, 2, figsize=(DOUBLE_WIDTH_IN, 2.8), sharey=True)35    for ax, same in zip(axes, (True, False)):36        pairs = _pairs(results, same)37        if not pairs:38            ax.set_title("no pairs of this kind in this run", loc="left")39            continue40        base = np.mean([results["conditions"]["base"][p.split("->")[1]]["target_iid"]["accuracy"] for p in pairs])41        for i, (cond, label) in enumerate(MAIN_CONDS):42            d = results["conditions"].get(cond, {})43            vals = [d[p] for p in pairs if p in d]44            if not vals:45                continue46            t = np.mean([v["target_iid"]["accuracy"] for v in vals])47            c = np.mean([v["control_iid"]["accuracy"] for v in vals])48            col = ENTITY_COLORS["graft"] if cond.startswith(("act_", "subspace", "adapter", "weight", "steer")) else (ENTITY_COLORS["sham"] if cond.startswith("ctl_") else ENTITY_COLORS["recipient"])49            ax.scatter(c, t, s=18, color=col, edgecolor="#ffffff", linewidth=0.5, zorder=3)50            if cond in LABELLED:51                dx, dy = LABEL_OFFSETS.get(cond, (-3, 3))52                ax.annotate(label, (c, t), fontsize=5.5, xytext=(dx, dy), textcoords="offset points", color=TEXT_SECONDARY, ha="right" if dx < 0 else "left")53        ax.axhline(base, color=ENTITY_COLORS["chance"], linewidth=0.8, linestyle=":", label="recipient baseline (target)")54        ax.set_xlabel("control-stratum accuracy (retention)")55        ax.set_title("same initialization seed" if same else "different initialization seed", loc="left")56        ax.set_xlim(-0.02, 1.02)57        ax.set_ylim(-0.02, 1.02)58    axes[0].set_ylabel("target-stratum accuracy")59    from matplotlib.lines import Line2D6061    h = [Line2D([], [], marker="o", linestyle="", color=ENTITY_COLORS["graft"], label="graft family"), Line2D([], [], marker="o", linestyle="", color=ENTITY_COLORS["sham"], label="control"),62         Line2D([], [], marker="o", linestyle="", color=ENTITY_COLORS["recipient"], label="conventional learning baseline"), Line2D([], [], linestyle=":", color=ENTITY_COLORS["chance"], label="recipient baseline")]63    axes[1].legend(handles=h, loc="lower left", fontsize=6)64    fig.savefig(out / "fig_pareto_tier0.pdf")65    plt.close(fig)666768def fig_layer_grid_dose(results: dict, out: Path) -> None:69    apply_style()70    fig, axes = plt.subplots(1, 2, figsize=(DOUBLE_WIDTH_IN * 0.85, 2.3))71    cells = results["layer_grid"]["cells"]72    L = int(max(int(k.split("->")[0]) for k in cells)) + 173    m = np.zeros((L, L))74    for k, v in cells.items():75        a, b = map(int, k.split("->"))76        m[a, b] = v["target"]77    heatmap(axes[0], m, [f"donor L{i}" for i in range(L)], [f"L{i}" for i in range(L)], cmap=SEQUENTIAL_CMAP, vmin=0, vmax=1, cbar_label="target accuracy", title=f"ridge graft, {results['layer_grid']['pair']}", annotate=True)78    axes[0].set_xlabel("recipient layer")79    d = results["dose"]["cells"]80    xs = sorted(float(k) for k in d)81    axes[1].plot(xs, [d[str(x) if str(x) in d else k]["target"] for x, k in zip(xs, sorted(d, key=float))], marker="o", markersize=3, color=ENTITY_COLORS["graft"], label="target")82    axes[1].plot(xs, [d[k]["control"] for k in sorted(d, key=float)], marker="s", markersize=3, color=ENTITY_COLORS["recipient"], label="control")83    axes[1].set_xlabel("graft strength α")84    axes[1].set_ylabel("accuracy")85    axes[1].set_ylim(-0.02, 1.02)86    axes[1].legend(loc="center right")87    axes[1].set_title("dose–response (ridge, L0→L0)", loc="left")88    fig.savefig(out / "fig_layer_dose_tier0.pdf")89    plt.close(fig)909192def fig_subspace(results: dict, out: Path) -> None:93    apply_style()94    pairs = _pairs(results, True)95    ks = sorted({int(c.split("subspace_k")[1]) for c in results["conditions"] if c.startswith("subspace_k")})96    fig, ax = plt.subplots(figsize=(DOUBLE_WIDTH_IN * 0.5, 2.3))97    for prefix, label, col in (("subspace_k", "Ω-contrast PCA subspace", ENTITY_COLORS["graft"]), ("subspace_random_k", "random subspace (same k)", ENTITY_COLORS["sham"])):98        ys = [np.mean([results["conditions"][f"{prefix}{k}"][p]["target_iid"]["accuracy"] for p in pairs if p in results["conditions"][f"{prefix}{k}"]]) for k in ks]99        ax.plot(ks, ys, marker="o", markersize=3, color=col, label=label)100    full = np.mean([results["conditions"]["act_unrelated_procrustes_all"][p]["target_iid"]["accuracy"] for p in pairs])101    ax.axhline(full, color=ENTITY_COLORS["graft"], linestyle=":", linewidth=0.8, label="full-dimensional graft (same map)")102    ax.set_xscale("log", base=2)103    ax.set_xlabel("graft dimensionality k")104    ax.set_ylabel("target accuracy")105    ax.set_ylim(-0.02, 1.02)106    ax.legend(fontsize=6)107    fig.savefig(out / "fig_subspace_tier0.pdf")108    plt.close(fig)109110111def fig_portability(results: dict, out: Path) -> None:112    """Donor × recipient matrix of target accuracy for the identity and Procrustes grafts."""113    apply_style()114    pairs = _pairs(results)115    ds = sorted({p.split("->")[0] for p in pairs})116    rs = sorted({p.split("->")[1] for p in pairs})117    fig, axes = plt.subplots(1, 3, figsize=(DOUBLE_WIDTH_IN, 2.1), sharey=True)118    for ax, cond in zip(axes, ("act_identity_all", "act_procrustes_all", "act_unrelated_procrustes_all")):119        m = np.full((len(ds), len(rs)), np.nan)120        for i, d in enumerate(ds):121            for j, r in enumerate(rs):122                v = results["conditions"][cond].get(f"{d}->{r}")123                if v:124                    m[i, j] = v["target_iid"]["accuracy"]125        title = {"act_identity_all": "identity graft", "act_procrustes_all": "Procrustes (task rows)", "act_unrelated_procrustes_all": "Procrustes (task-free rows)"}[cond]126        heatmap(ax, m, [x.replace("donor_s", "D") for x in ds], [x.replace("recipient_s", "R") for x in rs], cmap=SEQUENTIAL_CMAP, vmin=0, vmax=1, cbar_label="target accuracy", title=title, annotate=True)127    fig.savefig(out / "fig_portability_tier0.pdf")128    plt.close(fig)129130131def fig_tier1(results: dict, out: Path) -> None:132    """Per task: recipient zero-shot, few-shot, own vector, best transferred vector, controls, for each pair."""133    apply_style()134    pairs = list(results["pairs"])135    tasks = list(results["tasks"])136    fig, axes = plt.subplots(1, len(pairs), figsize=(DOUBLE_WIDTH_IN, 2.4), sharey=True)137    axes = np.atleast_1d(axes)138    series = [("recipient_zero", "zero-shot", ENTITY_COLORS["recipient"]), ("recipient_few", "few-shot (prompting)", ENTITY_COLORS["chance"]), ("recipient_own_vector", "own task vector", ENTITY_COLORS["donor"]),139              ("best_transfer", "transferred donor vector", ENTITY_COLORS["graft"]), ("best_control", "best control", ENTITY_COLORS["sham"])]140    w = 0.15141    for ax, p in zip(axes, pairs):142        pr = results["pairs"][p]143        x = np.arange(len(tasks))144        for k, (key, label, col) in enumerate(series):145            ys = []146            for t in tasks:147                tr = pr["tasks"][t]148                if key == "best_transfer":149                    ys.append(max(v["accuracy"] for c, v in tr["conditions"].items() if not c.startswith("ctl") and "random_orth" not in c and "sham" not in c and "shuffled" not in c))150                elif key == "best_control":  # null controls only (amendment A1: shuffled-query state is a diagnostic)151                    ys.append(max(v["accuracy"] for c, v in tr["conditions"].items() if (c.startswith("ctl") or "random_orth" in c or "sham" in c) and "shuffled" not in c))152                else:153                    ys.append(tr[key])154            ax.bar(x + (k - 2) * w, ys, width=w - 0.02, color=col, label=label, linewidth=0)155        ax.set_xticks(x)156        ax.set_xticklabels([{"antonym": "antonym", "capital": "capital", "morph": "morph.", "addthree": "add 3"}.get(t, t) for t in tasks], rotation=30, ha="right")157        ax.set_title(p.replace("->", " → "), loc="left")158        ax.set_ylim(0, 1.02)159    axes[0].set_ylabel("exact-match accuracy")160    axes[-1].legend(fontsize=5.5, loc="upper right")161    fig.savefig(out / f"fig_tier1_transfer_{results['run_id'].split('_')[-1][:6]}.pdf")162    plt.close(fig)163164165def fig_mechanism(results: dict, out: Path) -> None:166    """Ablation → loss → rescue chain and the rule-vs-cache test."""167    apply_style()168    fig, axes = plt.subplots(1, 2, figsize=(DOUBLE_WIDTH_IN, 2.4), gridspec_kw={"wspace": 0.25})169    c = results["conditions"]170    d0 = [k for k in c["base"] if k.startswith("donor_s0")][0]171    chain = [("intact", c["base"][d0]["target_iid"]["accuracy"], c["base"][d0]["control_iid"]["accuracy"]),172             ("L0 attn.\nablated", *[list(c["mech_donor_ablated_attn0"].values())[0][s]["accuracy"] for s in ("target_iid", "control_iid")]),173             ("rescue:\nintact donor", *[list(c["mech_rescue_from_intact_donor"].values())[0][s]["accuracy"] for s in ("target_iid", "control_iid")]),174             ("rescue:\nsham map", *[list(c["mech_rescue_sham"].values())[0][s]["accuracy"] for s in ("target_iid", "control_iid")])]175    x = np.arange(len(chain))176    axes[0].bar(x - 0.18, [v[1] for v in chain], width=0.34, color=ENTITY_COLORS["donor"], label="target stratum", linewidth=0)177    axes[0].bar(x + 0.18, [v[2] for v in chain], width=0.34, color=ENTITY_COLORS["chance"], label="control stratum", linewidth=0)178    axes[0].set_xticks(x)179    axes[0].set_xticklabels([v[0] for v in chain], fontsize=6)180    axes[0].set_ylim(0, 1.02)181    axes[0].set_ylabel("accuracy")182    axes[0].set_title("ablation → loss → rescue (donor seed 0)", loc="left", fontsize=7.5)183    axes[0].legend(fontsize=6, loc="upper center", ncol=2, bbox_to_anchor=(0.5, -0.32))184    h5 = results["rule_vs_cache"]185    conds = [("h5_act_procrustes_seenfit", "act. graft\nProcrustes"), ("h5_act_ridge_seenfit", "act. graft\nridge"), ("h5_adapter_seenfit", "donor-free\nadapter")]186    conds = [(k, l) for k, l in conds if k in h5]187    x = np.arange(len(conds))188    for i, (key, lab, col) in enumerate((("seen", "seen Ω inputs", ENTITY_COLORS["graft"]), ("unseen", "unseen Ω inputs", ENTITY_COLORS["sham"]))):189        vals = [h5[k][key]["estimate"] for k, _ in conds]190        err = [[h5[k][key]["estimate"] - h5[k][key]["low"] for k, _ in conds], [h5[k][key]["high"] - h5[k][key]["estimate"] for k, _ in conds]]191        axes[1].bar(x + (i - 0.5) * 0.36, vals, width=0.34, color=col, label=lab, linewidth=0, yerr=err, error_kw={"elinewidth": 0.6, "ecolor": TEXT_SECONDARY, "capsize": 1.5})192    axes[1].set_xticks(x)193    axes[1].set_xticklabels([l for _, l in conds], fontsize=6)194    axes[1].set_ylim(0, 1.02)195    axes[1].set_title("rule vs cache (Ω-step accuracy; maps fitted on 6 of 11 inputs)", loc="left", fontsize=7.5)196    axes[1].legend(fontsize=6, loc="upper center", ncol=2, bbox_to_anchor=(0.5, -0.32))197    fig.savefig(out / "fig_mechanism_tier0.pdf")198    plt.close(fig)199