1import json
  2from pathlib import Path
  3import numpy as np
  4from scipy.stats import norm
  5
  6
  7def newey_west(samples, lag=20):
  8    """Columnwise mean, Bartlett HAC variance, SE, z and lower-tail p-value."""
  9    x = np.asarray(samples, dtype=float)
 10    if x.ndim != 2 or x.shape[0] < 2:
 11        raise ValueError("samples must have shape [time, components], T >= 2")
 12    t = x.shape[0]
 13    lag = max(0, min(int(lag), t - 1))
 14    mean = x.mean(axis=0)
 15    c = x - mean
 16    hac = np.mean(c * c, axis=0)
 17    for k in range(1, lag + 1):
 18        gamma = np.mean(c[k:] * c[:-k], axis=0)
 19        hac += 2.0 * (1.0 - k / (lag + 1.0)) * gamma
 20    return mean, np.maximum(hac, 1e-12)
 21
 22
 23def confidence_scores(samples, delta, lag=20):
 24    mean, hac = newey_west(samples, lag)
 25    se = np.sqrt(hac / samples.shape[0])
 26    z = (mean - delta) / se
 27    return {"mean": mean, "hac": hac, "se": se, "z": z, "p": norm.cdf(z)}
 28
 29
 30def select_components(samples, delta, keep, lag=20):
 31    s = confidence_scores(samples, delta, lag)
 32    # Large p means evidence is not against useful contribution.
 33    indices = np.argsort(s["p"])[-keep:]
 34    return np.sort(indices), s
 35
 36
 37def ar1(mu, rho, sigma, t, rng):
 38    e = rng.normal(0.0, sigma * np.sqrt(1.0 - rho * rho), size=t)
 39    y = np.empty(t)
 40    y[0] = mu + rng.normal(0.0, sigma)
 41    for i in range(1, t):
 42        y[i] = mu + rho * (y[i - 1] - mu) + e[i]
 43    return y
 44
 45
 46def mechanism_sweep(seed=7):
 47    """Numerical checks of the three quantitative predictions in the idea."""
 48    rng = np.random.default_rng(seed)
 49    # Prediction 1: SE approximately proportional to T^-1/2 for iid observations.
 50    ts = np.array([100, 200, 400, 800, 1600])
 51    se_obs = []
 52    for t in ts:
 53        vals = np.stack([ar1(0.3, 0.0, 1.0, t, rng) for _ in range(250)], axis=1)
 54        se_obs.append(float(np.mean(confidence_scores(vals, 0.0, lag=0)["se"])))
 55    slope = float(np.polyfit(np.log(ts), np.log(se_obs), 1)[0])
 56
 57    # Prediction 2: AR(1) long-run variance inflation is (1+rho)/(1-rho).
 58    rhos = np.array([0.0, 0.3, 0.6, 0.8])
 59    infl_obs = []
 60    for rho in rhos:
 61        vals = np.stack([ar1(0.0, rho, 1.0, 3000, rng) for _ in range(35)], axis=1)
 62        infl_obs.append(float(np.mean(confidence_scores(vals, 0.0, lag=200)["hac"])))
 63    infl_obs = np.array(infl_obs)
 64    infl_obs /= infl_obs[0]
 65    infl_pred = (1 + rhos) / (1 - rhos)
 66    # normalize prediction consistently to the empirically estimated rho=0 baseline
 67    infl_pred /= infl_pred[0]
 68
 69    # Prediction 3: when true mean falls below Delta, rejection probability rises.
 70    delta = 0.25
 71    means = np.array([0.45, 0.30, 0.20, 0.05])
 72    reject = []
 73    for mu in means:
 74        count = 0
 75        for _ in range(180):
 76            x = ar1(mu, 0.5, 1.0, 300, rng)[:, None]
 77            if confidence_scores(x, delta, lag=20)["p"][0] < 0.05:
 78                count += 1
 79        reject.append(count / 180.0)
 80    return {
 81        "prediction_1_se_slope": {"predicted": -0.5, "observed": slope, "se_by_T": se_obs},
 82        "prediction_2_hac_inflation": {"rho": rhos.tolist(), "predicted": infl_pred.tolist(), "observed": infl_obs.tolist()},
 83        "prediction_3_rejection_power": {"true_mean": means.tolist(), "delta": delta, "rejection_rate": reject},
 84    }
 85
 86
 87def pruning_comparison(seed=11, repeats=30):
 88    """Fixed-budget proxy: compare history inference with latest-minibatch ranking."""
 89    n, keep, t = 24, 10, 160
 90    truth = np.linspace(0.04, 0.48, n)
 91    stat_utils, latest_utils = [], []
 92    for r in range(repeats):
 93        rng = np.random.default_rng(seed + r)
 94        histories = np.stack([ar1(mu, 0.75, 0.55, t, rng) for mu in truth], axis=1)
 95        stat_idx, stat = select_components(histories, 0.20, keep, lag=20)
 96        latest_idx = np.argsort(histories[-1])[-keep:]
 97        stat_utils.append(float(truth[stat_idx].sum()))
 98        latest_utils.append(float(truth[latest_idx].sum()))
 99    # Include one run's allocation and p-values for inspectability.
100    rng = np.random.default_rng(seed)
101    histories = np.stack([ar1(mu, 0.75, 0.55, t, rng) for mu in truth], axis=1)
102    stat_idx, stat = select_components(histories, 0.20, keep, lag=20)
103    latest_idx = np.argsort(histories[-1])[-keep:]
104    return {
105        "keep": keep, "delta": 0.20, "repeats": repeats,
106        "statistical_indices": stat_idx.tolist(),
107        "latest_indices": np.sort(latest_idx).tolist(),
108        "statistical_true_utility_mean": float(np.mean(stat_utils)),
109        "latest_true_utility_mean": float(np.mean(latest_utils)),
110        "statistical_win_rate": float(np.mean(np.array(stat_utils) > np.array(latest_utils))),
111        "all_true_utility": float(np.sort(truth)[-keep:].sum()),
112        "statistical_pvalues": stat["p"].tolist(),
113    }
114
115
116def main():
117    result = {"mechanism": mechanism_sweep(), "pruning": pruning_comparison()}
118    Path("results.json").write_text(json.dumps(result, indent=2))
119    print(json.dumps(result, indent=2))
120
121
122if __name__ == "__main__":
123    main()