1import json
  2import math
  3from itertools import permutations
  4import numpy as np
  5
  6
  7def stationary_solve(Q):
  8    """High-accuracy stationary distribution for numerical verification."""
  9    n = Q.shape[0]
 10    A = Q.T.copy()
 11    b = np.zeros(n)
 12    A[-1, :] = 1.0
 13    b[-1] = 1.0
 14    v = np.linalg.solve(A, b)
 15    v = np.maximum(v, 0.0)
 16    return v / v.sum()
 17
 18SEED = 278
 19rng = np.random.default_rng(SEED)
 20
 21
 22def stationary_power(Q, steps=20000):
 23    """Stationary distribution for a row generator using P=I+Q/c."""
 24    n = Q.shape[0]
 25    c = max(1.0, float(np.max(-np.diag(Q)))) * 1.000001
 26    P = np.eye(n) + Q / c
 27    v = np.ones(n) / n
 28    for _ in range(steps):
 29        nv = v @ P
 30        if np.max(np.abs(nv - v)) < 1e-14:
 31            v = nv
 32            break
 33        v = nv
 34    return np.maximum(v, 0) / np.sum(v)
 35
 36
 37def make_generator(scores, groups, gamma=100.0, tau=1.0):
 38    n = len(groups)
 39    R = np.zeros((n, n))
 40    for x in range(n):
 41        for y in range(n):
 42            if x != y:
 43                R[x, y] = math.exp(float(scores[x, y]) / tau) * (gamma if groups[x] == groups[y] else 1.0)
 44    Q = R.copy()
 45    np.fill_diagonal(Q, -np.sum(R, axis=1))
 46    return R, Q
 47
 48
 49def tree_stationary(R):
 50    """Enumerate directed spanning trees oriented toward each root (small sanity check)."""
 51    n = R.shape[0]
 52    weights = np.zeros(n)
 53    # Each non-root chooses one outgoing parent; retain choices with a path to root.
 54    for root in range(n):
 55        total = 0.0
 56        nonroots = [x for x in range(n) if x != root]
 57        for parents in __import__('itertools').product(range(n), repeat=len(nonroots)):
 58            nxt = {x: p for x, p in zip(nonroots, parents)}
 59            if any(x == nxt[x] for x in nonroots):
 60                continue
 61            valid = True
 62            product = 1.0
 63            for x in nonroots:
 64                product *= R[x, nxt[x]]
 65                z = x
 66                seen = set()
 67                while z != root:
 68                    if z in seen or z not in nxt:
 69                        valid = False
 70                        break
 71                    seen.add(z)
 72                    z = nxt[z]
 73                if not valid:
 74                    break
 75            if valid:
 76                total += product
 77        weights[root] = total
 78    return weights / weights.sum()
 79
 80
 81def fast_class_reduction(R, groups):
 82    unique = np.unique(groups)
 83    mus = []
 84    for g in unique:
 85        ix = np.flatnonzero(groups == g)
 86        local = R[np.ix_(ix, ix)]
 87        q = local.copy()
 88        np.fill_diagonal(q, -np.sum(local, axis=1))
 89        mus.append(stationary_solve(q))
 90    m = len(unique)
 91    QB = np.zeros((m, m))
 92    for i, g in enumerate(unique):
 93        ix = np.flatnonzero(groups == g)
 94        for j, h in enumerate(unique):
 95            if i != j:
 96                jx = np.flatnonzero(groups == h)
 97                QB[i, j] = np.sum(mus[i][:, None] * R[np.ix_(ix, jx)])
 98        QB[i, i] = -np.sum(QB[i])
 99    return stationary_solve(QB), mus, QB
100
101
102def route_experiment():
103    n_groups, per_group = 4, 3
104    n = n_groups * per_group
105    groups = np.repeat(np.arange(n_groups), per_group)
106    # Construct a nonuniform class prior from asymmetric cross-group rates.
107    # This prior counteracts an explicit raw-router bias toward later groups.
108    desired = np.array([0.40, 0.30, 0.20, 0.10])
109    pair_scores = rng.normal(0, 0.35, (n, n))
110    for i in range(n):
111        for j in range(n):
112            if i == j:
113                pair_scores[i, j] = -np.inf
114            elif groups[i] != groups[j]:
115                pair_scores[i, j] = np.log(desired[groups[j]])
116    R, Q = make_generator(pair_scores, groups, gamma=100.0)
117    bar_pi, mus, QB = fast_class_reduction(R, groups)
118
119    batches, tokens = 100, 256
120    # Domain composition drifts, creating a realistic source of group load oscillation.
121    base_means = np.full((n_groups, n), -2.0)
122    for d in range(n_groups):
123        base_means[d, groups == d] = 3.0
124        base_means[d, (groups == (d + 1) % n_groups)] = 1.0
125    cap = math.ceil(1.20 * tokens / n)
126    metrics = {"baseline": {"cv": [], "overflow": [], "churn": [], "correct": []},
127               "idea": {"cv": [], "overflow": [], "churn": [], "correct": []}}
128    previous = {"baseline": None, "idea": None}
129    beta = 0.5
130    eps = 1e-12
131    for b in range(batches):
132        # Slowly changing mixture with substantial batch-to-batch noise.
133        dominant = (b // 5) % n_groups
134        domprob = np.full(n_groups, 0.08)
135        domprob[dominant] = 0.76
136        domains = rng.choice(n_groups, size=tokens, p=domprob)
137        raw_bias = 0.5 * groups[None, :]
138        logits = base_means[domains] + raw_bias + rng.normal(0, 1.8, (tokens, n))
139        # Baseline: flat noisy router. Idea: stationary class prior before top-1 selection.
140        for name, adjusted in (("baseline", logits),
141                               ("idea", logits + beta * np.log(bar_pi[groups] + eps))):
142            selected = np.argmax(adjusted, axis=1)
143            loads = np.bincount(selected, minlength=n) / tokens
144            metrics[name]["cv"].append(float(np.std(loads) / (np.mean(loads) + eps)))
145            metrics[name]["overflow"].append(float(np.sum(np.maximum(np.bincount(selected, minlength=n) - cap, 0)) / tokens))
146            if previous[name] is not None:
147                metrics[name]["churn"].append(float(np.mean(selected != previous[name])))
148            previous[name] = selected
149            metrics[name]["correct"].append(float(np.mean(groups[selected] == domains)))
150    summary = {}
151    for name, vals in metrics.items():
152        summary[name] = {k: float(np.mean(v)) for k, v in vals.items()}
153    summary["capacity_per_expert"] = cap
154    summary["bar_pi"] = bar_pi.tolist()
155    return summary
156
157
158def main():
159    # Independent 4-state verification of the tree theorem.
160    s = rng.normal(0, 0.7, (4, 4))
161    s = np.where(np.eye(4, dtype=bool), -np.inf, s)
162    R4, Q4 = make_generator(s, np.zeros(4, dtype=int), gamma=1.0)
163    pi_power = stationary_power(Q4)
164    pi_tree = tree_stationary(R4)
165    tree_error = float(np.max(np.abs(pi_power - pi_tree)))
166
167    # Fast within-class theorem check: full expert stationary vs reduced class stationary.
168    groups = np.repeat(np.arange(3), 2)
169    scores = rng.normal(0, 0.5, (6, 6))
170    scores = np.where(np.eye(6, dtype=bool), -np.inf, scores)
171    R, Q = make_generator(scores, groups, gamma=1e5)
172    full_pi = stationary_solve(Q)
173    bar_pi, mus, QB = fast_class_reduction(R, groups)
174    full_group = np.array([full_pi[groups == g].sum() for g in range(3)])
175    reduction_error = float(np.max(np.abs(full_group - bar_pi)))
176    # Also report reduction error across scales; asymptotic reduction should improve as gamma grows.
177    reduction_errors = {}
178    for gamma in [10.0, 100.0, 1000.0, 1e5]:
179        RR, QQ = make_generator(scores, groups, gamma=gamma)
180        fp = stationary_solve(QQ)
181        bp, _, _ = fast_class_reduction(RR, groups)
182        reduction_errors[str(gamma)] = float(np.max(np.abs(np.array([fp[groups == g].sum() for g in range(3)]) - bp)))
183    out = {"seed": SEED, "tree_theorem_max_error": tree_error,
184           "fast_reduction_error_by_gamma": reduction_errors,
185           "fast_reduction_max_group_error_gamma_1e5": reduction_error,
186           "routing": route_experiment()}
187    with open('results.json', 'w') as f:
188        json.dump(out, f, indent=2, sort_keys=True)
189    print(json.dumps(out, indent=2, sort_keys=True))
190
191
192if __name__ == '__main__':
193    main()