spb/wpai0001_uqo

1#!/usr/bin/env python32"""Generate LaTeX table fragments in paper/tables/ from registered run results (no computation).34Usage: python scripts/build_tables.py --graft <run_id> [--pilot <run_id>] [--tier1 <run_id>] [--confirmatory <run_id>] [--tier1-confirmatory <run_id>]5Each fragment starts with a comment naming its source run so that every number is traceable.6"""7from __future__ import annotations89import argparse10import json11import sys12from pathlib import Path1314import numpy as np1516ROOT = Path(__file__).resolve().parents[1]17sys.path.insert(0, str(ROOT / "src"))18OUT = ROOT / "paper" / "tables"1920PRETTY = {21    "act_identity_all": "Activation graft, identity map", "act_procrustes_all": "Activation graft, Procrustes (task rows)",22    "act_unrelated_procrustes_all": "Activation graft, Procrustes (task-free rows)", "ctl_unrelated_data_map": "Activation graft, ridge (task-free rows)",23    "act_ridge_all": "Activation graft, ridge (task rows)", "act_cca_all": "Activation graft, CCA (task rows)",24    "ctl_sham_map": "Control: sham map (random orthogonal)", "ctl_norm_matched_random": "Control: norm-matched random rows",25    "ctl_shuffled_token": "Control: shuffled-item donor rows", "ctl_shuffled_layer": "Control: donor layer 3 rows (ridge)",26    "ctl_unrelated_donor": "Control: unrelated-skill donor (ridge)", "steer_a1.0": "Steering vector, $\\alpha=1$", "steer_a4.0": "Steering vector, $\\alpha=4$",27    "subspace_k8": "Subspace graft, $k=8$", "subspace_k32": "Subspace graft, $k=32$", "subspace_k128": "Subspace graft, $k=128$",28    "subspace_random_k32": "Control: random subspace, $k=32$", "adapter_graft": "Adapter graft (donor-free, ridge interface)",29    "baseline_supervised_adapter": "Baseline: supervised adapter (labels)", "baseline_distillation": "Baseline: distillation (donor outputs)",30    "weight_L0_mlp": "Weight graft: L0 MLP", "weight_L0_attn": "Weight graft: L0 attention", "weight_L0_block_ln": "Weight graft: L0 block",31    "weight_L012_blocks_embed": "Weight graft: L0--L2 blocks + $\\Omega$ embedding", "weight_embed_omega": "Weight graft: $\\Omega$ embedding only",32    "ctl_weight_shuffled_layer": "Control: donor L3 block into L0", "ctl_weight_unrelated_donor": "Control: unrelated-skill donor L0 block",33}343536def fmt(x, nd=3):37    return "---" if x is None else f"{x:.{nd}f}"383940def load(run):41    return json.loads((ROOT / "experiments" / "summaries" / run / "results.json").read_text())424344def pairs_of(r, same=None):45    pm = r["pair_meta"]46    return [p for p, m in pm.items() if same is None or m["same_init_seed"] == same]474849def mean_sd(r, cond, pairs, split="target_iid", key="accuracy"):50    v = [r["conditions"][cond][p][split][key] for p in pairs if p in r["conditions"].get(cond, {}) and split in r["conditions"][cond][p]]51    if not v:52        return None, None, 053    return float(np.mean(v)), float(np.std(v)), len(v)545556def table_main(r, run, name="tab_main_effects.tex"):57    same, diff = pairs_of(r, True), pairs_of(r, False)58    rows = []59    for cond in PRETTY:60        if cond not in r["conditions"]:61            continue62        ts, tsd, n1 = mean_sd(r, cond, same)63        cs, _, _ = mean_sd(r, cond, same, "control_iid")64        td, tdd, n2 = mean_sd(r, cond, diff)65        cd, _, _ = mean_sd(r, cond, diff, "control_iid")66        # control-adjusted effect vs sham on the matched pairs (mean of per-pair estimates and CI bounds)67        adj = [r["contrasts"][p][cond]["control_adjusted"] for p in same + diff if cond in r["contrasts"].get(p, {}) and "control_adjusted" in r["contrasts"][p][cond]]68        adj_s = f"{np.mean([a['estimate'] for a in adj]):+.3f} [{np.mean([a['low'] for a in adj]):+.3f}, {np.mean([a['high'] for a in adj]):+.3f}]" if adj else "---"69        rows.append(f"{PRETTY[cond]} & {fmt(ts)} ({fmt(tsd,2)}) & {fmt(cs)} & {fmt(td) if n2 else '---'} {('('+fmt(tdd,2)+')') if n2 else ''} & {fmt(cd) if n2 else '---'} & {adj_s} \\\\")70    base_t = [r["conditions"]["base"][k]["target_iid"]["accuracy"] for k in r["conditions"]["base"] if k.startswith("recipient_s") and "_s2_" not in k and "w96" not in k and "L6" not in k]71    hdr = f"""% source: experiments/summaries/{run}/results.json (generated by scripts/build_tables.py)72\\begin{{tabular}}{{lccccc}}73\\toprule74 & \\multicolumn{{2}}{{c}}{{same initialization ($n={len(same)}$ pairs)}} & \\multicolumn{{2}}{{c}}{{different initialization ($n={len(diff)}$)}} & control-adjusted \\\\75\\cmidrule(lr){{2-3}} \\cmidrule(lr){{4-5}}76Condition & target (sd) & control & target (sd) & control & $\\Delta_T^{{\\mathrm{{adj}}}}$ [95\\% CI] \\\\77\\midrule78Recipient, unmodified & {fmt(np.mean(base_t))} ({fmt(np.std(base_t),2)}) & 1.000 & & & \\\\79"""80    (OUT / name).write_text(hdr + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n")818283def table_maps(r, run):84    same, diff = pairs_of(r, True), pairs_of(r, False)85    rows = []86    for fam, label in (("identity", "identity"), ("procrustes", "Procrustes (task rows)"), ("unrelated_procrustes", "Procrustes (task-free rows)"), ("ridge", "ridge (task rows)"), ("unrelated_data", "ridge (task-free rows)"), ("cca", "CCA (task rows)"), ("random_orth", "random orthogonal (sham)")):87        def get(p):88            m = r["maps"][f"{p}|{fam}"]89            return m.get("r2_fit", m.get("r2_on_target_rows"))90        a = [get(p) for p in same]91        b = [get(p) for p in diff]92        npar = r["maps"][f"{same[0]}|{fam}"]["n_params"]93        rows.append(f"{label} & {npar} & {fmt(np.mean(a))} & {fmt(np.mean(b))} \\\\")94    txt = f"""% source: experiments/summaries/{run}/results.json95\\begin{{tabular}}{{lrcc}}96\\toprule97Map & parameters & $R^2$ on $\\Omega$ rows, same init & $R^2$, different init \\\\98\\midrule99""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"100    (OUT / "tab_maps.tex").write_text(txt)101102103def table_mismatch(r, run):104    rows = []105    for cond, d in r["conditions"].items():106        if cond.startswith("mismatch_") or cond.startswith("stride2_"):107            for p, v in d.items():108                rows.append(f"{p.replace('_', '\\_')} & {cond.replace('_', '\\_')} & {fmt(v['target_iid']['accuracy'])} & {fmt(v['target_iid'].get('omega_step_acc'))} & {fmt(v['control_iid']['accuracy'])} & {fmt(v.get('r2_fit'))} \\\\")109    for key in ("recipient_w96_s0", "recipient_L6_s0", "recipient_s2_s0", "donor_s2_s0"):110        if key in r["conditions"]["base"]:111            v = r["conditions"]["base"][key]112            rows.insert(0, f"{key.replace('_', '\\_')} & base & {fmt(v['target_iid']['accuracy'])} & {fmt(v['target_iid'].get('omega_step_acc'))} & {fmt(v['control_iid']['accuracy'])} & --- \\\\")113    txt = f"""% source: experiments/summaries/{run}/results.json114\\begin{{tabular}}{{llcccc}}115\\toprule116Pair & condition & target & $\\Omega$-step & control & map $R^2$ \\\\117\\midrule118""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"119    (OUT / "tab_mismatch.tex").write_text(txt)120121122def table_mechanism(r, run):123    rows = []124    for cond in ("mech_donor_ablated_attn0", "mech_rescue_from_intact_donor", "mech_rescue_sham", "mech_reverse_recipient_to_donor"):125        for p, v in r["conditions"].get(cond, {}).items():126            rows.append(f"{cond.replace('mech_', '').replace('_', ' ')} & {p.replace('_', '\\_')} & {fmt(v['target_iid']['accuracy'])} & {fmt(v['target_iid'].get('omega_step_acc'))} & {fmt(v['control_iid']['accuracy'])} \\\\")127    h5 = r["rule_vs_cache"]128    for k, v in h5.items():129        if isinstance(v, dict) and "seen" in v:130            rows.append(f"rule vs cache: {k.replace('h5_', '').replace('_', ' ')} & seen / unseen $\\Omega$ inputs & {fmt(v['seen']['estimate'])} [{fmt(v['seen']['low'])}, {fmt(v['seen']['high'])}] & {fmt(v['unseen']['estimate'])} [{fmt(v['unseen']['low'])}, {fmt(v['unseen']['high'])}] & --- \\\\")131    txt = f"""% source: experiments/summaries/{run}/results.json132\\begin{{tabular}}{{llccc}}133\\toprule134Test & pair / stratum & target & $\\Omega$-step & control \\\\135\\midrule136""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"137    (OUT / "tab_mechanism.tex").write_text(txt)138139140def table_tier1(r, run, name="tab_tier1.tex"):141    rows = []142    for p, pr in r["pairs"].items():143        for tn, tr in pr["tasks"].items():144            conds = tr["conditions"]145            best_name, best = max(((k, v) for k, v in conds.items() if not k.startswith("ctl") and "random_orth" not in k), key=lambda kv: kv[1]["accuracy"])146            ctl_name, ctl = max(((k, v) for k, v in conds.items() if k.startswith("ctl") or "random_orth" in k), key=lambda kv: kv[1]["accuracy"])147            g = best["gain_vs_zero"]148            rows.append(f"{p.replace('_', '\\_')} & {tn} & {fmt(tr['recipient_zero'],2)} & {fmt(tr['recipient_few'],2)} & {fmt(tr['recipient_own_vector'],2)} & {fmt(best['accuracy'],2)} & {best_name.replace('_', '\\_').replace('|', ' ')} & {fmt(ctl['accuracy'],2)} & {g['estimate']:+.2f} [{g['low']:+.2f}, {g['high']:+.2f}] \\\\")149    txt = f"""% source: experiments/summaries/{run}/results.json150\\begin{{tabular}}{{llccccccc}}151\\toprule152Pair & task & zero-shot & few-shot & own vector & best transfer & condition & best control & gain vs zero [95\\% CI] \\\\153\\midrule154""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"155    (OUT / name).write_text(txt)156157158def table_tier1_localization(r, run, name="tab_tier1_localization.tex"):159    rows = []160    for mk, m in r["models"].items():161        loc = m.get("localization")162        if not loc:163            continue164        for tn, t in loc["tasks"].items():165            rows.append(f"{mk} ({m['n_layers']} layers, $d={m['hidden']}$) & {tn} & {fmt(t['zero_shot'],2)} & {fmt(t['few_shot'],2)} & L{t['best_layer']} & {fmt(t['sufficiency_by_layer'][str(t['best_layer'])],2)} & {fmt(t['random_vector_same_norm_at_best'],2)} & {fmt(t['necessity_fewshot_minus_direction'],2)} \\\\")166    txt = f"""% source: experiments/summaries/{run}/results.json167\\begin{{tabular}}{{llcccccc}}168\\toprule169Model & task & zero-shot & few-shot & best layer & vector added (zero-shot) & random vector & few-shot minus direction \\\\170\\midrule171""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"172    (OUT / name).write_text(txt)173174175def table_tier1_confirmatory(run):176    c = json.loads((ROOT / "experiments" / "summaries" / run / "tier1_confirmatory.json").read_text())177    rows = []178    for p, pr in c["pairs"].items():179        for tn, t in pr["tasks"].items():180            if "error" in t:181                continue182            el = "yes" if t["eligible"] else "no"183            rows.append(f"{p.replace('->', ' $\\to$ ')} & {tn} & {el} & {t['selected'].replace('_', '\\_').replace('|', ' ')} & {fmt(t['conf_zero'],2)} & {fmt(t['conf_few'],2)} & {fmt(t['conf_own'],2)} & {fmt(t['conf_selected'],2)} & {fmt(t['best_control_acc'],2)} & {t['diff']:+.2f} [{t['ci'][0]:+.2f}, {t['ci'][1]:+.2f}]{' $\\checkmark$' if t['supported'] else ''} & {fmt(max(t['shuffled_query_diagnostic'].values()) if t['shuffled_query_diagnostic'] else None,2)} \\\\")184        rows.append(f"\\multicolumn{{11}}{{l}}{{\\emph{{{p.replace('->', ' $\\to$ ')}: H6 {pr['H6']} ({pr['n_supported_tasks']} task(s))}}}} \\\\ \\addlinespace")185    txt = f"""% source: experiments/summaries/{run}/tier1_confirmatory.json (scripts/tier1_analysis.py)186\\begin{{tabular}}{{llllcccccll}}187\\toprule188Pair & task & eligible & selected condition (dev) & zero & few & own & selected & best null control & selected $-$ control [95\\% CI] & shuffled-query diag. \\\\189\\midrule190""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"191    (OUT / f"tab_tier1_confirmatory_{run.split('_')[-1][:6]}.tex").write_text(txt)192193194def table_pilot(r, run):195    rows = []196    for en in ("target_iid", "control_iid", "comp_heldout", "len_ood"):197        d, rc = r["across_seeds"][en]["donor"], r["across_seeds"][en]["recipient"]198        g = r["seeds"]["0"]["gap"][en]199        rows.append(f"{en.replace('_', '\\_')} & {fmt(r['chance_level'][en])} & {fmt(d['mean'])} ({fmt(d['sd'] or 0,3)}) & {fmt(rc['mean'])} ({fmt(rc['sd'] or 0,3)}) & {g['estimate']:+.3f} [{g['low']:+.3f}, {g['high']:+.3f}] \\\\")200    txt = f"""% source: experiments/summaries/{run}/results.json201\\begin{{tabular}}{{lcccc}}202\\toprule203Stratum & chance & donor (sd) & recipient (sd) & paired gap, seed 0 [95\\% CI] \\\\204\\midrule205""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"206    (OUT / "tab_pilot_gap.tex").write_text(txt)207208209def table_models(index_path: Path):210    idx = json.loads(index_path.read_text())211    rows = []212    for k, v in idx.items():213        m = v["model"]214        t = v["task"]215        rows.append(f"{k.replace('_', '\\_')} & {m['n_layers']} & {m['d_model']} & {m['n_heads']} & {', '.join(t['excluded_operators']) or '---'} & {'vex' if 'vex' in t['operators'] else '---'} & {t.get('trace_stride', 1)} & {fmt(v['final_dev_acc'])} & {v['train_hash']} \\\\")216    txt = f"""% source: {index_path.relative_to(ROOT)}217\\begin{{tabular}}{{lcccccccl}}218\\toprule219Model & layers & $d$ & heads & excluded & extra op. & stride & dev acc. & train split hash \\\\220\\midrule221""" + "\n".join(rows) + "\n\\bottomrule\n\\end{tabular}\n"222    (OUT / "tab_models.tex").write_text(txt)223224225def main() -> None:226    ap = argparse.ArgumentParser()227    ap.add_argument("--graft")228    ap.add_argument("--pilot")229    ap.add_argument("--tier1", nargs="*", default=[])230    ap.add_argument("--confirmatory")231    ap.add_argument("--tier1-confirmatory", nargs="*", default=[])232    args = ap.parse_args()233    OUT.mkdir(parents=True, exist_ok=True)234    if args.graft:235        r = load(args.graft)236        table_main(r, args.graft)237        table_maps(r, args.graft)238        table_mismatch(r, args.graft)239        table_mechanism(r, args.graft)240    if args.confirmatory:241        table_main(load(args.confirmatory), args.confirmatory, "tab_confirmatory.tex")242    if args.pilot:243        table_pilot(load(args.pilot), args.pilot)244    for run in args.tier1:245        r = load(run)246        suf = run.split("_")[-1][:6]247        table_tier1(r, run, f"tab_tier1_{suf}.tex")248        table_tier1_localization(r, run, f"tab_tier1_localization_{suf}.tex")249    for run in args.tier1_confirmatory:250        table_tier1_confirmatory(run)251    idx = ROOT / "artifacts" / "models" / "matrix_v1" / "index.json"252    if idx.exists():253        table_models(idx)254    print("tables written:", sorted(p.name for p in OUT.glob("*.tex")))255256257if __name__ == "__main__":258    main()259