1#!/usr/bin/env python32"""Ablations (Phase 7, Table 6): (a) Hessian/score quantization bits; (b) exact vs stale vs single-machine Hessian; (c) split3patterns (1,2,4) vs (1,2,3) vs (1,3) and orders r = 1..4 with q = 0..3; on logistic (p = 8) and exponential models.4Paired replications (same sample for all variants). Output: results/raw/ablation_<model>.csv (one row per replication × variant).5"""6import sys, pathlib, time, hashlib, argparse7import numpy as np, pandas as pd8ROOT = pathlib.Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT / "src"))9from pamle.likelihoods import make_model, gather_blocks10from pamle.algorithms import local_split_fits, worker_corrected_estimates, pilot_estimator, global_newton, aggregated_score_hessian, quantize_relative, newton_step11from pamle.jackknife import DEFAULT_SPLITS1213PATTERNS = {"1": (1,), "1-2": (1, 2), "1-2-4": (1, 2, 4), "1-2-4-8": (1, 2, 4, 8), "1-2-3": (1, 2, 3), "1-3": (1, 3), "1-2-3-4": (1, 2, 3, 4)}14BITS = [2, 3, 4, 6, 8, 12, 16, 24, 52]151617def run(model_name, p, N, m, reps, theta0, seed, out):18 model = make_model(model_name, p); rng = np.random.default_rng(seed); K = N // m19 rows = []; t0 = time.time()20 for b in range(reps):21 data = model.simulate(rng, N, theta0) if model_name != "exponential" else {"x": rng.exponential(1 / theta0[0], size=N)}22 db = gather_blocks(data, np.arange(N).reshape(K, m))23 full = {k: v.reshape((1, N) + v.shape[2:]) for k, v in db.items()}24 cen = model.fit_batched(full, tol=1e-12, maxit=100)[0][0]25 needed = sorted({s for pat in PATTERNS.values() for s in pat if m % s == 0})26 fits = local_split_fits(model, db, tuple(needed))27 H_bar = fits[1][2][:, 0].mean(axis=0)28 base = {"rep": b, "N": N, "K": K, "m": m, "p": p, "model": model_name}29 def add(variant, th, **extra):30 row = dict(base, variant=variant, gap=float(np.sqrt(N) * np.linalg.norm(th - cen)), gap0=float(np.sqrt(N) * (th[0] - cen[0])), **extra)31 rows.append(row)32 for name, pat in PATTERNS.items():33 if any(m % s for s in pat):34 continue35 J, ok = worker_corrected_estimates({s: fits[s] for s in pat}, pat); th0 = pilot_estimator(J)36 r = len(pat)37 path, _ = global_newton(model, db, th0, q=3)38 for q in range(4):39 add(f"pat{name}_q{q}", path[q], pattern=name, r=r, q=q, hessian="exact", bits=64, valid_frac=float(ok.mean()))40 if name in ("1", "1-2", "1-2-4"):41 ps, _ = global_newton(model, db, th0, q=3, hessian_mode="stale", stale_hessian=H_bar)42 pg, _ = global_newton(model, db, th0, q=3, hessian_mode="single")43 for q in range(1, 4):44 add(f"pat{name}_q{q}_stale", ps[q], pattern=name, r=r, q=q, hessian="stale", bits=64)45 add(f"pat{name}_q{q}_single", pg[q], pattern=name, r=r, q=q, hessian="single", bits=64)46 if name in ("1", "1-2", "1-2-4"):47 g, H, _, _ = aggregated_score_hessian(model, db, th0)48 for bits in BITS:49 thH, _ = newton_step(th0, g, quantize_relative(H, bits)); add(f"pat{name}_q1_Hbits{bits}", thH, pattern=name, r=r, q=1, hessian="quantizedH", bits=bits)50 thg, _ = newton_step(th0, quantize_relative(g, bits), H); add(f"pat{name}_q1_gbits{bits}", thg, pattern=name, r=r, q=1, hessian="quantizedg", bits=bits)51 thb, _ = newton_step(th0, quantize_relative(g, bits), quantize_relative(H, bits)); add(f"pat{name}_q1_bothbits{bits}", thb, pattern=name, r=r, q=1, hessian="quantizedBoth", bits=bits)52 if (b + 1) % max(1, reps // 10) == 0:53 print(f"[{model_name} N={N} m={m}] {b+1}/{reps} {time.time()-t0:.0f}s", flush=True)54 pd.DataFrame(rows).to_csv(out, index=False)55 pd.DataFrame(rows).to_csv(out, index=False)565758if __name__ == "__main__":59 ap = argparse.ArgumentParser(); ap.add_argument("--model", default="logistic"); ap.add_argument("--reps", type=int, default=400); a = ap.parse_args()60 (ROOT / "results" / "raw").mkdir(parents=True, exist_ok=True)61 if a.model == "logistic":62 theta0 = np.array([0.3] + [0.5 * (-1) ** j / (1 + 0.15 * j) for j in range(7)])63 for K, m in [(2048, 96), (1024, 192), (512, 384), (256, 768), (4096, 192), (1024, 768)]:64 N = K * m65 seed = 20261028 + int(hashlib.sha256(f"abl logistic {N} {m}".encode()).hexdigest(), 16) % 10 ** 666 run("logistic", 8, N, m, a.reps, theta0, seed, ROOT / "results" / "raw" / f"ablation_logistic_N{N}_m{m}.csv")67 else:68 for K, m in [(16384, 12), (16384, 24), (16384, 48), (16384, 96), (16384, 192), (65536, 24), (65536, 96), (262144, 24)]:69 N = K * m70 seed = 20261028 + int(hashlib.sha256(f"abl exp {N} {m}".encode()).hexdigest(), 16) % 10 ** 671 run("exponential", 1, N, m, a.reps, np.array([1.0]), seed, ROOT / "results" / "raw" / f"ablation_exponential_N{N}_m{m}.csv")72