1import json
  2import math
  3import numpy as np
  4
  5SEED = 123
  6EPS = 1e-12
  7
  8
  9def bregman(y, x):
 10    y = np.asarray(y, dtype=np.float64)
 11    x = np.asarray(x, dtype=np.float64)
 12    return float(np.sum(-np.log(y / x) + (y - x) / x))
 13
 14
 15def barrier_step(p, g, alpha, tol=1e-13, max_iter=200):
 16    """Solve y_i=(1/p_i+alpha*(g_i+lambda))^-1, sum(y)=1."""
 17    p = np.asarray(p, dtype=np.float64)
 18    g = np.asarray(g, dtype=np.float64)
 19    if np.any(p <= 0) or abs(p.sum() - 1) > 1e-9:
 20        raise ValueError("p must be strictly positive and sum to one")
 21    invp = 1.0 / p
 22    # Domain is lambda > max_i(-g_i-invp_i/alpha).
 23    lo = float(np.max(-g - invp / alpha)) + 1e-14
 24    def F(lam):
 25        den = invp + alpha * (g + lam)
 26        if np.any(den <= 0):
 27            return np.inf
 28        return float(np.sum(1.0 / den) - 1.0)
 29    # F decreases continuously from +infinity to -1.
 30    left = lo
 31    right = max(1.0, float(np.max(np.abs(g)) + np.max(invp) / alpha + 1.0))
 32    while F(right) > 0:
 33        right = 2.0 * right + 1.0
 34        if right > 1e14:
 35            raise FloatingPointError("could not bracket lambda")
 36    for _ in range(max_iter):
 37        mid = (left + right) / 2.0
 38        if F(mid) > 0:
 39            left = mid
 40        else:
 41            right = mid
 42        if right - left <= tol * max(1.0, abs(mid)):
 43            break
 44    lam = (left + right) / 2.0
 45    den = invp + alpha * (g + lam)
 46    y = 1.0 / den
 47    y = y / y.sum()
 48    return y, lam, float(abs(y.sum() - 1.0)), float(np.min(den))
 49
 50
 51def math_checks():
 52    rng = np.random.default_rng(SEED)
 53    bvals = []
 54    residuals = []
 55    mins = []
 56    for _ in range(300):
 57        x = rng.dirichlet(np.ones(7) * .7)
 58        y = rng.dirichlet(np.ones(7) * .7)
 59        bvals.append(bregman(y, x))
 60        g = rng.normal(size=7)
 61        q, _, r, md = barrier_step(x, g, 0.37)
 62        residuals.append(r)
 63        mins.append((q > 0).all() and md > 0)
 64    return {"min_bregman": float(min(bvals)),
 65            "max_simplex_residual": float(max(residuals)),
 66            "all_denominators_positive": bool(all(mins))}
 67
 68
 69def prediction_sweeps():
 70    # Prediction 1: every finite alpha gives strictly positive coordinates and exact root.
 71    p = np.array([.55, .25, .15, .05])
 72    g = np.array([0., .4, 1.2, 2.0])
 73    alphas = np.logspace(-3, 3, 13)
 74    rows = []
 75    for a in alphas:
 76        q, lam, res, md = barrier_step(p, g, float(a))
 77        rows.append({"alpha": float(a), "min_q": float(q.min()),
 78                     "root_residual": res, "min_denominator": md})
 79
 80    # Prediction 2: for a boundary-seeking linear objective, 1/q_bad grows
 81    # approximately linearly in steps, with slope alpha (asymptotically).
 82    p0 = np.array([.5, .5])
 83    slopes = []
 84    for a in [0.5, 1., 2., 5., 10., 20.]:
 85        pcur = p0.copy(); inv_history = []
 86        for _ in range(80):
 87            pcur, _, _, _ = barrier_step(pcur, np.array([0., 1.]), a)
 88            inv_history.append(1.0 / pcur[1])
 89        t = np.arange(1, 81)
 90        slope = float(np.polyfit(t[30:], np.asarray(inv_history)[30:], 1)[0])
 91        slopes.append({"alpha": a, "observed_reciprocal_slope": slope,
 92                       "predicted_slope": a, "ratio": slope / a})
 93
 94    # Prediction 3: one step from p_bad has q_bad approximately p_bad/(1+alpha*p_bad)
 95    # when alpha*p_bad is large; report relative error across initial masses.
 96    scaling = []
 97    for pb in [.01, .03, .1, .3, .5]:
 98        pp = np.array([1-pb, pb]); aa = 100.
 99        q, _, _, _ = barrier_step(pp, np.array([0., 1.]), aa)
100        pred = pb / (1 + aa * pb)
101        scaling.append({"p_bad": pb, "observed_q_bad": float(q[1]),
102                        "approx_predicted_q_bad": float(pred),
103                        "relative_error": float(abs(q[1]-pred)/q[1])})
104    return {"positivity_and_root_sweep": rows,
105            "reciprocal_scaling_sweep": slopes,
106            "boundary_one_step_scaling": scaling}
107
108
109def adam_logits_step(z, grad, m, v, t, lr=.08):
110    b1, b2 = .9, .999
111    m = b1*m + (1-b1)*grad
112    v = b2*v + (1-b2)*grad*grad
113    mh = m/(1-b1**t); vh = v/(1-b2**t)
114    return z - lr*mh/(np.sqrt(vh)+1e-8), m, v
115
116
117def routing_benchmark():
118    # Convex cross entropy to an imbalanced useful routing prior. Same 160 steps.
119    target = np.array([.70, .20, .08, .019, .001])
120    target /= target.sum()
121    nsteps = 160
122    start = np.ones(5)/5
123    results = {}
124    # Barrier operates directly on p with exact gradient d(-r log p)/dp.
125    for name in ["barrier", "eg", "adam_logits"]:
126        p = start.copy(); z = np.log(start); m = np.zeros(5); v = np.zeros(5)
127        losses = []
128        for t in range(1, nsteps+1):
129            if name == "barrier":
130                g = -target / p
131                p, _, _, _ = barrier_step(p, g, .08)
132            elif name == "eg":
133                g = -target / p
134                p = p * np.exp(-.08*g); p /= p.sum()
135            else:
136                ex = np.exp(z-z.max()); p = ex/ex.sum()
137                gradz = p-target
138                z, m, v = adam_logits_step(z, gradz, m, v, t)
139                p = np.exp(z-z.max()); p /= p.sum()
140            loss = float(-np.sum(target*np.log(np.maximum(p, 1e-300))))
141            losses.append(loss)
142        entropy = float(-np.sum(p*np.log(np.maximum(p, 1e-300))))
143        results[name] = {"final_loss": losses[-1], "loss_step_20": losses[19],
144                         "final_entropy": entropy, "min_probability": float(p.min()),
145                         "losses": losses}
146    return {"target": target.tolist(), "steps": nsteps, "methods": results}
147
148
149def main():
150    out = {"seed": SEED, "math_checks": math_checks(),
151           "predictions": prediction_sweeps(), "routing_benchmark": routing_benchmark()}
152    with open("results.json", "w") as f:
153        json.dump(out, f, indent=2)
154    print(json.dumps(out, indent=2))
155
156if __name__ == "__main__":
157    main()