1"""Primary analysis: task-weighted completion by arm, prespecified contrasts (D−A primary; B−A, C−B, D−C secondary with Holm),2task-cluster bootstrap CIs and paired sign-flip permutation p-values. Also secondary endpoints by arm."""3import os, sys, json4sys.path.insert(0, os.path.dirname(__file__))5from common import *67def arm_table(df):8 rows = []9 for a in ARMS:10 d = df[df["treatment"] == a]11 rows.append({"arm": a, "n_runs": int(len(d)), "n_tasks": int(d["task_id"].nunique()), "n_evaluated": int((d["status"] == "evaluated").sum()),12 "Y_task_weighted": task_weighted_rate(d, a, "Y"), "Y_run_mean": float(d["Y"].mean()) if len(d) else np.nan,13 "correct_any_task_weighted": task_weighted_rate(d, a, "Y_any"), "violation_rate": float(d["viol"].mean()) if len(d) else np.nan,14 "n_runs_with_violation": int(d["viol"].sum()), "reproducible_rate": float(d["reproducibility"].mean()) if len(d) else np.nan,15 "quality_mean": float(d["quality_score"].mean()) if len(d) else np.nan, "q_correctness": float(d["q_correctness"].mean()) if len(d) else np.nan,16 "q_reporting": float(d["q_reporting"].mean()) if len(d) else np.nan, "q_baseline": float(d["q_baseline"].mean()) if len(d) else np.nan,17 "cost_mean_usd": float(d["spent_usd"].mean()) if len(d) else np.nan, "cost_median_usd": float(d["spent_usd"].median()) if len(d) else np.nan,18 "cost_p90_usd": float(d["spent_usd"].quantile(0.9)) if len(d) else np.nan, "cost_total_usd": float(d["spent_usd"].sum()),19 "reviewer_usd_total": float(d["reviewer_usd"].fillna(0).sum()), "wall_median_s": float(d["spent_wall_s"].median()) if len(d) else np.nan,20 "wall_p90_s": float(d["spent_wall_s"].quantile(0.9)) if len(d) else np.nan, "n_timeout": int((d["terminal_reason"] == "wall_clock_limit").sum()),21 "n_budget_exhausted": int(d["terminal_reason"].fillna("").str.contains("budget").sum()), "n_sessions_mean": float(d["n_sessions"].mean()) if len(d) else np.nan,22 "cost_per_success_usd": (float(d["spent_usd"].sum() / d["Y"].sum()) if d["Y"].sum() > 0 else None)})23 return pd.DataFrame(rows)2425def contrasts(df, col="Y", B=5000, n_perm=20000):26 out = []27 pv = {}28 for (a, b, role) in [("D", "A", "primary"), ("B", "A", "secondary"), ("C", "B", "secondary"), ("D", "C", "secondary"), ("C", "A", "exploratory"), ("D", "B", "exploratory")]:29 d = paired_task_diff(df, a, b, col); est, lo, hi = cluster_bootstrap(d, B); p = permutation_p(d, n_perm)30 out.append({"contrast": f"{a}-{b}", "role": role, "endpoint": col, "n_task_clusters": int(len(d)), "estimate": est, "ci_low": lo, "ci_high": hi, "perm_p": p,31 "n_tasks_favoring_first": int((d > 0).sum()), "n_tasks_favoring_second": int((d < 0).sum()), "n_tasks_tied": int((d == 0).sum())})32 if role == "secondary": pv[f"{a}-{b}"] = p33 adj = holm(pv) if pv else {}34 for r in out:35 r["holm_p"] = adj.get(r["contrast"])36 return pd.DataFrame(out)3738def main(split="confirmatory"):39 os.makedirs(EST, exist_ok=True)40 df = load_runs(split)41 if df.empty:42 print(f"no runs for split {split}"); return43 at = arm_table(df); at.to_csv(os.path.join(EST, f"{split}_arm_table.csv"), index=False)44 ct = pd.concat([contrasts(df, "Y"), contrasts(df, "Y_any"), contrasts(df, "viol"), contrasts(df, "quality_score")]); ct.to_csv(os.path.join(EST, f"{split}_contrasts.csv"), index=False)45 fam = df.groupby(["family", "treatment"]).apply(lambda d: pd.Series({"Y": d.groupby("task_id")["Y"].mean().mean(), "n_tasks": d["task_id"].nunique(), "n_runs": len(d), "viol": d["viol"].mean(), "cost_mean": d["spent_usd"].mean()}), include_groups=False).reset_index()46 fam.to_csv(os.path.join(EST, f"{split}_family_table.csv"), index=False)47 print(at.round(3).to_string()); print(); print(ct.round(4).to_string())4849if __name__ == "__main__":50 main(sys.argv[1] if len(sys.argv) > 1 else "confirmatory")51