Blame for tables/scripts/paper_numbers.py

"""Main manuscript tables and headline values, recomputed from the artifact.

Run this and diff against the draft. Validation-checkpoint epochs and epoch-matched
shares are emitted by the companion ``checkpoint_selection_audit.py``.

    python scripts/paper_numbers.py            # human-readable
    python scripts/paper_numbers.py --json     # machine-readable

Sources (all under ../drift_run unless noted):
    drift_out/drift_by_stage.csv       per (chain, layer, stage) drift, 3 inner products
    drift_out/drift_per_feature.csv    per-latent drift, density, persistence
    drift_out/base_feature_z.csv       per-latent logit-lens peakedness at the base SAE
    drift_out/null_replicates.csv      null-step replicates + real first transitions
    drift_out/null_trajectory_l18.csv  the originally reported 7-step null trajectory
    drift_out/null_trajectory_l18_seeds.csv  the same, 4 seeds (supersedes the above)
    null_chain_all/layer{L}_seed{S}/   per-layer null trajectories, 3 seeds each
    null_budget/                       epoch-budget sweep + budget-matched trajectories
    arch_null/arch_null_trajectory.csv TopK / BatchTopK / JumpReLU at matched recipe
    saes_local/, null_ctrl/, null_seeds/  SAE weights (for cosines and hashes)
    ../sae_rl/results/sae_checkpoint_metrics.csv    reconstruction fidelity
    ../sae_rl/results/task_performance.csv          GSM8K accuracy
    ../sae_rl/data/activations*/                    activation-set sizes

Sections needing SAE weights (direction, budget, trajectories, alignment, persistent
null) run on CPU by default and are skipped with a warning if the weights are absent,
so the script still works from the CSVs alone.
"""

import argparse
import glob
import json
import os
import pathlib
import re

import pandas as pd
from scipy.stats import spearmanr

DEV = "cpu"

HERE = pathlib.Path(__file__).resolve().parent


def _find(*parts):
    """Locate a data directory whether this paper tree sits inside sae_rl/ or beside it.

    The paper moved out of sae_rl/paper/ on 2026-08-20, so the data root is one
    level further away than it used to be. Try both rather than pinning either.
    """
    for base in (HERE.parent.parent, HERE.parent.parent / "sae_rl"):
        cand = base.joinpath(*parts)
        if cand.exists():
            return cand
    return HERE.parent.parent.joinpath(*parts)


DRIFT = _find("drift_run", "drift_out")
RESULTS = _find("sae_rl", "results")

ACC = "val-core/openai/gsm8k/acc/mean@1"
# Chains sharing a bit-identical base SAE, so cross-chain latent indices are
# comparable. Verified by hashing W_dec; see base_alignment().
ALIGNED = ["sae_sft", "sae_flexible", "sae_strict/k64"]


def _round(o, n=4):
    if isinstance(o, dict):
        return {k: _round(v, n) for k, v in o.items()}
    if isinstance(o, list):
        return [_round(v, n) for v in o]
    return round(o, n) if isinstance(o, float) else o


def artifact(drift):
    """Size of the artifact: SAEs, chains, layers."""
    return {
        "n_saes": int(len(drift)),
        "n_chains": int(drift.chain.nunique()),
        "layers": sorted(drift.layer.unique().tolist()),
        "chain_n_checkpoints": {c: int(g.stage.nunique())
                                for c, g in drift.groupby("chain")},
        "d_sae": int(drift.n_features.iloc[0]),
    }


def identity_vs_fidelity(drift, metrics):
    """Inherited-index alignment turns over while reconstruction stays flat."""
    out = {}
    fin = drift[(drift.chain == "sae_sft") & (drift.step == 348)]
    for _, r in fin.iterrows():
        out[f"L{int(r.layer)}"] = {
            "dec_cos_geom": r.dec_geom_all_cos_matched_mean,
            "dec_cos_readout": r.dec_readout_all_cos_matched_mean,
            "dec_cos_lens": r.dec_lens_all_cos_matched_mean,
            "tok_jaccard": r.tok_jaccard_matched_all,
            "frac_top1_identical": r.frac_identical_top1_all,
            "frac_top1_changed": 1 - r.frac_identical_top1_all,
        }
    f = metrics[(metrics.chain == "sft") & (metrics.layer == 18)]
    out["L18_fidelity_over_sft"] = {
        "nmse_min": f.nmse_recomputed.min(), "nmse_max": f.nmse_recomputed.max(),
        "frac_rec_min": f.frac_rec.min(), "frac_rec_max": f.frac_rec.max(),
        "dead_pct_max": f.dead_latents_pct.max(),
    }
    out["dead_pct_max_all_chains"] = metrics.dead_latents_pct.max()
    return out


def metric_agreement(drift):
    """The three inner products must agree, else the geometry story is metric-dependent."""
    cols = ["dec_geom_all_drift_norm", "dec_readout_all_drift_norm",
            "dec_lens_all_drift_norm"]
    d = drift[drift.step > 0].copy()
    d["spread"] = d[cols].max(axis=1) - d[cols].min(axis=1)
    per_cell = d.groupby(["chain", "layer"]).spread.max()
    return {
        "n_cells": int(len(per_cell)),
        "max_spread": per_cell.max(),
        "median_spread": per_cell.median(),
        "argmax_cell": "/".join(str(x) for x in per_cell.idxmax()),
    }


def null_step(reps):
    """How much of a real first transition a warm-started retraining reproduces."""
    nulls, out = {}, {}
    for L, g in reps[reps.kind == "nullstep"].groupby("layer"):
        nulls[L] = {"n": int(len(g)), "drift_mean": g.drift.mean(),
                    "drift_min": g.drift.min(), "drift_max": g.drift.max(),
                    "tokJ_mean": g.tok_jaccard.mean()}
    shares = []
    for _, r in reps[reps.kind == "real"].iterrows():
        n = nulls[r.layer]
        shares.append({
            "layer": int(r.layer), "chain": r.label,
            "real_drift": r.drift, "null_drift_mean": n["drift_mean"],
            "pct_mean": 100 * n["drift_mean"] / r.drift,
            "pct_lo": 100 * n["drift_min"] / r.drift,
            "pct_hi": 100 * n["drift_max"] / r.drift,
            # top-token turnover share: (1 - J_null) / (1 - J_real)
            "tok_pct": 100 * (1 - n["tokJ_mean"]) / (1 - r.tok_jaccard),
        })
    out["null_replicates"] = {f"L{k}": v for k, v in nulls.items()}
    out["share_of_real"] = shares
    # Headline range, excluding L12 (null replicates there are unstable: see
    # drift_min/drift_max) and excluding cells where the real transition drifts
    # *less* than the null step (pct > 100).
    core = [s for s in shares if s["layer"] != 12 and s["pct_mean"] <= 110]
    out["headline_pct_range"] = [min(s["pct_mean"] for s in core),
                                 max(s["pct_mean"] for s in core)]
    out["headline_tok_pct_range"] = [min(s["tok_pct"] for s in core),
                                     max(s["tok_pct"] for s in core)]
    out["exceeds_100"] = [f"{s['chain']}@L{s['layer']}={s['pct_mean']:.0f}%"
                          for s in shares if s["pct_mean"] > 110]
    return out


def null_trajectory(traj, drift):
    """A whole trajectory of retrainings, with the model held fixed throughout."""
    real = (drift[(drift.chain == "sae_sft") & (drift.layer == 18)]
            .sort_values("step").reset_index(drop=True))
    rows = []
    for i, r in traj.sort_values("null_step").iterrows():
        rr = real.iloc[i] if i < len(real) else None
        rows.append({
            "step_index": int(r.null_step),
            "null_dec_cos": r.dec_cos, "null_tok_jaccard": r.tok_jaccard,
            "real_step": int(rr.step) if rr is not None else None,
            "real_dec_cos": rr.dec_geom_all_cos_matched_mean if rr is not None else None,
            "real_tok_jaccard": rr.tok_jaccard_matched_all if rr is not None else None,
        })
    end = rows[-1]
    return {
        "per_step": rows,
        "final_null_cos": end["null_dec_cos"], "final_real_cos": end["real_dec_cos"],
        "final_null_tokJ": end["null_tok_jaccard"],
        "final_real_tokJ": end["real_tok_jaccard"],
        # Fraction of the real end-to-end drift reproduced with the model frozen.
        "pct_drift_reproduced": 100 * (1 - end["null_dec_cos"]) / (1 - end["real_dec_cos"]),
        "pct_tok_turnover_reproduced":
            100 * (1 - end["null_tok_jaccard"]) / (1 - end["real_tok_jaccard"]),
    }


def shuffled_control(drift, task):
    """A reward carrying no task information, at a matched optimiser step."""
    out = {"accuracy": {}}
    for c, g in task.groupby("chain"):
        g = g.sort_values("step")
        out["accuracy"][c] = {"start": g[ACC].iloc[0], "final": g[ACC].iloc[-1],
                              "max": g[ACC].max()}
    for step in (10, 30):
        p = (drift[drift.step == step]
             .pivot_table(index="chain", columns="layer",
                          values="dec_geom_all_drift_norm"))
        cell = {}
        for ref in ("sae_flexible", "sae_strict/k64"):
            ratios = (p.loc["sae_shuffled"] / p.loc[ref])
            cell[ref] = {f"L{int(k)}": v for k, v in ratios.items()}
        out[f"step{step}_shuffled_over_informative"] = cell
        out[f"step{step}_drift_norm"] = {
            ch: {f"L{int(k)}": v for k, v in row.items()} for ch, row in p.iterrows()}
    return out


def heterogeneity(per_feat, z):
    """Stability is not uniform, and our selection heuristic rides the tail."""
    d = per_feat.merge(z, on=["chain", "layer", "feature"])
    out = {"spearman_final_sft": {}, "selection_bias_final_sft": {}}
    for L in (6, 12, 18, 23):
        s = d[(d.chain == "sae_sft") & (d.layer == L) & (d.step == 348)]
        out["spearman_final_sft"][f"L{L}"] = {
            "cos_vs_density": spearmanr(s.dec_geom_cos, s.density_base).statistic,
            "cos_vs_peakedness": spearmanr(s.dec_geom_cos, s.logit_z).statistic,
            "tokJ_vs_density": spearmanr(s.tok_jaccard, s.density_base).statistic,
            "tokJ_vs_peakedness": spearmanr(s.tok_jaccard, s.logit_z).statistic,
        }
        # The sweep's own heuristic: frequency band, then top-n by peakedness.
        pool = s[(s.density_base > 0.05) & (s.density_base < 0.5)]
        sel = pool.nlargest(12, "logit_z")
        out["selection_bias_final_sft"][f"L{L}"] = {
            "n_pool": int(len(pool)),
            "sel_tokJ": sel.tok_jaccard.mean(), "pop_tokJ": s.tok_jaccard.mean(),
            "tokJ_ratio": sel.tok_jaccard.mean() / s.tok_jaccard.mean(),
            "sel_cos": sel.dec_geom_cos.mean(), "pop_cos": s.dec_geom_cos.mean(),
        }
    r = [v["tokJ_ratio"] for v in out["selection_bias_final_sft"].values()]
    out["tokJ_ratio_range"] = [min(r), max(r)]
    return out


def _torch():
    import torch
    return torch


def _wdec(path):
    """Decoder directions [d_sae, d_model] from any of the on-disk SAE formats."""
    import sys
    sys.path.insert(0, str(_find("sae_tuned_lens").parent))
    from sae_tuned_lens.sae_io import load_sae_full
    return load_sae_full(str(path))["W_dec"].float()


def _cos(A, B):
    return _torch().nn.functional.cosine_similarity(A, B, dim=-1)


ROOT = _find("drift_run")
SFT_SAES = _find("sae_rl", "checkpoints", "saes_sft")


def direction_of_drift():
    """Does the null step move the basis the same WAY, or only as far? (§4.5)

    A share of drift is a magnitude claim. If retraining and fine-tuning rotated
    latents by similar amounts into unrelated subspaces, the shares would hold while
    retraining explained none of the real chain's particular movement.
    """
    torch = _torch()
    base = _wdec(ROOT / "saes_local/sae_sft/sae_instruct_base_layer18.pt")
    out = {}
    cases = [
        ("one_retraining", ROOT / "null_ctrl/sae_instruct_base_layer18_nullstep.pt",
         ROOT / "saes_local/sae_sft/sae_sft_step29_layer18.pt"),
        ("seven_retrainings", ROOT / "null_chain/sae_nullchain7_layer18.pt",
         ROOT / "saes_local/sae_sft/sae_sft_step348_layer18.pt"),
    ]
    for label, npath, rpath in cases:
        if not (npath.exists() and rpath.exists()):
            out[label] = "weights absent"
            continue
        nul, real = _wdec(npath), _wdec(rpath)
        c_bn = _cos(base, nul).mean().item()
        c_br = _cos(base, real).mean().item()
        c_nr = _cos(nul, real).mean().item()
        n = base.shape[0]
        g = torch.Generator().manual_seed(0)
        perm = torch.randperm(n, generator=g)
        idx = torch.arange(n)
        perm[perm == idx] = (perm[perm == idx] + 1) % n
        out[label] = {
            "cos_base_null": c_bn, "cos_base_real": c_br,
            "cos_null_real_observed": c_nr,
            "cos_null_real_if_independent": c_bn * c_br,
            "random_pairing_floor": _cos(nul, real[perm]).mean().item(),
            "excess_over_independent": c_nr - c_bn * c_br,
        }
    return out


def activation_budgets():
    """Optimisation budget of the null step vs each chain's own transitions (§3.3)."""
    torch = _torch()
    act = _find("sae_rl", "data")
    out = {}
    for sub, stem in [("activations", "instruct_base_layer18_train.pt"),
                      ("activations", "ppo_step10_layer18_train.pt"),
                      ("activations_sft", "sft_step29_layer18_train.pt"),
                      ("activations_shuffled", "ppo_step10_layer18_train.pt"),
                      ("activations_kl0p025", "ppo_step10_layer18_train.pt")]:
        p = act / sub / stem
        if not p.exists():
            continue
        n = torch.load(p, map_location="cpu", mmap=True).shape[0]
        out[f"{sub}/{stem}"] = {"tokens": int(n), "batches_at_512": n // 512,
                                "adam_steps_20_epochs": (n // 512) * 20}
    base = out.get("activations/instruct_base_layer18_train.pt")
    sft = out.get("activations_sft/sft_step29_layer18_train.pt")
    if base and sft:
        out["null_step_shortfall_vs_sft_pct"] = 100 * (
            1 - base["adam_steps_20_epochs"] / sft["adam_steps_20_epochs"])
        out["epochs_to_match_sft"] = 20 * sft["batches_at_512"] / base["batches_at_512"]
    return out


def budget_sweep(drift):
    """Nominal budget sweep, interpreted as LR-horizon sensitivity (§4.4)."""
    real = drift[(drift.chain == "sae_sft") & (drift.layer == 18)].sort_values("step")
    real_drift = 1 - real.dec_geom_all_cos_matched_mean.iloc[1]
    base = ROOT / "saes_local/sae_sft/sae_instruct_base_layer18.pt"
    if not base.exists():
        return "weights absent"
    W0 = _wdec(base)
    rows = {}
    for p in sorted(glob.glob(str(ROOT / "null_budget/sae_instruct_base_layer18_e*_s*.pt"))):
        m = re.search(r"_e(\d+)_s(\d+)", p)
        rows.setdefault(int(m.group(1)), []).append(1 - _cos(W0, _wdec(p)).mean().item())
    out = {"real_first_transition_drift": real_drift, "by_epochs": {}}
    for E in sorted(rows):
        v = rows[E]
        out["by_epochs"][E] = {
            "adam_steps": E * 696, "n_seeds": len(v),
            "drift_mean": sum(v) / len(v), "drift_min": min(v), "drift_max": max(v),
            "pct_of_real": 100 * (sum(v) / len(v)) / real_drift,
        }
    return out


def trajectories_per_layer(drift):
    """Null trajectories against the real SFT chain, per layer, over seeds (§4.4)."""
    out = {}
    for L in (6, 12, 18, 23):
        base = ROOT / f"saes_local/sae_sft/sae_instruct_base_layer{L}.pt"
        if not base.exists():
            continue
        dirs = [d for d in sorted(glob.glob(str(ROOT / f"null_chain_all/layer{L}_seed*")))
                if os.path.exists(os.path.join(d, "DONE"))]
        if L == 18:  # the originally reported chain belongs to this population
            if (ROOT / "null_chain").exists():
                dirs = [str(ROOT / "null_chain")] + dirs
        if not dirs:
            continue
        W0 = _wdec(base)
        real = (drift[(drift.chain == "sae_sft") & (drift.layer == L)]
                .sort_values("step").reset_index(drop=True))
        rf = real.dec_geom_all_cos_matched_mean.iloc[-1]
        ends = []
        for d in dirs:
            ch = sorted(glob.glob(f"{d}/sae_nullchain*_layer{L}.pt"),
                        key=lambda q: int(re.search(r"nullchain(\d+)", q).group(1)))
            if len(ch) < len(real):
                continue
            ends.append(_cos(W0, _wdec(ch[len(real) - 1])).mean().item())
        if not ends:
            continue
        pct = [100 * (1 - e) / (1 - rf) for e in ends]
        out[f"L{L}"] = {
            "n_chains": len(ends), "real_end_cos": rf,
            "null_end_cos_mean": sum(ends) / len(ends),
            "null_end_cos_min": min(ends), "null_end_cos_max": max(ends),
            "pct_reproduced_mean": sum(pct) / len(pct),
            "pct_reproduced_min": min(pct), "pct_reproduced_max": max(pct),
        }
    return out


def persistent_null(drift, per_feat):
    """The null step restricted to evaluation-persistent latents (§4.3).

    The mask is based on firing density over shared frozen-base evaluation
    activations; it does not record whether a latent was resampled in training.

    Without this, a persistent-subset real drift gets divided by an all-latent null and
    the share is wrong.
    """
    torch = _torch()
    first = {"sae_sft": 29, "sae_flexible": 10, "sae_strict/k64": 10}
    out = {}
    for L in (6, 12, 18, 23):
        base = ROOT / f"saes_local/sae_sft/sae_instruct_base_layer{L}.pt"
        if not base.exists():
            continue
        W0 = _wdec(base)
        # Every seed on disk, not a fixed subset: the seed count is what makes the
        # published shares reproducible, and s10/s11 sort after s9 only numerically.
        paths = [ROOT / f"null_ctrl/sae_instruct_base_layer{L}_nullstep.pt"]
        paths += [pathlib.Path(p) for p in
                  sorted(glob.glob(str(ROOT / f"null_seeds/sae_instruct_base_layer{L}_null_s*.pt")),
                         key=lambda q: int(re.search(r"_null_s(\d+)", q).group(1)))]
        cs = [_cos(W0, _wdec(p)) for p in paths if p.exists()]
        if not cs:
            continue
        for ch, st in first.items():
            row = drift[(drift.chain == ch) & (drift.layer == L) & (drift.step == st)]
            m = per_feat[(per_feat.chain == ch) & (per_feat.layer == L) & (per_feat.step == st)]
            if not len(row) or not len(m):
                continue
            row = row.iloc[0]
            idx = torch.tensor(m.sort_values("feature").persistent.values.astype(bool))
            na = sum(1 - c.mean().item() for c in cs) / len(cs)
            npr = sum(1 - c[idx].mean().item() for c in cs) / len(cs)
            ra = 1 - row.dec_geom_all_cos_matched_mean
            rp = 1 - row.dec_geom_persist_cos_matched_mean
            out[f"L{L}/{ch}"] = {
                "n_persist": int(idx.sum()), "null_all": na, "null_persist": npr,
                "real_all": ra, "real_persist": rp,
                "pct_all": 100 * na / ra, "pct_persist": 100 * npr / rp,
            }
    return out


def architectures():
    """Does the confound depend on the SAE parameterisation? (Jacob's grid, L12.)"""
    p = ROOT / "arch_null/arch_null_trajectory.csv"
    if not p.exists():
        return "not run"
    a = pd.read_csv(p)
    last = a.null_step.max()
    return {r.arm: {"arch": r.arch, "final_dec_cos": r.dec_cos,
                    "final_tok_jaccard": r.tok_jaccard, "mean_l0": r.mean_l0,
                    "val_mse": r.val_mse, "steps": int(last)}
            for _, r in a[a.null_step == last].iterrows()}


def fidelity_detail(metrics, drift):
    """The Table 1 / §4.1 numbers the earlier version of this script did not emit."""
    fin = metrics[(metrics.chain == "sft")].copy()
    out = {"sft_final_nmse_by_layer": {}, "sft_mse_base_to_final": {},
           "persist_counts": {}, "layer23_nmse_all_runs": {}}
    for L in (6, 12, 18, 23):
        s = fin[fin.layer == L].sort_values("step")
        if not len(s):
            continue
        out["sft_final_nmse_by_layer"][f"L{L}"] = s.nmse_recomputed.iloc[-1]
        out["sft_mse_base_to_final"][f"L{L}"] = [s.mse.iloc[0], s.mse.iloc[-1]]
    n = drift.n_persist
    out["persist_counts"] = {"min": int(n.min()), "max": int(n.max()),
                             "d_sae": int(drift.n_features.iloc[0])}
    l23 = metrics[metrics.layer == 23]
    out["layer23_nmse_all_runs"] = {"min": l23.nmse_recomputed.min(),
                                    "max": l23.nmse_recomputed.max()}
    return out


def base_alignment():
    """Cross-chain latent indices are only comparable if the base SAE is shared."""
    import hashlib
    import sys
    sys.path.insert(0, str(_find("sae_tuned_lens").parent))
    from sae_tuned_lens.sae_io import load_sae_full

    root = _find("drift_run", "saes_local")
    chains = ["sae_sft", "sae_flexible", "sae_strict/k64", "sae_shuffled", "sae_kl0p025"]
    out = {}
    for L in (6, 12, 18, 23):
        hashes = {}
        for c in chains:
            p = root / c / f"sae_instruct_base_layer{L}.pt"
            if not p.exists():
                hashes[c] = None
                continue
            W = load_sae_full(str(p))["W_dec"].float().cpu().numpy()
            hashes[c] = hashlib.md5(W.tobytes()).hexdigest()[:12]
        ref = hashes["sae_sft"]
        out[f"L{L}"] = {"hashes": hashes,
                        "shares_base": [c for c, h in hashes.items() if h == ref],
                        "own_init": [c for c, h in hashes.items()
                                     if h is not None and h != ref]}
    always_own = [c for c in chains
                  if all(c in out[f"L{L}"]["own_init"] for L in (6, 12, 18, 23))]
    out["own_init_at_every_layer"] = always_own
    out["n_chains_own_init"] = len(always_own)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--json", action="store_true")
    ap.add_argument("--skip-weights", "--skip-alignment", dest="skip_weights",
                    action="store_true",
                    help="skip every section that loads SAE weights (hashes, drift "
                         "direction, budget sweep, per-layer trajectories, persistent null)")
    args = ap.parse_args()

    drift = pd.read_csv(DRIFT / "drift_by_stage.csv")
    per_feat = pd.read_csv(DRIFT / "drift_per_feature.csv")
    z = pd.read_csv(DRIFT / "base_feature_z.csv")
    reps = pd.read_csv(DRIFT / "null_replicates.csv")
    traj = pd.read_csv(DRIFT / "null_trajectory_l18.csv")
    metrics = pd.read_csv(RESULTS / "sae_checkpoint_metrics.csv")
    task = pd.read_csv(RESULTS / "task_performance.csv")

    report = {
        "artifact": artifact(drift),
        "identity_vs_fidelity": identity_vs_fidelity(drift, metrics),
        "fidelity_detail": fidelity_detail(metrics, drift),
        "metric_agreement": metric_agreement(drift),
        "null_step": null_step(reps),
        "null_trajectory_l18": null_trajectory(traj, drift),
        "shuffled_control": shuffled_control(drift, task),
        "heterogeneity": heterogeneity(per_feat, z),
        "activation_budgets": activation_budgets(),
        "architectures": architectures(),
    }
    if not args.skip_weights:
        report["direction_of_drift"] = direction_of_drift()
        report["budget_sweep"] = budget_sweep(drift)
        report["trajectories_per_layer"] = trajectories_per_layer(drift)
        report["persistent_null"] = persistent_null(drift, per_feat)
        report["base_alignment"] = base_alignment()

    report = _round(report)
    if args.json:
        print(json.dumps(report, indent=2, default=str))
    else:
        for section, body in report.items():
            print(f"\n{'=' * 70}\n{section}\n{'=' * 70}")
            print(json.dumps(body, indent=2, default=str))


if __name__ == "__main__":
    main()