1import json
  2import math
  3import numpy as np
  4
  5
  6def project_simplex(v, z=1.0):
  7    """Euclidean projection onto {x >= 0, sum(x)=z}."""
  8    v = np.asarray(v, dtype=float)
  9    u = np.sort(v)[::-1]
 10    cssv = np.cumsum(u) - z
 11    rho = np.nonzero(u - cssv / (np.arange(len(v)) + 1) > 0)[0]
 12    if len(rho) == 0:
 13        return np.full_like(v, z / len(v))
 14    theta = cssv[rho[-1]] / (rho[-1] + 1)
 15    return np.maximum(v - theta, 0.0)
 16
 17
 18def payoff(x, a, b):
 19    # x: groups x experts; a: groups x experts; b: experts
 20    D = b + x.sum(axis=0)
 21    return float(np.sum(a * x / D[None, :]))
 22
 23
 24def equimarginal_route(a, b, budget=1.0, steps=20, eta=0.5,
 25                       tol=1e-10, backtrack=True, init=None):
 26    """Simultaneous projected ascent on each group's own payoff.
 27
 28    The update is x_j <- simplex_project(x_j + eta*g_j), where
 29    g_ji = a_ji*(D_i-x_ji)/D_i^2. Backtracking keeps batch payoff monotone.
 30    """
 31    a = np.asarray(a, dtype=float)
 32    b = np.asarray(b, dtype=float)
 33    m, n = a.shape
 34    x = np.full((m, n), budget / n) if init is None else np.array(init, float, copy=True)
 35    history = [payoff(x, a, b)]
 36    for _ in range(steps):
 37        D = b + x.sum(axis=0)
 38        g = a * (D[None, :] - x) / (D[None, :] ** 2)
 39        local_eta = eta
 40        old = history[-1]
 41        while True:
 42            y = np.vstack([project_simplex(row + local_eta * grad, budget)
 43                           for row, grad in zip(x, g)])
 44            new = payoff(y, a, b)
 45            if (not backtrack) or new >= old - 1e-12 or local_eta < 1e-9:
 46                break
 47            local_eta *= 0.5
 48        x = y
 49        history.append(new)
 50        if np.max(np.abs(g - g.mean(axis=1, keepdims=True))) < tol:
 51            break
 52    return x, np.asarray(history)
 53
 54
 55def independent_softmax(a, temperature=1.0, budget=1.0):
 56    z = a / max(temperature, 1e-12)
 57    z -= z.max(axis=1, keepdims=True)
 58    p = np.exp(z)
 59    p /= p.sum(axis=1, keepdims=True)
 60    return budget * p
 61
 62
 63def cv(v):
 64    v = np.asarray(v)
 65    return float(v.std() / (v.mean() + 1e-12))
 66
 67
 68def transition_check():
 69    # One-player water filling has x_i=sqrt(a_i*b_i/c)-b_i and
 70    # a new expert enters when c=a_i/b_i. These are exact predictions.
 71    a = np.array([[4.0, 2.0, 1.0, 0.25]])
 72    b = np.ones(4)
 73    ratios = (a[0] / b)
 74    predicted = []
 75    for k in range(1, 4):
 76        c = ratios[k]
 77        r = np.sum(np.sqrt(a[0, :k] * b[:k] / c) - b[:k])
 78        predicted.append(float(r))
 79    # Probe just below/above each threshold and detect positive support.
 80    observed = []
 81    for r0 in predicted:
 82        for r in (r0 * (1 - 2e-3), r0 * (1 + 2e-3)):
 83            x, _ = equimarginal_route(a, b, budget=r, steps=500, eta=0.8)
 84            observed.append({"budget": float(r), "support": int((x[0] > 1e-6).sum())})
 85    return predicted, observed
 86
 87
 88def marginal_check():
 89    rng = np.random.default_rng(7)
 90    a = np.exp(rng.normal(0, .7, size=(5, 6)))
 91    b = np.exp(rng.normal(0, .3, size=6))
 92    x, hist = equimarginal_route(a, b, steps=80, eta=1.0)
 93    D = b + x.sum(axis=0)
 94    g = a * (D[None, :] - x) / D[None, :] ** 2
 95    residuals = []
 96    for j in range(len(a)):
 97        active = x[j] > 1e-5
 98        c = np.median(g[j, active]) if active.any() else 0
 99        # KKT residual: active rates should equal c; inactive rates <= c.
100        active_err = np.max(np.abs(g[j, active] - c)) if active.any() else 0
101        inactive_violation = np.max(np.maximum(g[j, ~active] - c, 0)) if (~active).any() else 0
102        residuals.append(float(max(active_err, inactive_violation)))
103    return float(hist[0]), float(hist[-1]), float(max(residuals)), x
104
105
106def load_sweep():
107    # Same affinity pattern for all groups: independent routing overloads the
108    # high-affinity experts; congestion routing's predicted response is to
109    # spread mass as group count raises congestion.
110    rng = np.random.default_rng(11)
111    n, mmax = 8, 64
112    b = np.ones(n) * 0.25
113    base = np.array([3.0, 2.6, 2.2, 1.8, 1.4, 1.1, .9, .7])
114    rows = []
115    for m in [4, 8, 16, 32, 64]:
116        noise = rng.normal(0, .05, size=(m, n))
117        a = np.maximum(base[None, :] * np.exp(noise), 1e-5)
118        x_c, h = equimarginal_route(a, b, steps=30, eta=1.0)
119        x_s = independent_softmax(np.log(a), temperature=.7)
120        cap = (m / n) * 1.15
121        rows.append({"groups": m,
122                     "congestion_cv": cv(x_c.sum(0)),
123                     "softmax_cv": cv(x_s.sum(0)),
124                     "congestion_overflow": float(np.maximum(x_c.sum(0)-cap, 0).sum()),
125                     "softmax_overflow": float(np.maximum(x_s.sum(0)-cap, 0).sum()),
126                     "congestion_payoff": payoff(x_c, a, b),
127                     "softmax_payoff": payoff(x_s, a, b),
128                     "iterations": len(h)-1})
129    return rows
130
131
132def main():
133    pred, obs = transition_check()
134    initial, final, residual, x = marginal_check()
135    loads = load_sweep()
136    result = {"transition_predicted_budgets": pred,
137              "transition_observed": obs,
138              "marginal_initial_payoff": initial,
139              "marginal_final_payoff": final,
140              "marginal_kkt_residual": residual,
141              "marginal_row_sum_max_error": float(np.max(np.abs(x.sum(1)-1))),
142              "load_sweep": loads}
143    print(json.dumps(result, indent=2))
144
145if __name__ == '__main__':
146    main()