MachineSex/figures/stats_llm_curriculum.py

"""Statistics for the six-generation language-model population and its two controls.

One loader for every curriculum bundle (the arm label is set HERE by directory, never trusted from
the parquet alone, because the veto arm is recorded as ``society`` with ``allow_veto`` on), and the
pre-registered readouts for the two 2026-09-11 controls (tasks/prereg-llm-society-v4.md on the dev branch §8h):

1. **Forced stop at generation 3** (``llm_curriculum_v5_stop3``): per-seed paired contrasts of the
   best lineage's final all-family accuracy, veto − stop3, stop3 − isolated, stop3 − society.
2. **Decorrelated curriculum** (``llm_curriculum_v5_decor``): partial Spearman correlation of the
   fraction of merges declined with partner complementarity, controlling for generation, pooled over
   both curricula (Latin square + decorrelated), with a seed-clustered bootstrap CI; and the mirror
   partial correlation with generation controlling for complementarity.

Reads committed artifacts only. Missing bundles are skipped, so the script runs at any stage of the
campaign and reports what exists.

Usage:  python figures/stats_llm_curriculum.py
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
from scipy.stats import rankdata, spearmanr

ROOT = Path(__file__).resolve().parents[1]
RES = ROOT / "results"
SEEDS = (1, 2, 3)

# experiment directory -> (curriculum label, {recorded arm -> reported arm})
RELABEL = {
    "llm_curriculum_v5": ("latin", {}),
    "llm_curriculum_v5_veto": ("latin", {"society": "veto"}),
    "llm_curriculum_v5_stop3": ("latin", {"society": "society_stop3"}),
    "llm_curriculum_v5_decor": ("decor", {"isolated": "decor_isolated", "society": "decor_veto"}),
    # conflict-arrival curricula (2026-09-12): conflicting pair first (early) or last (late)
    "llm_curriculum_v5_early": ("early", {"isolated": "early_isolated", "society": "early_veto"}),
    "llm_curriculum_v5_early_obl": ("early", {"society": "early_society"}),
    "llm_curriculum_v5_late": ("late", {"isolated": "late_isolated", "society": "late_veto"}),
    "llm_curriculum_v5_late_obl": ("late", {"society": "late_society"}),
    # differential reproduction (2026-09-12): Latin square with truncation selection
    "llm_curriculum_v5_cull": ("latin", {"isolated": "cull_isolated", "society": "cull_veto"}),
}
VETO_ARMS = ("veto", "decor_veto", "early_veto", "late_veto", "cull_veto")
# Generation (0-based) from which BOTH conflicting families (boolq yes/no, winogrande 1/2) are
# present in every lineage of each curriculum: the conflict_present indicator of the timing test.
CONFLICT_FROM = {"latin": 4, "decor": 3, "early": 1, "late": 5}


def _bundles(exp: str) -> list[tuple[int, Path]]:
    """(seed, parquet) pairs for one experiment directory, in every layout the campaign used.

    ``s1/`` or top-level for seed 1 (local runs), ``s{seed}/`` or ``s{seed}_<arm>/`` for the HPC array
    elements. The seed is read from the frame itself, so the directory name only locates the file.
    """
    d = RES / exp
    if not d.exists():
        return []
    out = []
    for p in sorted(d.glob("results.parquet")) + sorted(d.glob("s[0-9]*/results.parquet")):
        seeds = pd.read_parquet(p, columns=["seed"])["seed"].unique()
        out += [(int(s), p) for s in seeds]
    return out


def load_curriculum() -> pd.DataFrame:
    """Every curriculum bundle as one long-form frame with ``curriculum`` and relabelled ``arm``."""
    frames = []
    for exp, (curriculum, relabel) in RELABEL.items():
        for _, p in _bundles(exp):
            d = pd.read_parquet(p)
            d["arm"] = d["arm"].map(lambda a: relabel.get(a, a))
            d["curriculum"] = curriculum
            frames.append(d)
    if not frames:
        raise FileNotFoundError("no curriculum bundles under results/")
    return pd.concat(frames, ignore_index=True).drop_duplicates(
        ["curriculum", "arm", "seed", "generation", "model", "metric"])


def best_lineage(df: pd.DataFrame, metric: str = "all_families") -> pd.DataFrame:
    """Best lineage per (curriculum, arm, seed, generation) on ``metric`` (the paper's readout)."""
    sub = df[(df["metric"] == metric) & (df["generation"] >= 0) & df["model"].str.startswith("lineage")]
    return sub.groupby(["curriculum", "arm", "seed", "generation"])["value"].max().reset_index()


def final_contrasts(best: pd.DataFrame, pairs: list[tuple[str, str]]) -> pd.DataFrame:
    """Per-seed paired differences at the final generation, one row per contrast."""
    g_last = best["generation"].max()
    fin = best[best["generation"] == g_last].pivot_table(index="seed", columns="arm", values="value")
    rows = []
    for a, b in pairs:
        if a not in fin or b not in fin:
            continue
        d = (fin[a] - fin[b]).dropna()
        rows.append({"contrast": f"{a} − {b}", "n_seeds": len(d),
                     **{f"s{s}": round(v, 3) for s, v in d.items()},
                     "mean": round(d.mean(), 3),
                     "ci95": round(1.96 * d.std(ddof=1) / np.sqrt(len(d)), 3) if len(d) > 1 else np.nan})
    return pd.DataFrame(rows)


def decline_table(df: pd.DataFrame) -> pd.DataFrame:
    """Mean fraction of merges declined per (curriculum, seed, generation), with complementarity."""
    veto_arms = df["arm"].isin(list(VETO_ARMS))
    v = (df[veto_arms & (df["metric"] == "veto_used")]
         .groupby(["curriculum", "seed", "generation"])["value"].mean().rename("declined"))
    c = (df[veto_arms & (df["metric"] == "complementarity")]
         .groupby(["curriculum", "seed", "generation"])["value"].mean().rename("complementarity"))
    tab = pd.concat([v, c], axis=1).dropna().reset_index()
    tab["conflict_present"] = (tab["generation"] >= tab["curriculum"].map(CONFLICT_FROM)).astype(float)
    return tab


def conflict_timing_test(tab: pd.DataFrame, B: int = 4000, seed: int = 0) -> dict:
    """Does the decline rate track the ARRIVAL of conflicting conventions once generation is
    controlled? Partial ρ(declined, conflict_present | generation) pooled over the curricula in
    ``tab`` (the early/late pair decorrelates the two by design), seed-clustered bootstrap CI."""
    rng = np.random.default_rng(seed)
    seeds = tab["seed"].unique()
    x, c, z = tab["declined"], tab["conflict_present"], tab["generation"]
    out = {"n_points": len(tab), "n_curricula": tab["curriculum"].nunique(),
           "rho_partial_conflict": partial_spearman(c, x, z),
           "rho_partial_generation": partial_spearman(z, x, c),
           "rho_raw_conflict": float(spearmanr(c, x)[0])}
    if len(seeds) > 1:
        groups = {s: tab[tab["seed"] == s] for s in seeds}
        boots = []
        for _ in range(B):
            bs = pd.concat([groups[s] for s in rng.choice(seeds, size=len(seeds), replace=True)])
            boots.append(partial_spearman(bs["conflict_present"], bs["declined"], bs["generation"]))
        boots = np.array(boots)
        out["ci95_partial_conflict"] = (float(np.nanpercentile(boots, 2.5)),
                                        float(np.nanpercentile(boots, 97.5)))
    return out


def partial_spearman(x, y, z) -> float:
    """Spearman correlation of x and y after rank-regressing both on z."""
    rx, ry, rz = rankdata(x), rankdata(y), rankdata(z)
    Z = np.column_stack([np.ones_like(rz), rz])
    res = lambda r: r - Z @ np.linalg.lstsq(Z, r, rcond=None)[0]
    return float(spearmanr(res(rx), res(ry))[0])


def decline_test(tab: pd.DataFrame, B: int = 4000, seed: int = 0) -> dict:
    """The pre-registered primary readout: partial ρ(declined, complementarity | generation), pooled
    over curricula, with a seed-clustered percentile bootstrap; plus the mirror partial correlation."""
    rng = np.random.default_rng(seed)
    seeds = tab["seed"].unique()
    x, y, z = tab["declined"], tab["complementarity"], tab["generation"]
    out = {"n_points": len(tab), "n_curricula": tab["curriculum"].nunique(), "n_seeds": len(seeds),
           "rho_partial_complementarity": partial_spearman(y, x, z),
           "rho_partial_generation": partial_spearman(z, x, y),
           "rho_raw_complementarity": float(spearmanr(y, x)[0]),
           "rho_raw_generation": float(spearmanr(z, x)[0])}
    if len(seeds) > 1:
        boots = []
        groups = {s: tab[tab["seed"] == s] for s in seeds}
        for _ in range(B):
            bs = pd.concat([groups[s] for s in rng.choice(seeds, size=len(seeds), replace=True)])
            boots.append(partial_spearman(bs["complementarity"], bs["declined"], bs["generation"]))
        boots = np.array(boots)
        out["ci95_partial_complementarity"] = (float(np.nanpercentile(boots, 2.5)),
                                               float(np.nanpercentile(boots, 97.5)))
    return out


def main() -> None:
    df = load_curriculum()
    best = best_lineage(df)
    print("bundles loaded — arms × seeds:")
    print(best.groupby(["curriculum", "arm"])["seed"].nunique().to_string(), "\n")

    print("## Final-generation best-lineage accuracy (all six families), mean over seeds")
    fin = best[best["generation"] == best["generation"].max()]
    print(fin.groupby(["curriculum", "arm"])["value"].agg(["mean", "count"]).round(3).to_string(), "\n")

    print("## Pre-registered contrasts (per seed; mean ± 95% CI over seeds)")
    pairs = [("veto", "society_stop3"), ("society_stop3", "isolated"), ("society_stop3", "society"),
             ("veto", "isolated"), ("decor_veto", "decor_isolated"),
             ("early_veto", "early_isolated"), ("late_veto", "late_isolated"),
             ("early_society", "early_isolated"), ("late_society", "late_isolated"),
             ("cull_veto", "cull_isolated"), ("cull_veto", "veto"), ("cull_isolated", "isolated")]
    print(final_contrasts(best, pairs).to_string(index=False), "\n")

    tab = decline_table(df)
    if len(tab):
        print("## Fraction of merges declined vs partner complementarity")
        print(tab.groupby(["curriculum", "generation"])[["declined", "complementarity"]]
              .mean().round(2).to_string(), "\n")
        res = decline_test(tab)
        print("## Partial-correlation test (pooled over curricula; controls: generation)")
        for k, v in res.items():
            print(f"  {k}: {np.round(v, 3) if not isinstance(v, tuple) else tuple(round(t, 3) for t in v)}")
        timing = tab[tab["curriculum"].isin(["early", "late"])]
        if timing["curriculum"].nunique() == 2:
            print("\n## Conflict-timing test (early + late curricula; controls: generation)")
            for k, v in conflict_timing_test(timing).items():
                print(f"  {k}: {np.round(v, 3) if not isinstance(v, tuple) else tuple(round(t, 3) for t in v)}")
            print("## Same test pooled over all four curricula")
            for k, v in conflict_timing_test(tab).items():
                print(f"  {k}: {np.round(v, 3) if not isinstance(v, tuple) else tuple(round(t, 3) for t in v)}")
        if res["n_curricula"] < 2:
            print("  (one curriculum only: complementarity and generation are collinear; the partial"
                  " correlation is not interpretable until the decorrelated bundle exists)")


if __name__ == "__main__":
    main()