1import json
  2import numpy as np
  3
  4
  5def score_distribution(a, u, n, kappa=0.05, epsilon=0.05):
  6    score = (np.abs(a) + kappa) * u / np.sqrt(np.maximum(n, 1.0))
  7    q = score / score.sum()
  8    return (1.0 - epsilon) * q + epsilon / len(q)
  9
 10
 11def allocate_controller(a, u, budget, batch=16, kappa=0.05, epsilon=0.05, seed=0):
 12    rng = np.random.default_rng(seed)
 13    c = len(u)
 14    n = np.ones(c, dtype=int)
 15    while n.sum() < budget:
 16        q = score_distribution(a, u, n, kappa, epsilon)
 17        take = min(batch, budget - n.sum())
 18        chosen = rng.choice(c, size=take, p=q)
 19        n += np.bincount(chosen, minlength=c)
 20    return n
 21
 22
 23def allocate_fixed(weights, budget):
 24    weights = np.asarray(weights, dtype=float)
 25    raw = budget * weights / weights.sum()
 26    n = np.floor(raw).astype(int)
 27    n = np.maximum(n, 1)
 28    while n.sum() < budget:
 29        n[np.argmax(raw - n)] += 1
 30    while n.sum() > budget:
 31        eligible = np.where(n > 1)[0]
 32        j = eligible[np.argmax(n[eligible] - raw[eligible])]
 33        n[j] -= 1
 34    return n
 35
 36
 37def variance_prediction(u, n):
 38    return float(np.sum(np.asarray(u) ** 2 / np.asarray(n)))
 39
 40
 41def empirical_variance(u, n, reps=4000, seed=11):
 42    rng = np.random.default_rng(seed)
 43    errors = np.zeros(reps)
 44    for c in range(len(u)):
 45        errors += rng.normal(0.0, u[c], size=(reps, int(n[c]))).mean(axis=1)
 46    return float(np.var(errors, ddof=1))
 47
 48
 49def main():
 50    rng = np.random.default_rng(947)
 51    C = 32
 52    budget = 4096
 53    # Rare families have small probability but remain important signed terms.
 54    P = np.exp(np.linspace(0.0, -5.0, C)); P /= P.sum()
 55    s = np.where(np.arange(C) % 3 == 0, -1.0, 1.0)
 56    E = 0.8 + 1.2 * rng.random(C)
 57    a = s * P * E
 58    u = 0.15 * (0.4 + 2.5 * np.sqrt(P / P.min()))
 59    u *= (0.8 + 0.4 * rng.random(C))
 60
 61    # Prediction 1: independent-noise variance should equal sum u^2/n.
 62    n_test = allocate_fixed(np.ones(C), budget)
 63    pred_var = variance_prediction(u, n_test)
 64    obs_var = empirical_variance(u, n_test)
 65    rel_err = abs(obs_var - pred_var) / pred_var
 66
 67    # Prediction 2: optimal allocation n proportional to u; compare a sweep.
 68    budgets = [512, 1024, 2048, 4096]
 69    opt_rows = []
 70    for b in budgets:
 71        nu = allocate_fixed(u, b)
 72        uniform = allocate_fixed(np.ones(C), b)
 73        opt_rows.append({
 74            "budget": b,
 75            "uniform_predicted_variance": variance_prediction(u, uniform),
 76            "u_proportional_predicted_variance": variance_prediction(u, nu),
 77            "ratio_uniform_over_u_proportional": variance_prediction(u, uniform) / variance_prediction(u, nu),
 78        })
 79
 80    # Prediction 3: controller score should preferentially increase count of
 81    # high |a|u families, while epsilon enforces a nonzero floor.
 82    eps_rows = []
 83    target = (np.abs(a) + 0.05) * u
 84    for eps in [0.0, 0.01, 0.05, 0.2, 0.5]:
 85        n_ctrl = allocate_controller(a, u, budget, epsilon=eps, seed=947)
 86        corr = float(np.corrcoef(n_ctrl, target)[0, 1])
 87        min_share = float(np.min(n_ctrl / n_ctrl.sum()))
 88        eps_rows.append({"epsilon": eps, "count_score_correlation": corr, "minimum_family_share": min_share, "predicted_variance": variance_prediction(u, n_ctrl)})
 89
 90    # Mode-collapse test: deliberately hide a high-uncertainty family from the
 91    # controller. With epsilon=0 it receives only its initial sample; an
 92    # exploration mixture should restore samples and reduce true variance.
 93    hidden = int(np.argmax(u))
 94    u_bad = u.copy()
 95    u_bad[hidden] = u.min() * 0.01
 96    collapse_rows = []
 97    for eps in [0.0, 0.01, 0.05, 0.2]:
 98        n_bad = allocate_controller(a, u_bad, budget, epsilon=eps, seed=947)
 99        collapse_rows.append({
100            "epsilon": eps,
101            "hidden_family": hidden,
102            "hidden_family_count": int(n_bad[hidden]),
103            "true_predicted_variance": variance_prediction(u, n_bad),
104        })
105
106    # Secondary equal-budget comparison: uniform, magnitude-only, and controller.
107    n_uniform = allocate_fixed(np.ones(C), budget)
108    n_mag = allocate_fixed(np.abs(a) + 0.05, budget)
109    n_ctrl = allocate_controller(a, u, budget, epsilon=0.05, seed=947)
110    comparison = {
111        "uniform": {"predicted_variance": variance_prediction(u, n_uniform), "empirical_variance": empirical_variance(u, n_uniform, seed=21)},
112        "magnitude_only": {"predicted_variance": variance_prediction(u, n_mag), "empirical_variance": empirical_variance(u, n_mag, seed=22)},
113        "uncertainty_guided_controller": {"predicted_variance": variance_prediction(u, n_ctrl), "empirical_variance": empirical_variance(u, n_ctrl, seed=23)},
114    }
115    result = {
116        "setup": {"families": C, "budget_including_initial_counts": budget, "P_min": float(P.min()), "P_max": float(P.max())},
117        "prediction_1_variance_identity": {"predicted": pred_var, "observed": obs_var, "relative_error": rel_err},
118        "prediction_2_optimal_allocation_sweep": opt_rows,
119        "prediction_3_exploration_sweep": eps_rows,
120        "prediction_3_mode_collapse_sweep": collapse_rows,
121        "equal_budget_comparison": comparison,
122    }
123    with open("results.json", "w") as f:
124        json.dump(result, f, indent=2)
125    print(json.dumps(result, indent=2))
126
127
128if __name__ == "__main__":
129    main()