1"""F2-conf-01: school-level programme effect with cluster-robust inference."""2import json, os3import numpy as np, pandas as pd4import statsmodels.formula.api as smf5from scipy import stats67ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))8d = pd.read_csv(os.path.join(ROOT, "data", "schools.csv"))9n_rows = len(d)10issues = []11# --- data checks ---12if d.isna().any().any():13 issues.append({"issue": "missing values", "handling": "rows with missing values dropped"})14 d = d.dropna()15if d.student_id.duplicated().any():16 issues.append({"issue": "duplicate student_id", "handling": "duplicates dropped (first kept)"})17 d = d.drop_duplicates("student_id")18bad = d.groupby("school_id").treated.nunique()19if (bad > 1).any():20 issues.append({"issue": "treatment varies within school", "handling": "schools excluded"})21 d = d[~d.school_id.isin(bad[bad > 1].index)]22sizes = d.groupby("school_id").size()23assert sizes.nunique() == 1, "unbalanced"24G = d.school_id.nunique(); m = int(sizes.iloc[0]); N = len(d)25G1 = int(d.groupby("school_id").treated.first().sum())2627# --- OLS, iid SE ---28ols = smf.ols("score ~ treated", data=d).fit()29tau = float(ols.params["treated"]); se_iid = float(ols.bse["treated"])30# --- cluster-robust (CR1, Stata-style small-sample correction), t with G-1 df ---31cl = smf.ols("score ~ treated", data=d).fit(cov_type="cluster",32 cov_kwds={"groups": d.school_id.values, "use_correction": True})33se_cl = float(cl.bse["treated"])34df = G - 135tstat = tau / se_cl36p = float(2 * stats.t.sf(abs(tstat), df))37q = stats.t.ppf(0.975, df)38lo, hi = tau - q * se_cl, tau + q * se_cl3940# cross-check: two-sample t-test on school means (exact analogue under balance)41sm = d.groupby("school_id").agg(y=("score", "mean"), t=("treated", "first"))42tt = stats.ttest_ind(sm.y[sm.t == 1], sm.y[sm.t == 0], equal_var=True)43se_means = float(tau / tt.statistic)4445# --- ICC of residual: one-way ANOVA estimator on OLS residuals ---46d = d.assign(r=ols.resid)47rm = d.groupby("school_id").r.transform("mean")48ssw = ((d.r - rm) ** 2).sum(); msw = ssw / (N - G)49ssb = m * (d.groupby("school_id").r.mean() ** 2).sum(); msb = ssb / (G - 2) # 2 params removed50s2b = max((msb - msw) / m, 0.0)51icc = float(s2b / (s2b + msw))52deff = 1 + (m - 1) * icc5354res = {55 "task_id": "F2-conf-01", "tau_hat": tau, "se_iid": se_iid, "se_cluster": se_cl,56 "p_value": p, "tau_ci_low": float(lo), "tau_ci_high": float(hi),57 "effect_significant_5pct": bool(p < 0.05), "icc_hat": icc,58 "n_schools": int(G), "n_treated_schools": G1, "n_analysed": int(N),59 "feasibility": "complete", "feasibility_reason": None, "data_issues": issues,60 "n_input_rows": {"schools.csv": int(n_rows)}, "random_seed": None,61 "_diagnostics": {"t_stat_cluster": float(tstat), "df": df, "students_per_school": m,62 "se_school_means_ttest": se_means, "p_school_means_ttest": float(tt.pvalue),63 "p_iid": float(ols.pvalues["treated"]), "design_effect_from_icc": deff,64 "se_ratio_cluster_over_iid": se_cl / se_iid},65}66os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)67with open(os.path.join(ROOT, "out", "results.json"), "w") as f:68 json.dump(res, f, indent=2)69print(json.dumps(res, indent=2))70