1#!/usr/bin/env python32"""Monte Carlo driver for WP28. One *cell* = (model, p, N, K, DGP options, reps, seed). Every estimator is computed on the3same simulated sample and the same partition (paired design). Output: one CSV per cell in results/raw/ with one row per4replication, columns <estimator>_<j> for parameter j, plus diagnostics and timings, and a JSON manifest.56Usage7 python3 experiments/run_monte_carlo.py --config experiments/configs/main_exponential.yaml [--cells 0,3,5] [--reps 200]8 python3 experiments/run_monte_carlo.py --list --config ...9"""10from __future__ import annotations11import argparse, hashlib, json, os, pathlib, platform, socket, sys, time12import numpy as np, yaml1314ROOT = pathlib.Path(__file__).resolve().parents[1]15sys.path.insert(0, str(ROOT / "src"))16from pamle.likelihoods import make_model, gather_blocks, STATUS_OK # noqa: E40217from pamle.algorithms import estimate_all, aggregated_score_hessian, local_split_fits # noqa: E40218from pamle.partitioning import skewed_sizes # noqa: E40219from pamle import __version__ # noqa: E4022021SE_ESTIMATORS = ("central", "r1_q0", "r1_q1", "r1_q2", "r2_q0", "r2_q1", "r3_q1", "cs_q1")222324def cell_seed(base: int, cell: dict) -> int:25 key = json.dumps({k: cell[k] for k in sorted(cell) if k not in ("reps",)}, sort_keys=True, default=str)26 return (base + int(hashlib.sha256(key.encode()).hexdigest(), 16)) % (2 ** 63)272829def expand_cells(cfg: dict):30 cells = []31 for g in cfg["grids"]:32 for N in g["N"]:33 for m in g["m"]:34 if N % m:35 continue36 for p in g.get("p", [1]):37 if m < g.get("min_m_over_p", 4) * p:38 continue39 c = {"model": g["model"], "p": p, "N": int(N), "K": int(N // m), "m": int(m), "reps": int(g["reps"]),40 "splits_max": int(g.get("splits_max", 3)), "q_max": int(g.get("q_max", 2)),41 "dgp": g.get("dgp", {}), "label": g.get("label", g["model"]), "extras": bool(g.get("extras", True))}42 cells.append(c)43 return cells444546def true_theta(cell):47 p = cell["p"]; dgp = cell["dgp"]48 if cell["model"] == "exponential":49 return np.array([float(dgp.get("theta0", 1.0))])50 if "theta0" in dgp:51 return np.asarray(dgp["theta0"], float)52 # default: intercept 0.3, slopes alternating +-0.5 decaying53 th = np.array([0.3] + [0.5 * (-1) ** j / (1 + 0.15 * j) for j in range(p - 1)])54 if cell["model"] == "poisson":55 th = th * 0.556 return th575859def simulate(model, cell, theta0, rng):60 dgp = cell["dgp"]; N = cell["N"]; p = cell["p"]61 kind = dgp.get("kind", "well_specified")62 if cell["model"] == "exponential":63 x = rng.exponential(1.0 / theta0[0], size=N)64 if kind == "contaminated":65 eps = float(dgp.get("eps", 0.01)); scale = float(dgp.get("scale", 10.0))66 mask = rng.random(N) < eps67 x[mask] = rng.exponential(scale / theta0[0], size=mask.sum())68 return {"x": x}69 cov = None70 if dgp.get("cond", None) is not None and p > 1: # covariate covariance with prescribed condition number71 q = p - 1; ev = np.geomspace(1.0, 1.0 / float(dgp["cond"]), q)72 cov = np.diag(ev)73 data = model.simulate(rng, N, theta0, cov=cov)74 if kind == "misspecified": # omitted quadratic term in the index75 X = data["X"]; eta = X @ theta0 + float(dgp.get("omitted", 0.5)) * (X[:, 1] ** 2 - 1.0)76 if cell["model"] == "logistic":77 data["y"] = (rng.random(N) < 1 / (1 + np.exp(-eta))).astype(float)78 else:79 data["y"] = rng.poisson(np.exp(eta)).astype(float)80 return data818283def run_cell(cell, out_dir: pathlib.Path, base_seed: int, reps_override=None, verbose=True):84 model = make_model(cell["model"], cell["p"]); p = model.p85 reps = reps_override or cell["reps"]86 seed = cell_seed(base_seed, cell); rng = np.random.default_rng(seed)87 theta0 = true_theta(cell)88 N, K, m = cell["N"], cell["K"], cell["m"]89 skew = float(cell["dgp"].get("skew", 0.0))90 name = f'{cell["label"]}_p{p}_N{N}_K{K}'91 rows = []; t_start = time.time()92 for b in range(reps):93 data = simulate(model, cell, theta0, rng)94 t0 = time.perf_counter()95 if skew > 1.0:96 # unequal contiguous blocks; estimate_all needs equal shapes -> use K blocks of sizes from skewed_sizes,97 # handled by padding is not possible: we instead run a dedicated weighted/naive comparison (see run_skewed)98 raise NotImplementedError("use run_skewed.py for unequal partitions")99 db = gather_blocks(data, np.arange(N).reshape(K, m))100 out, diag = estimate_all(model, db, splits_max=cell["splits_max"], q_max=cell["q_max"], extras=cell["extras"])101 t_est = time.perf_counter() - t0102 row = {"rep": b, "seed": seed, "N": N, "K": K, "m": m, "p": p, "model": cell["model"], "label": cell["label"],103 "t_estimate_s": t_est, "central_status": diag["central_status"]}104 for s, cnt in diag["local_status_counts"].items():105 row[f"local_fail_s{s}"] = int(sum(cnt[1:]))106 for r in range(1, cell["splits_max"] + 1):107 row[f"valid_r{r}"] = int(diag[f"valid_r{r}"].sum())108 nd = diag[f"newton_r{r}"]109 for q, d in enumerate(nd, start=1):110 row[f"cond_r{r}_q{q}"] = d["cond"]; row[f"step_r{r}_q{q}"] = d["step_norm"]; row[f"grad_r{r}_q{q}"] = d["grad_norm"]111 for k, v in out.items():112 for j in range(p):113 row[f"{k}_{j}"] = float(v[j])114 # standard errors (Fisher at the estimate; sandwich at the centralized estimate)115 for k in SE_ESTIMATORS:116 if k in out and np.all(np.isfinite(out[k])):117 g, H, gk, Hk = aggregated_score_hessian(model, db, out[k])118 try:119 cov = np.linalg.inv(-H)120 for j in range(p):121 row[f"se_{k}_{j}"] = float(np.sqrt(max(cov[j, j], 0) / N))122 except np.linalg.LinAlgError:123 pass124 th = np.broadcast_to(out["central"], (K, p))125 J = model.outer_scores(th, db).mean(axis=0); g, H, _, _ = aggregated_score_hessian(model, db, out["central"])126 Hi = np.linalg.inv(H); S = Hi @ J @ Hi127 for j in range(p):128 row[f"se_sand_central_{j}"] = float(np.sqrt(max(S[j, j], 0) / N))129 for j in range(p):130 row[f"theta0_{j}"] = float(theta0[j])131 rows.append(row)132 if verbose and (b + 1) % max(1, reps // 10) == 0:133 print(f"[{name}] {b+1}/{reps} reps, {time.time()-t_start:.1f}s", flush=True)134 import pandas as pd135 df = pd.DataFrame(rows)136 out_dir.mkdir(parents=True, exist_ok=True)137 f = out_dir / f"{name}.csv"; df.to_csv(f, index=False)138 sha = hashlib.sha256(f.read_bytes()).hexdigest()139 manifest = {"run_id": name, "cell": cell, "reps": reps, "seed": seed, "pamle_version": __version__,140 "python": platform.python_version(), "numpy": np.__version__, "host": socket.gethostname(),141 "machine": platform.machine(), "wall_seconds": time.time() - t_start, "raw_output_sha256": sha,142 "git_commit": os.popen(f"git -C {ROOT} rev-parse --short HEAD").read().strip(), "status": "completed"}143 (out_dir / f"{name}.manifest.json").write_text(json.dumps(manifest, indent=1))144 return f, manifest145146147def main():148 ap = argparse.ArgumentParser()149 ap.add_argument("--config", required=True)150 ap.add_argument("--cells", default=None, help="comma-separated cell indices (default all)")151 ap.add_argument("--reps", type=int, default=None)152 ap.add_argument("--out", default=str(ROOT / "results" / "raw"))153 ap.add_argument("--list", action="store_true")154 a = ap.parse_args()155 cfg = yaml.safe_load(open(a.config))156 cells = expand_cells(cfg)157 if a.list:158 for i, c in enumerate(cells):159 print(i, c["label"], "p", c["p"], "N", c["N"], "K", c["K"], "m", c["m"], "reps", c["reps"])160 return161 idx = range(len(cells)) if a.cells is None else [int(x) for x in a.cells.split(",")]162 for i in idx:163 f, man = run_cell(cells[i], pathlib.Path(a.out), int(cfg.get("seed_base", 20261028)), a.reps)164 print("wrote", f, f"{man['wall_seconds']:.1f}s")165166167if __name__ == "__main__":168 main()169