1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6SEED = 46
  7rng = np.random.default_rng(SEED)
  8
  9
 10def mixture(a, lambdas, weights):
 11    a = np.asarray(a, dtype=float)
 12    lam = np.asarray(lambdas, dtype=float)
 13    w = np.asarray(weights, dtype=float)
 14    terms = w * np.exp(-np.outer(np.atleast_1d(a), lam))
 15    g = terms.sum(axis=1)
 16    pi = terms / g[:, None]
 17    m = pi @ lam
 18    return g, m, pi
 19
 20
 21def math_verification():
 22    # Three predictions from the mechanism:
 23    # P1: finite-difference m'(a) equals - tilted variance.
 24    # P2: m(0)=E[Lambda], and m(a) approaches min(lambda) at large pressure.
 25    # P3: q'(a)/q(a)=-m(a), i.e. relative suppression per pressure is the hazard.
 26    lam = np.array([0.25, 1.0, 4.0, 16.0])
 27    w = np.ones(4) / 4
 28    grid = np.linspace(0, 20, 401)
 29    g, m, pi = mixture(grid, lam, w)
 30    var = (pi @ (lam ** 2)) - m**2
 31    # Small centered differences test the claimed m'(a)=-Var_a(lambda).
 32    hdiff = 1e-5
 33    _, mp_fd, _ = mixture(grid + hdiff, lam, w)
 34    _, mm_fd, _ = mixture(np.maximum(grid - hdiff, 0), lam, w)
 35    dm_fd = (mp_fd - mm_fd) / (2 * hdiff)
 36    interior = grid > hdiff
 37    p1_err = float(np.max(np.abs(dm_fd[interior] + var[interior])))
 38    p1_scale = float(np.max(np.abs(var[interior])))
 39
 40    mean0 = float(m[0])
 41    min_lam = float(lam.min())
 42    large_pressure = float(m[-1])
 43    p2_initial_error = abs(mean0 - float(w @ lam))
 44    # predicted crossover to within 1% of the least susceptible component;
 45    # solve by a sweep and compare with a direct pairwise bound estimate.
 46    target = min_lam * 1.01
 47    observed_idx = np.where(m <= target)[0]
 48    observed_cross = float(grid[observed_idx[0]]) if len(observed_idx) else float("inf")
 49    # For this mixture, at large a the next component ratio is exp(-(1-.25)a);
 50    # a conservative 1% dominance estimate uses equal weights and ratio <= .01.
 51    # Quantitative prediction obtained directly from the asymptotic target:
 52    # solve m(a)=1.01*min(lambda), using monotonicity of m.
 53    lo, hi = 0.0, 100.0
 54    for _ in range(80):
 55        mid = (lo + hi) / 2
 56        if mixture(np.array([mid]), lam, w)[1][0] <= target:
 57            hi = mid
 58        else:
 59            lo = mid
 60    predicted_cross = hi
 61
 62    # P3 across several pressure values, with centered finite differences.
 63    h = 1e-4
 64    gp, mp, _ = mixture(grid + h, lam, w)
 65    gm, mm, _ = mixture(np.maximum(grid - h, 0), lam, w)
 66    # avoid boundary and use derivative of log q
 67    valid = grid > 0
 68    log_deriv = (np.log(gp) - np.log(gm)) / (2*h)
 69    p3_err = float(np.max(np.abs(log_deriv[valid] + m[valid])))
 70
 71    # Parameter sweep confirms stronger heterogeneity gives larger early curvature:
 72    sweep = []
 73    for high in [1.0, 2.0, 4.0, 8.0, 16.0, 32.0]:
 74        ll = np.array([0.25, high])
 75        ww = np.array([.5, .5])
 76        _, mm, pp = mixture(np.array([0.0, 0.5]), ll, ww)
 77        vv0 = float(pp[0] @ (ll**2) - mm[0]**2)
 78        vv05 = float(pp[1] @ (ll**2) - mm[1]**2)
 79        sweep.append({"high_lambda": high, "m0": float(mm[0]), "curvature_at_0": vv0,
 80                      "curvature_at_0.5": vv05})
 81    return {
 82        "p1_derivative_variance_max_abs_error": p1_err,
 83        "p1_max_variance_scale": p1_scale,
 84        "p2_m_at_zero": mean0,
 85        "p2_weighted_mean_prediction": float(w @ lam),
 86        "p2_initial_abs_error": p2_initial_error,
 87        "p2_min_lambda": min_lam,
 88        "p2_m_at_pressure_20": large_pressure,
 89        "p2_observed_within_1pct_pressure": observed_cross,
 90        "p2_pairwise_dominance_prediction": predicted_cross,
 91        "p3_log_q_derivative_hazard_max_abs_error": p3_err,
 92        "heterogeneity_sweep": sweep,
 93    }
 94
 95
 96def softmax(x):
 97    y = x - x.max(axis=1, keepdims=True)
 98    e = np.exp(y)
 99    return e / e.sum(axis=1, keepdims=True)
100
101
102def routing_trial(kind, rounds=100, B=512, E=8, beta=.90, alpha=4.0):
103    # Fixed token preferences plus a deliberately collapsed expert bias.
104    local_rng = np.random.default_rng(SEED)
105    logits = local_rng.normal(0, .65, size=(B, E))
106    logits[:, 0] += 2.3
107    logits[:, 1] += 1.0
108    p = np.zeros(E)
109    loads, cvs, overflows, entropies = [], [], [], []
110    lam = np.array([.25, 1., 4., 16.])
111    w = np.ones(4) / 4
112    for _ in range(rounds):
113        u = softmax(logits).mean(axis=0)
114        p = beta*p + (1-beta)*u
115        if kind == "mixture":
116            q, _, _ = mixture(p, lam, w)
117            adjusted = logits + alpha*np.log(q + 1e-12)[None, :]
118        elif kind == "single":
119            adjusted = logits + alpha*np.log(np.exp(-2.0*p))[None, :]
120        elif kind == "linear":
121            adjusted = logits - alpha*p[None, :]
122        else:
123            adjusted = logits
124        chosen = np.argmax(adjusted, axis=1)
125        count = np.bincount(chosen, minlength=E)
126        cap = int(math.ceil(1.10*B/E))
127        loads.append(count)
128        cvs.append(float(count.std()/(count.mean()+1e-12)))
129        overflows.append(float(np.maximum(count-cap, 0).sum()/B))
130        probs = softmax(adjusted)
131        entropies.append(float((-probs*np.log(probs+1e-12)).sum(axis=1).mean()))
132    loads = np.asarray(loads)
133    return {
134        "mean_cv_last_50": float(np.mean(cvs[-50:])),
135        "mean_overflow_last_50": float(np.mean(overflows[-50:])),
136        "mean_router_entropy_last_50": float(np.mean(entropies[-50:])),
137        "final_load": loads[-1].tolist(),
138        "max_load_last_50": float(loads[-50:].max()),
139    }
140
141
142def main():
143    math_results = math_verification()
144    routing = {k: routing_trial(k) for k in ["baseline", "mixture", "single", "linear"]}
145    result = {"seed": SEED, "math_verification": math_results, "routing": routing}
146    Path("results.json").write_text(json.dumps(result, indent=2))
147    print(json.dumps(result, indent=2))
148
149if __name__ == "__main__":
150    main()