1import json, math, os
  2import numpy as np
  3from sklearn.linear_model import LogisticRegression
  4from sklearn.metrics import accuracy_score
  5
  6SEED = 2392
  7rng = np.random.default_rng(SEED)
  8
  9
 10def allocation(q, bar_sigma2):
 11    q = np.asarray(q, dtype=float)
 12    v = bar_sigma2 * len(q) * (1.0 / q) / np.sum(1.0 / q)
 13    return v
 14
 15
 16def toy_sweep():
 17    # Gaussian latent z~N(0,I), scalar task t=a^T z+eta.  In this toy,
 18    # s_k=|a_k| is the exact task sensitivity proxy.
 19    K = 8
 20    bar = 0.25
 21    contrasts = [1, 2, 4, 8, 16, 32]
 22    rows = []
 23    for c in contrasts:
 24        # Two high-relevance and six low-relevance coordinates.
 25        a = np.array([c, c] + [1.0] * (K - 2))
 26        q = (np.abs(a) + 1e-8) / (np.mean(np.abs(a)) + 1e-8)
 27        v = allocation(q, bar)
 28        # Exact predicted task-noise variance, and Monte Carlo observed value.
 29        pred = float(np.sum(a * a * v))
 30        uniform = float(np.sum(a * a * bar))
 31        z = rng.normal(size=(400000, K))
 32        eps = rng.normal(size=(400000, K)) * np.sqrt(v)
 33        observed = float(np.var((z + eps) @ a - z @ a))
 34        rows.append({
 35            "contrast": c,
 36            "predicted_total_variance": float(K * bar),
 37            "observed_total_variance": float(np.mean(np.sum(eps * eps, axis=1))),
 38            "predicted_task_noise": pred,
 39            "observed_task_noise": observed,
 40            "uniform_task_noise": uniform,
 41            "adaptive_over_uniform": pred / uniform,
 42            "high_relevance_noise_fraction": float(np.sum(v[:2]) / np.sum(v)),
 43        })
 44    # Separate direct checks of the formula over random relevance vectors.
 45    budget_errors = []
 46    ordering_ok = []
 47    for _ in range(100):
 48        s = np.exp(rng.normal(size=K))
 49        q = (s + 1e-6) / (np.mean(s) + 1e-6)
 50        v = allocation(q, bar)
 51        budget_errors.append(abs(np.mean(v) - bar))
 52        ordering_ok.append(np.all(np.argsort(s) == np.argsort(-v)))
 53    return {
 54        "K": K, "bar_sigma2": bar, "rows": rows,
 55        "budget_max_abs_error": float(max(budget_errors)),
 56        "inverse_ordering_fraction": float(np.mean(ordering_ok)),
 57    }
 58
 59
 60def make_data(n, d, weights, noise, seed):
 61    r = np.random.default_rng(seed)
 62    x = r.normal(size=(n, d))
 63    score = x @ weights + r.normal(scale=noise, size=n)
 64    y = (score > 0).astype(np.int64)
 65    return x, y
 66
 67
 68def gradient_relevance(x, y):
 69    # A task probe supplies the operational MI sensitivity estimate:
 70    # average absolute gradient of per-example BCE wrt each latent coordinate.
 71    probe = LogisticRegression(C=10.0, max_iter=300, random_state=SEED)
 72    probe.fit(x, y)
 73    p = probe.predict_proba(x)[:, 1]
 74    w = probe.coef_[0]
 75    grad = np.abs((p - y)[:, None] * w[None, :])
 76    return grad.mean(axis=0), np.abs(w)
 77
 78
 79def evaluate(xtr, ytr, xte, yte, method, bar_sigma2, repeats=8):
 80    s, magnitude = gradient_relevance(xtr, ytr)
 81    d = xtr.shape[1]
 82    if method == "mi_gradient":
 83        q = (s + 1e-8) / (np.mean(s) + 1e-8)
 84    elif method == "magnitude":
 85        q = (magnitude + 1e-8) / (np.mean(magnitude) + 1e-8)
 86    elif method == "random":
 87        q = np.random.default_rng(SEED + 17).permutation((magnitude + 1e-8) / (np.mean(magnitude) + 1e-8))
 88    else:
 89        q = np.ones(d)
 90    v = allocation(q, bar_sigma2)
 91    acc = []
 92    for j in range(repeats):
 93        r = np.random.default_rng(SEED + 1000 + j)
 94        noisy_tr = xtr + r.normal(size=xtr.shape) * np.sqrt(v)
 95        noisy_te = xte + r.normal(size=xte.shape) * np.sqrt(v)
 96        clf = LogisticRegression(C=10.0, max_iter=300, random_state=SEED)
 97        clf.fit(noisy_tr, ytr)
 98        acc.append(accuracy_score(yte, clf.predict(noisy_te)))
 99    return {
100        "accuracy_mean": float(np.mean(acc)), "accuracy_std": float(np.std(acc)),
101        "mean_variance": float(np.mean(v)), "variances": v.tolist(),
102        "scores": s.tolist(), "probe_magnitude": magnitude.tolist(),
103    }
104
105
106def mini_experiment():
107    d = 8
108    true_w = np.array([3.0, 2.5, 1.8, 1.2, 0.7, 0.45, 0.25, 0.1])
109    xtr, ytr = make_data(12000, d, true_w, 0.7, SEED)
110    xte, yte = make_data(5000, d, true_w, 0.7, SEED + 1)
111    out = {}
112    for method in ["uniform", "random", "magnitude", "mi_gradient"]:
113        out[method] = evaluate(xtr, ytr, xte, yte, method, bar_sigma2=0.8)
114    return out
115
116
117def main():
118    result = {"seed": SEED, "toy": toy_sweep(), "mini": mini_experiment()}
119    with open("results.json", "w") as f:
120        json.dump(result, f, indent=2)
121    print(json.dumps(result, indent=2))
122
123if __name__ == "__main__":
124    main()