#!/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()