scripts/analysis/run.py
Served verbatim from the project repository. Internal working document conventions apply: documents may reference file paths, branch names, and findings-ledger anchors from the repo.
#!/usr/bin/env python
"""Driver: run all 10 approaches x 3 outcomes through the validity battery."""
import numpy as np
import pandas as pd
import bakeoff as B
pd.set_option("display.width", 160)
pd.set_option("display.max_columns", 30)
def eta2(labels, y):
"""Between-group variance fraction (one-way ANOVA R^2)."""
d = pd.DataFrame({"g": labels, "y": y}).dropna()
if d["g"].nunique() < 2 or len(d) < 10:
return np.nan
grand = d["y"].mean()
ss_tot = ((d["y"] - grand) ** 2).sum()
ss_b = d.groupby("g")["y"].apply(lambda s: len(s) * (s.mean() - grand) ** 2).sum()
return ss_b / ss_tot if ss_tot > 0 else np.nan
def within_sd(labels, y):
d = pd.DataFrame({"g": labels, "y": y}).dropna()
if len(d) < 10: return np.nan
return d.groupby("g")["y"].std().mean()
# ------- method runners: return (expected Series, labels Series or None) per year df
def run_method(name, df, ycol):
"""Return dict: expected (Series), labels (Series|None)."""
if name == "01_knn_demo":
e, l = B.expected_knn(df, ycol, B.DEMOG); return e, None
if name == "05_knn_prior":
e, l = B.expected_knn(df, ycol, B.DEMOG + ["prior"]); return e, None
if name == "06_need_strata":
# need-index decile within band = comparison stratum (partition)
lab = pd.Series("", index=df.index)
for band, g in df.groupby("grade_band"):
dec = B.ntile(B.index_meanpctrank(g), 10)
lab.loc[g.index] = [f"{band}:{d}" for d in dec]
e = pd.Series(B.loo_group_mean(lab.values, df[ycol].values), index=df.index)
return e, lab
if name in B.PARTITION_METHODS:
e, l = B.expected_partition(df, B.PARTITION_METHODS[name], ycol); return e, l
if name in B.INDEX_METHODS:
e, dec, idx = B.expected_index(df, B.INDEX_METHODS[name], ycol); return e, dec
if name == "10_reg_intake":
return B.expected_regression(df, ycol, B.DEMOG), None
if name == "11_reg_value_added":
return B.expected_regression(df, ycol, B.DEMOG + ["prior"]), None
raise ValueError(name)
def oos_method(name, train, test, ycol):
if name == "01_knn_demo":
return B.oos_knn(train, test, ycol, B.DEMOG)
if name == "05_knn_prior":
return B.oos_knn(train, test, ycol, B.DEMOG + ["prior"])
if name == "06_need_strata":
return B.oos_index(train, test, B.index_meanpctrank, ycol)
if name in B.PARTITION_METHODS:
return B.oos_partition(train, test, B.PARTITION_METHODS[name], ycol)
if name in B.INDEX_METHODS:
return B.oos_index(train, test, B.INDEX_METHODS[name], ycol)
if name == "10_reg_intake":
return B.oos_regression(train, test, ycol, B.DEMOG)
if name == "11_reg_value_added":
return B.oos_regression(train, test, ycol, B.DEMOG + ["prior"])
raise ValueError(name)
METHODS = ["01_knn_demo", "02_kmeans", "03_gmm", "04_ward", "05_knn_prior",
"06_need_strata", "07_meanpctrank", "07b_meanpctrank7",
"08_pca", "08b_pca3", "09_eni",
"10_reg_intake", "11_reg_value_added"]
NEEDS_PRIOR = {"05_knn_prior", "11_reg_value_added"}
def main():
panel = B.load_panel()
ins_rows, oos_rows, stab_rows = [], [], []
for oc, (ycol, hib, bands) in B.OUTCOMES.items():
# universe eligible for this outcome
uni = panel[panel["grade_band"].isin(bands)].copy()
years = sorted(uni.loc[uni[ycol].notna(), "year"].unique())
# prior-year outcome feature (per outcome) for the whole eligible universe
uni["prior"] = B.build_prior(uni, ycol)
print(f"\n=== OUTCOME {oc} ({ycol}) bands={sorted(bands)} years={years} ===")
# per-year in-sample + store residuals for stability
resid_store = {m: {} for m in METHODS} # m -> {year -> Series(resid by dbn)}
for yr in years:
dfy = uni[uni["year"] == yr].copy()
n = dfy[ycol].notna().sum()
if n < 30:
continue
# COMMON-SUPPORT mask: schools usable by EVERY method incl. prior-needing
# ones, so 01 vs 05 / 10 vs 11 are scored on the identical, equally-hard
# subset (fixes the "prior methods get an easier subset" confound).
common = (dfy[ycol].notna() & dfy["prior"].notna()).values
n_common = int(common.sum())
for m in METHODS:
if m in NEEDS_PRIOR and dfy["prior"].notna().sum() < 30:
continue
d = dfy if m not in NEEDS_PRIOR else dfy[dfy["prior"].notna()].copy()
exp, lab = run_method(m, d, ycol)
val = B.r2(d[ycol].values, exp.values)
# fair R^2 on the common subset (reindex expected back onto dfy)
exp_full = exp.reindex(dfy.index)
val_common = B.r2(dfy[ycol].values[common], exp_full.values[common]) \
if n_common >= 30 else np.nan
e2 = eta2(lab.values, d[ycol].values) if lab is not None else np.nan
wsd = within_sd(lab.values, d[ycol].values) if lab is not None else np.nan
ins_rows.append(dict(outcome=oc, method=m, year=yr, n=int(n),
n_common=n_common, r2_loo=val,
r2_loo_common=val_common, eta2=e2, within_sd=wsd))
resid = pd.Series(d[ycol].values - exp.values, index=d["dbn"].values)
resid_store[m][yr] = resid.dropna()
# residual stability: corr of residual across consecutive years
for m in METHODS:
ys = sorted(resid_store[m].keys())
for i in range(1, len(ys)):
a, b = resid_store[m][ys[i - 1]], resid_store[m][ys[i]]
common = a.index.intersection(b.index)
if len(common) < 30:
continue
c = np.corrcoef(a.loc[common], b.loc[common])[0, 1]
stab_rows.append(dict(outcome=oc, method=m,
year_pair=f"{ys[i-1]}->{ys[i]}", n=len(common),
resid_corr=c))
# out-of-sample: fit year T, predict T+1
for i in range(1, len(years)):
ytr, yte = years[i - 1], years[i]
tr = uni[uni["year"] == ytr].copy()
te = uni[uni["year"] == yte].copy()
if tr[ycol].notna().sum() < 50 or te[ycol].notna().sum() < 50:
continue
for m in METHODS:
if m in NEEDS_PRIOR and (tr["prior"].notna().sum() < 50 or te["prior"].notna().sum() < 50):
continue
trp = tr if m not in NEEDS_PRIOR else tr[tr["prior"].notna()].copy()
tep = te if m not in NEEDS_PRIOR else te[te["prior"].notna()].copy()
exp = oos_method(m, trp, tep, ycol)
val = B.r2(tep[ycol].values, exp.reindex(tep.index).values)
oos_rows.append(dict(outcome=oc, method=m,
year_pair=f"{ytr}->{yte}", n=int(tep[ycol].notna().sum()),
r2_oos=val))
ins = pd.DataFrame(ins_rows); oos = pd.DataFrame(oos_rows); stab = pd.DataFrame(stab_rows)
ins.to_csv(f"{B.DATA}/results_insample.csv", index=False)
oos.to_csv(f"{B.DATA}/results_oos.csv", index=False)
stab.to_csv(f"{B.DATA}/results_stability.csv", index=False)
# leaderboard: mean across years per (outcome, method)
lb = (ins.groupby(["outcome", "method"])
.agg(r2_loo=("r2_loo", "mean"), r2_common=("r2_loo_common", "mean"),
eta2=("eta2", "mean"), within_sd=("within_sd", "mean")).reset_index())
lb = lb.merge(oos.groupby(["outcome", "method"])["r2_oos"].mean().reset_index(), how="left")
lb = lb.merge(stab.groupby(["outcome", "method"])["resid_corr"].mean().reset_index(), how="left")
lb.to_csv(f"{B.DATA}/leaderboard.csv", index=False)
print("\n================ LEADERBOARD (mean across years) ================")
for oc in B.OUTCOMES:
sub = lb[lb["outcome"] == oc].sort_values("r2_common", ascending=False)
print(f"\n--- {oc} (sorted by r2_common = fair common-subset R^2) ---")
print(sub.to_string(index=False,
formatters={c: (lambda v: f"{v:.3f}" if pd.notna(v) else " -")
for c in ["r2_loo", "r2_common", "eta2", "within_sd", "r2_oos", "resid_corr"]}))
print("\nWrote results_insample.csv, results_oos.csv, results_stability.csv, leaderboard.csv")
if __name__ == "__main__":
main()