1"""T2-simulation: finite-sample coverage of CR0, CR1 and wild cluster bootstrap-t CIs for a mean.2Run from workspace root: ~/wp29/venv/bin/python out/analysis.py3Reads only data/design.csv; writes out/results.json, out/table1.csv, out/figure1.png, out/exclusions.csv.4Random numbers: per cell numpy default_rng(seed from design.csv); draw order per cell:5 u = sqrt(icc)*N(0,1) of shape (R,G); e = sqrt(1-icc)*N(0,1) of shape (R,N);6 Rademacher weights W = 2*integers(0,2,(R,B,G))-1.7"""8import json, os9import numpy as np10import pandas as pd11from scipy import stats12import matplotlib13matplotlib.use("Agg")14import matplotlib.pyplot as plt1516ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))17OUT = os.path.join(ROOT, "out")18design = pd.read_csv(os.path.join(ROOT, "data", "design.csv"))19n_raw = len(design)2021# ---------------- design audit ----------------22data_issues = []23cells = []24for _, row in design.iterrows():25 sizes = np.array([int(s) for s in str(row["cluster_sizes"]).split("|")])26 assert len(sizes) == row["G"], f"{row.cell_id}: len(cluster_sizes) != G"27 assert sizes.sum() == row["N"], f"{row.cell_id}: sum(cluster_sizes) != N"28 assert int(row["R"]) == 2000 and int(row["B"]) == 39929 cells.append(dict(cell_id=row["cell_id"], G=int(row["G"]), icc=float(row["icc"]), sizes=sizes,30 N=int(row["N"]), R=int(row["R"]), B=int(row["B"]), seed=int(row["seed"]),31 mu=float(row["mu"]), level=float(row["nominal_level"])))32assert design["cell_id"].is_unique and not design.isna().any().any()33if (design["G"] == 5).any():34 data_issues.append({35 "issue": "With G=5 Rademacher weights take only 2^5=32 distinct values (16 distinct |t*| up to sign symmetry, two of which (all weights equal) give t*=0); the B=399 bootstrap draws therefore repeat these few values and the WCB critical value is coarse.",36 "handling": "Kept as specified (B=399 i.i.d. Rademacher draws, not full enumeration); noted as a limitation in the report."})3738# exclusion ledger: nothing is excluded39pd.DataFrame([{"step": 1, "rule": "keep all design rows passing consistency checks (G=len(sizes), N=sum(sizes), R=2000, B=399)",40 "n_before": n_raw, "n_after": n_raw, "n_removed": 0,41 "rationale": "all 9 cells consistent; no filtering"}]).to_csv(os.path.join(OUT, "exclusions.csv"), index=False)4243def key(est, c):44 return f"{est}_G{c['G']}_icc{c['icc']:.1f}"4546CHUNK = 25047rows, coverage, mcse, relbias, base_cov = [], {}, {}, {}, {}48relbias_se = {}49for c in cells:50 G, N, R, B, icc, mu, sizes = c["G"], c["N"], c["R"], c["B"], c["icc"], c["mu"], c["sizes"]51 rng = np.random.default_rng(c["seed"])52 gid = np.repeat(np.arange(G), sizes)53 u = np.sqrt(icc) * rng.standard_normal((R, G))54 e = np.sqrt(1 - icc) * rng.standard_normal((R, N))55 W = 2 * rng.integers(0, 2, size=(R, B, G)) - 156 y = mu + u[:, gid] + e57 ybar = y.mean(axis=1)58 resid = y - ybar[:, None]59 S = np.zeros((R, G))60 for g in range(G):61 S[:, g] = resid[:, gid == g].sum(axis=1) # cluster residual sums62 V0 = (S ** 2).sum(axis=1) / N ** 263 V1 = G / (G - 1) * V064 var_exact = (icc * (sizes ** 2).sum() + (1 - icc) * N) / N ** 265 z = 1.96 # as specified in TASK.md66 tcrit = stats.t.ppf(0.5 + c["level"] / 2, G - 1)67 # wild cluster bootstrap-t: y* = ybar + w_g * resid; cluster sums of y*-ybar* are w_g S_g - n_g D68 q = np.empty(R)69 for s in range(0, R, CHUNK):70 Wc, Sc = W[s:s + CHUNK].astype(float), S[s:s + CHUNK]71 WS = Wc * Sc[:, None, :] # (r,B,G)72 D = WS.sum(axis=2) / N # ybar* - ybar73 Sstar = WS - sizes[None, None, :] * D[:, :, None]74 V1s = G / (G - 1) * (Sstar ** 2).sum(axis=2) / N ** 275 with np.errstate(divide="ignore", invalid="ignore"):76 tstar = np.abs(D / np.sqrt(V1s))77 tstar = np.where(np.isfinite(tstar), tstar, 0.0) # guard only (0/0 does not occur with continuous data)78 # 0.95 empirical quantile = inverted-CDF definition: the ceil(0.95*B)=380th order statistic79 q[s:s + CHUNK] = np.quantile(tstar, c["level"], axis=1, method="inverted_cdf")80 # naive i.i.d. baseline (ignores clustering)81 s2 = resid.var(axis=1, ddof=1)82 Viid = s2 / N83 hw = {"cr0": z * np.sqrt(V0), "cr1": tcrit * np.sqrt(V1), "wcb": q * np.sqrt(V1),84 "iid": z * np.sqrt(Viid)}85 err = np.abs(ybar - mu)86 for est in ["cr0", "cr1", "wcb", "iid"]:87 p = float(np.mean(err <= hw[est]))88 se = float(np.sqrt(p * (1 - p) / R))89 if est == "iid":90 base_cov[key("iid", c)] = p91 continue92 k = key(est, c)93 coverage[k], mcse[k] = p, se94 rb = None95 if est in ("cr0", "cr1"):96 rb = float(np.mean(V0 if est == "cr0" else V1) / var_exact)97 relbias[k] = rb98 Vk = V0 if est == "cr0" else V199 relbias_se[k] = float(Vk.std(ddof=1) / np.sqrt(R) / var_exact)100 rows.append(dict(estimator=est, cell_id=c["cell_id"], G=G, icc=icc, coverage=p, mc_se_coverage=se,101 rel_bias_var=rb, mean_ci_halfwidth=float(hw[est].mean()), R=R))102 print(f"{c['cell_id']} G={G} icc={icc}: " + ", ".join(f"{k2}={coverage[key(k2, c)]:.4f}" for k2 in ["cr0", "cr1", "wcb"])103 + f", iid={base_cov[key('iid', c)]:.4f}")104105table = pd.DataFrame(rows)106table["_o"] = table["estimator"].map({"cr0": 0, "cr1": 1, "wcb": 2})107table = table.sort_values(["_o", "G", "icc"]).drop(columns="_o")108table.to_csv(os.path.join(OUT, "table1.csv"), index=False, float_format="%.6f")109110# analytic E[V_CR0]/Var(ybar) (verification of the Monte Carlo bias estimates):111# S = A y with A = C - n 1'/N (C: G x N cluster indicator); E[S'S] = tr(A Sigma A'),112# Sigma = icc * C'C + (1-icc) I.113exact_rb = {}114for c in cells:115 G, N, icc, sizes = c["G"], c["N"], c["icc"], c["sizes"]116 C = (np.repeat(np.arange(G), sizes)[None, :] == np.arange(G)[:, None]).astype(float)117 A = C - np.outer(sizes, np.ones(N)) / N118 Sigma = icc * C.T @ C + (1 - icc) * np.eye(N)119 v_exact = (icc * (sizes ** 2).sum() + (1 - icc) * N) / N ** 2120 ev0 = np.trace(A @ Sigma @ A.T) / N ** 2 / v_exact121 exact_rb[key("cr0", c)] = float(ev0)122 exact_rb[key("cr1", c)] = float(G / (G - 1) * ev0)123124failure = sorted(k for k, p in coverage.items() if p < 0.90)125nominal = sorted(k for k, p in coverage.items() if abs(p - 0.95) <= 2 * mcse[k])126small = {est: float(np.mean([coverage[key(est, c)] for c in cells if c["G"] == 5])) for est in ["cr0", "cr1", "wcb"]}127best_small = max(small, key=small.get)128129results = {130 "task_id": "T2-simulation",131 "R": int(cells[0]["R"]), "B": int(cells[0]["B"]), "n_cells": len(cells),132 "coverage": coverage, "mc_se_coverage": mcse, "rel_bias_var": relbias,133 "failure_cells": failure, "n_failure_cells": len(failure), "nominal_cells": nominal,134 "best_procedure_small_G": best_small,135 "feasibility": "complete", "feasibility_reason": None,136 "data_issues": data_issues,137 "n_input_rows": {"design.csv": n_raw},138 "n_analysed": int(sum(c["R"] for c in cells)),139 "random_seed": int(cells[0]["seed"]),140 "baseline_coverage_iid": base_cov,141 "mean_coverage_small_G": small,142 "mc_se_rel_bias_var": relbias_se,143 "mean_ci_halfwidth": {f"{r['estimator']}_G{r['G']}_icc{r['icc']:.1f}": r["mean_ci_halfwidth"] for r in rows},144 "exact_rel_bias_var": exact_rb,145 "z_rel_bias_mc_vs_exact": {k: float((relbias[k] - exact_rb[k]) / relbias_se[k]) for k in relbias},146 "exact_var_ybar": {f"G{c['G']}_icc{c['icc']:.1f}": float((c['icc'] * (c['sizes'] ** 2).sum() + (1 - c['icc']) * c['N']) / c['N'] ** 2) for c in cells},147}148with open(os.path.join(OUT, "results.json"), "w") as f:149 json.dump(results, f, indent=2)150151# figure152fig, axes = plt.subplots(1, 3, figsize=(13, 4.2), sharey=True)153Gs = sorted({c["G"] for c in cells})154for ax, icc in zip(axes, [0.0, 0.2, 0.5]):155 for est, mk in [("cr0", "o"), ("cr1", "s"), ("wcb", "^"), ("iid", "x")]:156 src = base_cov if est == "iid" else coverage157 p = np.array([src[f"{est}_G{G}_icc{icc:.1f}"] for G in Gs])158 se = np.sqrt(p * (1 - p) / 2000)159 ax.errorbar(Gs, p, yerr=2 * se, marker=mk, capsize=3,160 label=est.upper() + (" (naive baseline)" if est == "iid" else ""),161 linestyle=":" if est == "iid" else "-")162 ax.axhline(0.95, color="k", lw=1, label="nominal 0.95")163 ax.axhline(0.90, color="r", lw=1, ls="--", label="failure line 0.90")164 ax.set_xscale("log"); ax.set_xticks(Gs); ax.set_xticklabels(Gs); ax.minorticks_off()165 ax.set_title(f"ICC = {icc}"); ax.set_xlabel("number of clusters G")166axes[0].set_ylabel("coverage (±2 MC SE)")167axes[0].legend(fontsize=8, loc="lower right")168fig.suptitle("Coverage of nominal 95% intervals for the mean (R=2000 per cell)")169fig.tight_layout()170fig.savefig(os.path.join(OUT, "figure1.png"), dpi=130)171print("failure cells:", failure)172print("best procedure at G=5:", best_small)173