1#!/usr/bin/env python32"""Unequal block sizes (Proposition 4): size-weighted vs equal-weight aggregation of local half-block jackknife estimates, with3q = 0, 1 Newton steps; logistic p = 3 and exponential. Output: results/raw/skew_<model>.csv"""4import sys, pathlib, time, hashlib5import numpy as np, pandas as pd6ROOT = pathlib.Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT / "src"))7from pamle.likelihoods import make_model8from pamle.partitioning import skewed_sizes, unequal_blocks9from pamle.jackknife import vandermonde_weights101112def run(model_name, p, N, K, ratio, reps, theta0, seed, out):13 model = make_model(model_name, p); rng = np.random.default_rng(seed)14 sizes = skewed_sizes(N, K, ratio); sizes = (sizes // 4) * 4; sizes[-1] += N - sizes.sum(); sizes[-1] -= sizes[-1] % 4; N_eff = int(sizes.sum())15 blocks = unequal_blocks(sizes); w2 = vandermonde_weights((1, 2))16 rows = []; t0 = time.time()17 for b in range(reps):18 data = model.simulate(rng, N_eff, theta0) if model_name != "exponential" else {"x": rng.exponential(1 / theta0[0], size=N_eff)}19 full = {k: v[None] for k, v in data.items()}; cen = model.fit_batched(full, tol=1e-12, maxit=100)[0][0]20 J = np.zeros((K, p)); raw = np.zeros((K, p)); g_k = np.zeros((K, p)); H_k = np.zeros((K, p, p))21 for k, idx in enumerate(blocks):22 dk = {kk: v[idx][None] for kk, v in data.items()}23 thf = model.fit_batched(dk)[0][0]; half = len(idx) // 224 ha = model.fit_batched({kk: v[idx[:half]][None] for kk, v in data.items()})[0][0]; hb = model.fit_batched({kk: v[idx[half:]][None] for kk, v in data.items()})[0][0]25 raw[k] = thf; J[k] = w2[0] * thf + w2[1] * 0.5 * (ha + hb)26 for agg, wts in (("weighted", sizes / sizes.sum()), ("equal", np.ones(K) / K)):27 for name, est in (("raw", raw), ("jack2", J)):28 th0 = wts @ est29 th = th0.copy()30 res = {"rep": b, "N": N_eff, "K": K, "ratio": ratio, "agg": agg, "pilot": name, "model": model_name}31 for q in range(3):32 res[f"gap_q{q}"] = float(np.sqrt(N_eff) * np.linalg.norm(th - cen))33 # aggregated derivatives (always size-weighted = full-sample derivatives)34 g = np.zeros(p); H = np.zeros((p, p))35 for k, idx in enumerate(blocks):36 gk, Hk = model.score_hessian(th[None], {kk: v[idx][None] for kk, v in data.items()}); g += sizes[k] / N_eff * gk[0]; H += sizes[k] / N_eff * Hk[0]37 try:38 th = th + np.linalg.solve(-H, g)39 except np.linalg.LinAlgError: # singular curvature at a wild pilot: record as failure40 th = np.full(p, np.nan)41 for qq in range(q + 1, 3):42 res[f"gap_q{qq}"] = float("nan")43 break44 rows.append(res)45 if (b + 1) % max(1, reps // 5) == 0:46 print(f"[{model_name} K={K} ratio={ratio}] {b+1}/{reps} {time.time()-t0:.0f}s", flush=True); pd.DataFrame(rows).to_csv(out, index=False)47 pd.DataFrame(rows).to_csv(out, index=False)484950if __name__ == "__main__":51 import sys52 only = sys.argv[1] if len(sys.argv) > 1 else None53 (ROOT / "results" / "raw").mkdir(parents=True, exist_ok=True)54 for model_name, p, theta0 in (("exponential", 1, np.array([1.0])), ("logistic", 3, np.array([0.5, 1.0, -1.0]))):55 for K, ratio in ((256, 1.0), (256, 4.0), (256, 16.0), (1024, 4.0), (1024, 16.0)):56 seed = 20261028 + int(hashlib.sha256(f"skew {model_name} {K} {ratio}".encode()).hexdigest(), 16) % 10 ** 657 if only and only != f"{model_name}_K{K}_r{int(ratio)}":58 continue59 run(model_name, p, 2 ** 18, K, ratio, 300 if model_name == "logistic" else 1000, theta0, seed, ROOT / "results" / "raw" / f"skew_{model_name}_K{K}_r{int(ratio)}.csv")60