1import json
  2import numpy as np
  3
  4SEED = 1037
  5
  6
  7def route(u, lam):
  8    return np.argmax(u - lam[None, :], axis=1)
  9
 10
 11def dual_bound(u, lam, cap):
 12    return float(lam @ cap + np.max(u - lam[None, :], axis=1).sum())
 13
 14
 15def accepted_stats(u, a, cap):
 16    counts = np.bincount(a, minlength=u.shape[1])
 17    # Capacity rule: keep highest-utility tokens within each expert.
 18    keep = np.zeros(len(a), dtype=bool)
 19    for e in range(u.shape[1]):
 20        ids = np.flatnonzero(a == e)
 21        if len(ids):
 22            ids = ids[np.argsort(u[ids, e])[::-1][:int(cap[e])]]
 23            keep[ids] = True
 24    p = float(u[np.arange(len(a))[keep], a[keep]].sum())
 25    return counts, keep, p
 26
 27
 28def math_checks(rng):
 29    # Prediction 1: weak duality has no violations for any nonnegative prices.
 30    violations = []
 31    gaps = []
 32    for _ in range(1000):
 33        n, e = 37, 5
 34        u = rng.uniform(0.05, 2.0, (n, e))
 35        cap = np.full(e, n / e)
 36        lam = rng.exponential(1.0, e)
 37        a = route(u, lam)
 38        _, keep, p = accepted_stats(u, a, cap)
 39        L = dual_bound(u, lam, cap)
 40        gaps.append(L - p)
 41        violations.append(max(0.0, p - L))
 42    # Prediction 2: local stability boundary for iid two-expert differences is rho*n < 2.
 43    # Around equal prices, load difference has slope approximately -n*(price difference),
 44    # so the linearized price-difference multiplier is 1-rho*n.
 45    n = 200
 46    rhos = [0.001, 0.005, 0.01, 0.02, 0.04]
 47    dyn = []
 48    for rho in rhos:
 49        lam = np.zeros(2)
 50        abs_imb = []
 51        for t in range(500):
 52            d = rng.uniform(-1, 1, n)
 53            u = np.column_stack([d, np.zeros(n)]) + 2.0
 54            a = route(u, lam)
 55            counts = np.bincount(a, minlength=2)
 56            lam = np.maximum(0, lam + rho * (counts - n / 2))
 57            if t >= 250:
 58                abs_imb.append(abs(counts[0] - counts[1]))
 59        dyn.append({"rho": rho, "predicted_multiplier": 1-rho*n,
 60                    "mean_abs_load_imbalance": float(np.mean(abs_imb)),
 61                    "max_abs_load_imbalance": int(np.max(abs_imb)),
 62                    "final_price_difference": float(lam[0]-lam[1])})
 63    # Prediction 3: at zero prices, routing is utility greedy; positive prices emerge
 64    # exactly when an expert is persistently over capacity.
 65    cap = np.array([50., 50.])
 66    n = 100
 67    u = np.zeros((n, 2)); u[:70, 0] = 1.0; u[70:, 1] = 1.0
 68    lam = np.zeros(2)
 69    zero_counts = np.bincount(route(u, lam), minlength=2)
 70    for _ in range(30):
 71        a = route(u, lam); c = np.bincount(a, minlength=2)
 72        lam = np.maximum(0, lam + .02*(c-cap))
 73    return {"weak_duality": {"trials":1000, "max_violation":float(max(violations)),
 74                             "min_gap":float(min(gaps)), "mean_gap":float(np.mean(gaps))},
 75            "stability_sweep": dyn,
 76            "zero_price_prediction": {"initial_greedy_counts":zero_counts.tolist(),
 77                                       "final_prices":lam.tolist(),
 78                                       "final_counts":np.bincount(route(u,lam),minlength=2).tolist(),
 79                                       "predicted": "overloaded expert gets positive price and its routed load falls"}}
 80
 81
 82def mini_experiment(rng):
 83    E, n, cap_each, batches = 4, 128, 32, 500
 84    bias = np.array([0.75, 0.35, 0.05, -0.2])
 85    def batch():
 86        return rng.normal(0, 0.35, (n,E)) + bias[None,:]
 87    results = {}
 88    for name, rho in [("baseline_greedy", None), ("dual_rho_0.001", .001),
 89                      ("dual_rho_0.01", .01), ("dual_rho_0.1", .1)]:
 90        lam = np.zeros(E); vals=[]; over=[]; vars_=[]; gaps=[]; prices=[]
 91        for t in range(batches):
 92            u = batch()
 93            a = route(u, np.zeros(E) if rho is None else lam)
 94            counts, keep, p = accepted_stats(u, a, np.full(E, cap_each))
 95            vals.append(p); over.append(int(np.maximum(counts-cap_each,0).sum()))
 96            vars_.append(float(np.var(counts)))
 97            if rho is not None:
 98                gaps.append(dual_bound(u,lam,np.full(E,cap_each))-p)
 99                lam = np.maximum(0, lam + rho*(counts-cap_each))
100                prices.append(lam.copy())
101        results[name] = {"mean_accepted_utility":float(np.mean(vals)),
102                         "mean_overflow_tokens":float(np.mean(over)),
103                         "mean_load_variance":float(np.mean(vars_)),
104                         "final_prices":(lam.tolist() if rho is not None else None),
105                         "mean_dual_gap":(float(np.mean(gaps)) if gaps else None),
106                         "p95_dual_gap":(float(np.percentile(gaps,95)) if gaps else None)}
107    return results
108
109
110def main():
111    rng = np.random.default_rng(SEED)
112    out = {"seed":SEED, "math":math_checks(rng), "mini_experiment":mini_experiment(rng)}
113    with open("results.json", "w") as f: json.dump(out, f, indent=2)
114    print(json.dumps(out, indent=2))
115
116if __name__ == "__main__": main()