SPB Git

spb/wp42_uqo

Public
14.8 KB · 210 lines python
1"""Agent G — preregistered estimation from data/validated/runs_<tag>.parquet.2Outputs: analysis/tables/*.csv (+ .tex) and analysis/derived/estimates_<tag>.json. All numbers in the manuscript come from here.3Usage: .venv/bin/python analysis/econometrics/estimate.py --tag pilot [--q 70] [--B 5000] [--perm 20000]"""4from __future__ import annotations5import argparse, json, os, warnings6import numpy as np, pandas as pd7warnings.filterwarnings("ignore")8REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))9TAB = os.path.join(REPO, "analysis", "tables"); DER = os.path.join(REPO, "analysis", "derived")10ARMS = ["A0", "A1", "A2", "A3", "A4"]11CONTRASTS = [("A3", "A1", "primary"), ("A1", "A0", "secondary"), ("A4", "A3", "secondary"), ("A2", "A1", "secondary")]1213def load(tag, q):14    df = pd.read_parquet(os.path.join(REPO, "data", "validated", f"runs_{tag}.parquet"))15    df = df[df["status"].isin(["collected", "evaluated"])].copy()   # metrics are computed before the run is marked evaluated16    df["S"] = ((df["det_completion"].astype(bool)) & (~df["critical_error"].astype(bool)) & (df["quality"].astype(float) >= q)).astype(float)17    df.loc[df["quality"].isna(), "S"] = np.nan18    df["D"] = df["det_completion"].astype(float); df["E"] = df["critical_error"].astype(float)19    df["C"] = df["total_cost_usd"].astype(float); df["T"] = df["wall_total_s"].astype(float)20    df["censored"] = df["timed_out"].astype(bool)21    # Killed sessions emit no `result` event, so their marginal cost is unobserved (cost 0 recorded). Impute it from the session's22    # token counts with per-token prices fitted by least squares on the completed runs of the same model (pricing is linear).23    tok = df["tokens"].apply(lambda x: json.loads(x) if isinstance(x, str) else (x or {}))24    X = pd.DataFrame({"input": tok.apply(lambda t: t.get("input", 0)), "output": tok.apply(lambda t: t.get("output", 0)),25                      "cache_read": tok.apply(lambda t: t.get("cache_read", 0)), "cache_create": tok.apply(lambda t: t.get("cache_create", 0))}, index=df.index).astype(float)26    df["cost_observed"] = df["C"]; df["cost_imputed_flag"] = False27    # prices are fitted on the completed pilot runs of the same model (the recovery stage has no completed first session)28    try:29        ref = pd.read_parquet(os.path.join(REPO, "data", "validated", "runs_pilot.parquet")); ref = ref[ref["cost_usd"].fillna(0) > 0]30    except Exception: ref = df[df["cost_usd"].fillna(0) > 0]  # noqa: BLE00131    rtok = ref["tokens"].apply(lambda x: json.loads(x) if isinstance(x, str) else (x or {}))32    RX = pd.DataFrame({"input": rtok.apply(lambda t: t.get("input", 0)), "output": rtok.apply(lambda t: t.get("output", 0)), "cache_read": rtok.apply(lambda t: t.get("cache_read", 0)), "cache_create": rtok.apply(lambda t: t.get("cache_create", 0))}, index=ref.index).astype(float)33    for mdl, g in df.groupby("model"):34        fit = ref[ref["model"] == mdl]35        miss = g[(g["cost_usd"].fillna(0) == 0) & g["censored"]]36        if len(fit) >= 5 and len(miss):37            beta, *_ = np.linalg.lstsq(RX.loc[fit.index].values, fit["cost_usd"].values, rcond=None)38            imput = X.loc[miss.index].values @ beta39            df.loc[miss.index, "C"] = np.maximum(0, imput) + df.loc[miss.index, "recovery_cost_usd"].fillna(0).astype(float)40            df.loc[miss.index, "cost_imputed_flag"] = True41    return df4243def task_means(df, col, arm):44    return df[df.arm == arm].groupby("task_id")[col].mean()4546def weighted_rate(df, col, arm, weights=None):47    tm = task_means(df, col, arm).dropna()48    if weights is None: return float(tm.mean())49    w = weights.reindex(tm.index).fillna(0); return float((tm * w).sum() / w.sum())5051def cluster_bootstrap(df, col, a, b, B, rng, weights=None):52    ta = task_means(df, col, a); tb = task_means(df, col, b); tasks = sorted(set(ta.index) & set(tb.index))53    d = np.array([ta[t] - tb[t] for t in tasks]); w = np.array([weights[t] if weights is not None else 1.0 for t in tasks])54    est = float((d * w).sum() / w.sum()); boots = []55    for _ in range(B):56        idx = rng.integers(0, len(tasks), len(tasks)); boots.append((d[idx] * w[idx]).sum() / w[idx].sum())57    lo, hi = np.percentile(boots, [2.5, 97.5]); return est, float(lo), float(hi), len(tasks)5859def permutation_p(df, col, a, b, n_perm, rng):60    ta = task_means(df, col, a); tb = task_means(df, col, b); tasks = sorted(set(ta.index) & set(tb.index))61    d = np.array([ta[t] - tb[t] for t in tasks]); obs = d.mean(); cnt = 062    for _ in range(n_perm):63        s = rng.choice([-1, 1], len(d)); cnt += abs((d * s).mean()) >= abs(obs) - 1e-1264    return float(cnt / n_perm)6566def holm(pvals):67    idx = np.argsort(pvals); m = len(pvals); adj = np.empty(m); run = 068    for r, i in enumerate(idx): run = max(run, (m - r) * pvals[i]); adj[i] = min(1.0, run)69    return adj.tolist()7071def mixed_logit(df, col):72    try:73        import statsmodels.formula.api as smf74        d = df.dropna(subset=[col]).copy(); d["y"] = d[col].astype(int)75        md = smf.mixedlm if False else None76        import statsmodels.api as sm77        X = pd.get_dummies(d["arm"], drop_first=False).astype(float); X = X[[a for a in ARMS if a in X.columns]]78        m = sm.GLM(d["y"], X, family=sm.families.Binomial()).fit(cov_type="cluster", cov_kwds={"groups": pd.factorize(d["task_id"])[0]})79        out = {a: {"logit": float(m.params[a]), "se_cluster": float(m.bse[a]), "p": float(1 / (1 + np.exp(-m.params[a])))} for a in X.columns}80        return {"model": "cluster-robust logit (arm dummies, no intercept), clustered by scenario", "params": out}81    except Exception as e:  # noqa: BLE00182        return {"error": repr(e)}8384def survival(df):85    try:86        from lifelines import KaplanMeierFitter87        out = {}88        for a in ARMS:89            d = df[df.arm == a]90            if len(d) == 0: continue91            km = KaplanMeierFitter(); km.fit(d["T"], event_observed=~d["censored"])92            out[a] = {"median_s": float(km.median_survival_time_), "n": int(len(d)), "censored": int(d["censored"].sum()),93                      "p_done_600": float(1 - km.survival_function_at_times(600).iloc[0]), "p_done_1200": float(1 - km.survival_function_at_times(1200).iloc[0])}94        return out95    except Exception as e:  # noqa: BLE00196        return {"error": repr(e)}9798def alf_curves(df, budgets, weights=None, h_allow=0.0, q_col="S"):99    """ALF_o(B,H,q) = Σ_j w_oj P(S=1, C≤B, H≤H). H = 5 min per unanswered question (assumption) + logged interventions."""100    rows = []101    for occ, g in df.groupby("occupation"):102        for a in ARMS:103            ga = g[g.arm == a]104            if len(ga) == 0: continue105            for B in budgets:106                ok = (ga[q_col] == 1) & (ga["C"] <= B) & (ga["H_min"] <= h_allow)107                tm = ok.groupby(ga["task_id"]).mean()108                w = (weights.reindex(tm.index).fillna(0) if weights is not None else pd.Series(1.0, index=tm.index))109                rows.append({"occupation": occ, "arm": a, "B": B, "H": h_allow, "ALF": float((tm * w).sum() / w.sum())})110    return pd.DataFrame(rows)111112def agreement(jdg):113    """Krippendorff's alpha (interval) across judges on the per-run quality scores + pairwise Pearson."""114    piv = jdg.pivot_table(index="run_id", columns="judge_idx", values="score")115    try:116        vals = piv.values.astype(float); n_units = vals.shape[0]117        obs = []; allv = vals[~np.isnan(vals)]118        De = np.var(allv) * 2 if len(allv) > 1 else np.nan119        Do = 0.0; m = 0120        for row in vals:121            r = row[~np.isnan(row)]122            if len(r) > 1:123                Do += np.sum((r[:, None] - r[None, :]) ** 2) / (len(r) - 1); m += len(r)124        Do = Do / m if m else np.nan125        alpha = 1 - Do / De if De and not np.isnan(De) else np.nan126    except Exception: alpha = np.nan  # noqa: BLE001127    corr = piv.corr().round(3).to_dict() if piv.shape[1] > 1 else {}128    return {"krippendorff_alpha_interval": float(alpha) if alpha == alpha else None, "pairwise_pearson": corr, "n_runs": int(piv.shape[0])}129130def main():131    ap = argparse.ArgumentParser(); ap.add_argument("--tag", default="pilot"); ap.add_argument("--q", type=float, default=70); ap.add_argument("--B", type=int, default=5000); ap.add_argument("--perm", type=int, default=20000); ap.add_argument("--seed", type=int, default=42)132    a = ap.parse_args(); rng = np.random.default_rng(a.seed); os.makedirs(TAB, exist_ok=True); os.makedirs(DER, exist_ok=True)133    df = load(a.tag, a.q)134    # human-time proxy: 5 min per question in COMPLETION.json (assumption) + logged interventions135    def nq(x):136        try: return len(json.loads(x)) if isinstance(x, str) else 0137        except Exception: return 0  # noqa: BLE001138    df["H_min"] = df["human_intervention_min"].astype(float) + 5.0 * df.get("n_questions", pd.Series(0, index=df.index)).fillna(0).astype(float) if "n_questions" in df else df["human_intervention_min"].astype(float)139    # task weights: equal (primary) and professional minutes (sensitivity)140    tasks = {t: json.load(open(os.path.join(REPO, "benchmark", "scenarios", t, "task.json"))) for t in df.task_id.unique()}141    w_min = pd.Series({t: float(v.get("professional_minutes_estimate", 60)) for t, v in tasks.items()})142    # --- Table 1: outcomes by arm143    rows = []144    for arm in ARMS:145        d = df[df.arm == arm]146        if len(d) == 0: continue147        S = d["S"]; acc = S.sum(); rows.append({"arm": arm, "n_runs": len(d), "n_tasks": d.task_id.nunique(), "acceptance_taskw": weighted_rate(df, "S", arm), "acceptance_minw": weighted_rate(df, "S", arm, w_min),148            "det_completion": weighted_rate(df, "D", arm), "critical_error_rate": float(d["E"].mean()), "quality_mean": float(d["quality"].mean()), "quality_sd": float(d["quality"].std()),149            "timeout_rate": float(d["censored"].mean()), "cost_mean": float(d["C"].mean()), "cost_sd": float(d["C"].std()), "cost_per_accepted": float(d["C"].sum() / acc) if acc else np.nan,150            "wall_median_s": float(d["T"].median()), "tool_calls_mean": float(d["tool_calls"].mean()), "tool_errors_mean": float(d["tool_errors"].mean()), "app_transitions_mean": float(d["app_transitions"].mean()),151            "abstention_rate": float((d["abstention_status"].isin(["needs_clarification", "abstained"])).mean())})152    t1 = pd.DataFrame(rows); t1.to_csv(os.path.join(TAB, f"table1_outcomes_by_arm_{a.tag}.csv"), index=False)153    # --- Table 2: contrasts154    crows = []; pv = []155    for x, y, kind in CONTRASTS:156        for col in ("S", "D", "E", "quality"):157            if col == "quality":158                est, lo, hi, J = cluster_bootstrap(df, "quality", x, y, a.B, rng)159                p = permutation_p(df, "quality", x, y, a.perm, rng)160            else:161                est, lo, hi, J = cluster_bootstrap(df, col, x, y, a.B, rng); p = permutation_p(df, col, x, y, a.perm, rng)162            est_w = cluster_bootstrap(df, col, x, y, 200, rng, w_min)[0]163            crows.append({"contrast": f"{x}-{y}", "kind": kind, "outcome": col, "estimate": est, "ci_lo": lo, "ci_hi": hi, "perm_p": p, "J": J, "estimate_minutes_weighted": est_w})164            if col == "S" and kind == "secondary": pv.append(p)165    t2 = pd.DataFrame(crows)166    if pv:167        adj = holm(pv); k = 0168        t2["holm_p"] = np.nan169        for i, r in t2.iterrows():170            if r["outcome"] == "S" and r["kind"] == "secondary": t2.loc[i, "holm_p"] = adj[k]; k += 1171    t2.to_csv(os.path.join(TAB, f"table2_contrasts_{a.tag}.csv"), index=False)172    # --- Table 3: by occupation × arm173    t3 = df.groupby(["occupation", "arm"]).agg(n=("S", "size"), acceptance=("S", "mean"), det=("D", "mean"), quality=("quality", "mean"), critical=("E", "mean"), cost=("C", "mean"), tool_calls=("tool_calls", "mean")).reset_index()174    t3.to_csv(os.path.join(TAB, f"table3_occupation_arm_{a.tag}.csv"), index=False)175    # --- dependency: acceptance vs designed depth and vs tool calls (within arm)176    df["depth"] = df.task_id.map({t: v.get("designed_dependency_depth") for t, v in tasks.items()})177    dep = df.groupby(["arm", "depth"]).agg(n=("S", "size"), acceptance=("S", "mean"), tool_calls=("tool_calls", "mean")).reset_index(); dep.to_csv(os.path.join(TAB, f"table4_dependency_{a.tag}.csv"), index=False)178    # --- recovery index (A4)179    a4 = df[df.arm == "A4"]; fail_first = a4[a4["selfcheck_pass_first"] == 0]180    ari = {"n_A4": int(len(a4)), "n_failed_first_selfcheck": int(len(fail_first)), "n_recovery_attempted": int(a4["recovery_attempted"].astype(bool).sum()),181           "n_recovered": int((fail_first["selfcheck_pass_final"] == 1).sum()), "ARI": float((fail_first["selfcheck_pass_final"] == 1).mean()) if len(fail_first) else None,182           "acceptance_given_recovery": float(fail_first["S"].mean()) if len(fail_first) else None, "mean_recovery_cost": float(a4["recovery_cost_usd"].fillna(0).mean())}183    # --- ALF curves184    budgets = [0.25, 0.5, 0.75, 1, 1.5, 2, 3, 4]185    alf0 = alf_curves(df, budgets, None, 0.0); alf0["weights"] = "equal"; alf1 = alf_curves(df, budgets, w_min, 0.0); alf1["weights"] = "minutes"186    alfH = alf_curves(df, budgets, None, 15.0); alfH["weights"] = "equal_H15"187    alf = pd.concat([alf0, alf1, alfH]); alf.to_csv(os.path.join(TAB, f"table5_alf_{a.tag}.csv"), index=False)188    # --- sensitivity to q189    qs = []190    for q in (50, 60, 70, 80, 90):191        d2 = df.copy(); d2["S"] = ((d2["D"] == 1) & (d2["E"] == 0) & (d2["quality"] >= q)).astype(float)192        for arm in ARMS:193            if (d2.arm == arm).any(): qs.append({"q": q, "arm": arm, "acceptance": weighted_rate(d2, "S", arm)})194    pd.DataFrame(qs).to_csv(os.path.join(TAB, f"table6_q_sensitivity_{a.tag}.csv"), index=False)195    # --- judges agreement196    jdg = pd.read_parquet(os.path.join(REPO, "data", "validated", f"judgments_{a.tag}.parquet"))197    agr = agreement(jdg)198    jm = jdg.groupby("judge_model")["score"].agg(["mean", "std", "count"]).reset_index(); jm.to_csv(os.path.join(TAB, f"table7_judges_{a.tag}.csv"), index=False)199    # det vs judge-acceptable agreement200    df["judge_majority_acceptable"] = df["judge_acceptable_votes"].apply(lambda s: (sum(1 for v in json.loads(s) if v) >= 2) if isinstance(s, str) else None)201    cross = pd.crosstab(df["D"], df["judge_majority_acceptable"]); cross.to_csv(os.path.join(TAB, f"table8_det_vs_judges_{a.tag}.csv"))202    out = {"tag": a.tag, "q": a.q, "n_runs": int(len(df)), "n_tasks": int(df.task_id.nunique()), "total_cost_usd": float(df["C"].sum()), "by_arm": t1.to_dict("records"), "contrasts": t2.to_dict("records"),203           "mixed_logit_S": mixed_logit(df, "S"), "survival": survival(df), "recovery": ari, "agreement": agr, "models": sorted(df["model"].dropna().unique().tolist()),204           "claude_code_versions": sorted(df["claude_code_version"].dropna().unique().tolist())}205    json.dump(out, open(os.path.join(DER, f"estimates_{a.tag}.json"), "w"), indent=1, default=str)206    print(t1.round(3).to_string()); print(t2.round(3).to_string()); print(json.dumps(ari, indent=1)); print(json.dumps(agr, indent=1, default=str))207208if __name__ == "__main__":209    main()210
Files · e5075d96fc9d1a321f82c0c9974e505eee466003
Menu

Theme