1"""T2-simulation: coverage of CR0, CR1 and wild-cluster-bootstrap-t intervals for a mean.2Run from workspace root: ~/wp29/venv/bin/python out/analysis.py3Regenerates out/results.json, out/table1.csv, out/figure1.png from data/design.csv only.4"""5import json, os6import numpy as np7import pandas as pd8from scipy import stats9import matplotlib10matplotlib.use("Agg")11import matplotlib.pyplot as plt1213DATA = "data/design.csv"14OUT = "out"15design = pd.read_csv(DATA)16data_issues = []1718def key(est, G, icc):19 return f"{est}_G{int(G)}_icc{float(icc):.1f}"2021def simulate_cell(row):22 G, icc, R, B = int(row.G), float(row.icc), int(row.R), int(row.B)23 mu, level = float(row.mu), float(row.nominal_level)24 n = np.array([int(s) for s in str(row.cluster_sizes).split("|")])25 N = int(n.sum())26 assert len(n) == G, "cluster_sizes length != G"27 assert N == int(row.N), "sum(cluster_sizes) != N"28 var_exact = (icc * np.sum(n**2) + (1 - icc) * N) / N**229 rng = np.random.default_rng(int(row.seed))30 # draw order (fixed, deterministic): U (R,G), E (R,N), bootstrap weights (R,B,G)31 U = rng.normal(0.0, np.sqrt(icc), size=(R, G)) if icc > 0 else np.zeros((R, G))32 E = rng.normal(0.0, np.sqrt(1 - icc), size=(R, N))33 W = rng.choice(np.array([-1.0, 1.0]), size=(R, B, G))34 cl = np.repeat(np.arange(G), n)35 Y = mu + U[:, cl] + E # (R,N)36 ybar = Y.mean(axis=1) # (R,)37 ehat = Y - ybar[:, None]38 S = np.zeros((R, G))39 for g in range(G):40 S[:, g] = ehat[:, cl == g].sum(axis=1) # cluster residual sums41 V0 = (S**2).sum(axis=1) / N**242 V1 = G / (G - 1) * V043 z = 1.96 # as specified in TASK.md44 tcrit = stats.t.ppf(1 - (1 - level) / 2, G - 1)45 hw0 = z * np.sqrt(V0)46 hw1 = tcrit * np.sqrt(V1)47 # Wild cluster bootstrap-t (unrestricted, Rademacher), using cluster sums:48 # ybar* - ybar = D = sum_g w_g S_g / N ; bootstrap residual cluster sums = w_g S_g - n_g D49 WS = W * S[:, None, :] # (R,B,G)50 D = WS.sum(axis=2) / N # (R,B)51 Sstar = WS - n[None, None, :] * D[:, :, None]52 V1s = G / (G - 1) * (Sstar**2).sum(axis=2) / N**253 tstar = np.abs(D) / np.sqrt(V1s)54 q = np.quantile(tstar, level, axis=1) # numpy default (linear interpolation) 0.95 quantile55 hwb = q * np.sqrt(V1)56 out = {}57 for est, hw, V in [("cr0", hw0, V0), ("cr1", hw1, V1), ("wcb", hwb, None)]:58 cov = np.mean(np.abs(ybar - mu) <= hw)59 out[est] = dict(coverage=float(cov), mc_se=float(np.sqrt(cov * (1 - cov) / R)),60 rel_bias=None if V is None else float(V.mean() / var_exact),61 rel_bias_mcse=None if V is None else float(V.std(ddof=1) / np.sqrt(R) / var_exact),62 hw=float(hw.mean()))63 # analytic E[V_CR0]/Var(ybar): T_g = cluster sum of (y-mu), a_g = Var(T_g), S_g = T_g - (n_g/N) sum T64 a = icc * n**2 + (1 - icc) * n65 ES2 = a - 2 * (n / N) * a + (n / N) ** 2 * a.sum()66 rb0_exact = float(ES2.sum() / N**2 / var_exact)67 extra = dict(rel_bias_cr0_analytic=rb0_exact, rel_bias_cr1_analytic=rb0_exact * G / (G - 1), var_exact=float(var_exact), emp_var_ybar=float(ybar.var(ddof=1)),68 n_distinct_abs_t_patterns=int(2 ** (G - 1)), mean_q=float(q.mean()), tcrit=float(tcrit))69 return out, extra7071rows, cov, se, rb, rbse, extras = [], {}, {}, {}, {}, {}72for _, row in design.iterrows():73 res, extra = simulate_cell(row)74 extras[row.cell_id] = extra75 for est in ["cr0", "cr1", "wcb"]:76 r = res[est]; k = key(est, row.G, row.icc)77 cov[k] = r["coverage"]; se[k] = r["mc_se"]78 if r["rel_bias"] is not None:79 rb[k] = r["rel_bias"]; rbse[k] = r["rel_bias_mcse"]80 rows.append(dict(estimator=est, cell_id=row.cell_id, G=int(row.G), icc=float(row.icc),81 coverage=r["coverage"], mc_se_coverage=r["mc_se"],82 rel_bias_var=r["rel_bias"], mean_ci_halfwidth=r["hw"], R=int(row.R)))8384if (design.G < 10).any():85 data_issues.append({"issue": "With G=5 and Rademacher weights there are only 2^5=32 distinct weight vectors "86 "(16 distinct |t*| values), so the B=399 bootstrap draws repeat these few values and the "87 "0.95 quantile of |t*| is coarse.",88 "handling": "Kept the specified design (B=399, Rademacher) unchanged; flagged as a limitation in the report."})8990table = pd.DataFrame(rows)91table.to_csv(os.path.join(OUT, "table1.csv"), index=False)9293failure = sorted([k for k, v in cov.items() if v < 0.90])94nominal = sorted([k for k, v in cov.items() if abs(v - 0.95) <= 2 * se[k]])95g5 = {est: float(np.mean([cov[key(est, 5, i)] for i in [0.0, 0.2, 0.5]])) for est in ["cr0", "cr1", "wcb"]}96best = max(g5, key=g5.get)9798results = {99 "task_id": "T2-simulation",100 "R": int(design.R.iloc[0]), "B": int(design.B.iloc[0]), "n_cells": int(len(design)),101 "coverage": cov, "mc_se_coverage": se, "rel_bias_var": rb,102 "failure_cells": failure, "n_failure_cells": len(failure), "nominal_cells": nominal,103 "best_procedure_small_G": best,104 "feasibility": "complete", "feasibility_reason": None,105 "data_issues": data_issues,106 "n_input_rows": {"design.csv": int(len(design))},107 "n_analysed": int(design.R.sum()),108 "random_seed": int(design.loc[design.cell_id == "C1", "seed"].iloc[0]),109 # supplementary diagnostics110 "supplementary": {"rel_bias_var_mc_se": rbse, "mean_coverage_G5": g5, "cell_diagnostics": extras},111}112with open(os.path.join(OUT, "results.json"), "w") as f:113 json.dump(results, f, indent=2)114115fig, axes = plt.subplots(1, 3, figsize=(13, 4.2), sharey=True)116Gs = sorted(design.G.unique())117for ax, icc in zip(axes, [0.0, 0.2, 0.5]):118 for est, m in [("cr0", "o"), ("cr1", "s"), ("wcb", "^")]:119 p = np.array([cov[key(est, G, icc)] for G in Gs]); s = np.array([se[key(est, G, icc)] for G in Gs])120 ax.errorbar(Gs, p, yerr=1.96 * s, marker=m, capsize=3, label=est.upper())121 ax.axhline(0.95, color="k", ls="--", lw=1, label="nominal 0.95")122 ax.axhline(0.90, color="r", ls=":", lw=1, label="failure line 0.90")123 ax.set_xscale("log"); ax.set_xticks(Gs); ax.set_xticklabels(Gs)124 ax.set_title(f"ICC = {icc}"); ax.set_xlabel("number of clusters G")125axes[0].set_ylabel("coverage of nominal 95% CI (±1.96 MC SE)")126axes[0].legend(fontsize=8, loc="lower right")127fig.tight_layout(); fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)128print(json.dumps({k: results[k] for k in ["coverage", "rel_bias_var", "failure_cells", "nominal_cells", "best_procedure_small_G"]}, indent=1))129print(json.dumps(extras, indent=1)); print(g5)130