1import json
  2import numpy as np
  3from pathlib import Path
  4
  5np.set_printoptions(precision=6, suppress=True)
  6
  7
  8def entropy_point(p, g, lam):
  9    """Simplex entropy mirror map: normalize p*exp(-lam*g), stably."""
 10    z = np.log(p) - lam * g
 11    z -= z.max()
 12    w = np.exp(z)
 13    return w / w.sum()
 14
 15
 16def entropy_root(p, g, delta, max_expand=80, bisect=80):
 17    """Solve g.(p-x(lam))=delta, or return failure if outside reachable range."""
 18    if delta <= 0:
 19        return 0.0, p.copy(), True
 20    gmin = g.min()
 21    reachable = float(np.dot(g, p) - gmin)
 22    if delta > reachable + 1e-12:
 23        return np.nan, p.copy(), False
 24    def phi(lam):
 25        x = entropy_point(p, g, lam)
 26        return float(np.dot(g, p-x) - delta)
 27    lo, hi = 0.0, 1.0
 28    while phi(hi) < 0 and hi < 1e12:
 29        hi *= 2
 30    if phi(hi) < 0:
 31        return np.nan, p.copy(), False
 32    for _ in range(bisect):
 33        mid = (lo + hi) / 2
 34        if phi(mid) >= 0:
 35            hi = mid
 36        else:
 37            lo = mid
 38    lam = (lo + hi) / 2
 39    x = entropy_point(p, g, lam)
 40    return lam, x, abs(float(np.dot(g, p-x)-delta)) < 1e-10
 41
 42
 43def simplex_project(v):
 44    u = np.sort(v)[::-1]
 45    cssv = np.cumsum(u) - 1
 46    ind = np.arange(1, len(v)+1)
 47    rho = np.nonzero(u - cssv / ind > 0)[0][-1]
 48    theta = cssv[rho] / (rho + 1)
 49    return np.maximum(v-theta, 0)
 50
 51
 52def core_sweeps():
 53    # Fixed non-symmetric point and gradient makes the predictions identifiable.
 54    p = np.array([0.55, 0.30, 0.15])
 55    g = np.array([0.8, -0.3, 1.4])
 56    mean = np.dot(p, g)
 57    variance = np.dot(p, (g-mean)**2)
 58    reachable = np.dot(p, g) - g.min()
 59
 60    # Prediction 1: as delta -> 0, lambda/delta -> 1/Var_p(g).
 61    deltas = np.logspace(-7, -2, 8)
 62    small = []
 63    for d in deltas:
 64        lam, x, ok = entropy_root(p, g, d)
 65        small.append((float(d), float(lam), float(lam/d), bool(ok)))
 66    asymptotic = 1.0 / variance
 67    observed_small = small[0][2]
 68
 69    # Prediction 2: simultaneous g, delta scaling leaves x unchanged and lambda -> lambda/a.
 70    # This is the exact homogeneity law; fixed absolute delta is not invariant.
 71    d0 = 0.23 * reachable
 72    l1, x1, ok1 = entropy_root(p, g, d0)
 73    scaling = []
 74    for a in [0.25, 0.5, 1.0, 2.0, 4.0]:
 75        la, xa, oka = entropy_root(p, a*g, a*d0)
 76        scaling.append({"a": a, "lambda": float(la), "predicted_lambda": float(l1/a),
 77                        "lambda_ratio": float(la/(l1/a)), "x_error": float(np.max(abs(xa-x1))), "ok": bool(oka)})
 78
 79    # A local fixed-delta prediction also follows from the Taylor expansion:
 80    # lambda/delta ~ 1/(a^2 Var_p(g)) as delta -> 0.
 81    fixed_delta_scaling = []
 82    tiny_delta = 1e-7
 83    for a in [0.25, 0.5, 1.0, 2.0, 4.0]:
 84        la, xa, oka = entropy_root(p, a*g, tiny_delta)
 85        predicted = tiny_delta / (a*a*variance)
 86        fixed_delta_scaling.append({"a": a, "lambda": float(la),
 87                                    "predicted_lambda": float(predicted),
 88                                    "ratio": float(la/predicted), "ok": bool(oka)})
 89    # Prediction 3: root exists iff delta <= limiting displacement, and fails above it.
 90    boundary = []
 91    for frac in [0.0, 0.25, 0.5, 0.9, 0.999999, 1.001, 1.5]:
 92        d = frac * reachable
 93        lam, x, ok = entropy_root(p, g, d)
 94        boundary.append({"fraction": frac, "delta": float(d), "success": bool(ok),
 95                         "lambda": None if not ok else float(lam)})
 96    return {"p": p.tolist(), "g": g.tolist(), "weighted_variance": float(variance),
 97            "reachable_delta": float(reachable), "small_gap": small,
 98            "small_gap_prediction_lambda_over_delta": float(asymptotic),
 99            "small_gap_observed_first": float(observed_small),
100            "scaling": scaling, "fixed_delta_local_scaling": fixed_delta_scaling, "boundary": boundary}
101
102
103def optimization_comparison(seed=7, n=8, steps=120):
104    rng = np.random.default_rng(seed)
105    q = rng.dirichlet(np.ones(n)*1.5)
106    p0 = np.ones(n)/n
107    # f=1/2||p-q||², f*=0, so the exact Polyak gap is known.
108    pe, pu = p0.copy(), p0.copy()
109    ent_losses, euc_losses = [], []
110    ent_resid, euc_resid = [], []
111    for _ in range(steps):
112        for p, kind, losses, residuals in [(pe, 'entropy', ent_losses, ent_resid), (pu, 'euclidean', euc_losses, euc_resid)]:
113            g = p-q
114            f = 0.5*np.dot(g,g)
115            delta = f
116            if kind == 'entropy':
117                lam, x, ok = entropy_root(p, g, delta)
118                if not ok: x = p
119            else:
120                # Euclidean Polyak step, with simplex projection for a fair feasible baseline.
121                alpha = delta / max(np.dot(g,g), 1e-30)
122                x = simplex_project(p-alpha*g)
123            residuals.append(abs(np.dot(g, p-x)-delta))
124            losses.append(f)
125            if kind == 'entropy': pe = x
126            else: pu = x
127    return {"target_q": q.tolist(), "steps": steps,
128            "entropy_loss_start": ent_losses[0], "entropy_loss_final": ent_losses[-1],
129            "euclidean_loss_start": euc_losses[0], "euclidean_loss_final": euc_losses[-1],
130            "entropy_loss_at_10": ent_losses[9], "euclidean_loss_at_10": euc_losses[9],
131            "entropy_max_halfspace_residual": max(ent_resid),
132            "euclidean_max_halfspace_residual_after_projection": max(euc_resid)}
133
134
135def main():
136    out = {"core": core_sweeps(), "optimization": optimization_comparison()}
137    Path("results.json").write_text(json.dumps(out, indent=2))
138    print(json.dumps(out, indent=2))
139
140if __name__ == '__main__':
141    main()