1import json, math
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2078
  6np.set_printoptions(precision=6, suppress=True)
  7
  8
  9def prox(p, c, eta):
 10    """KL prox: argmin_q <c,q> + KL(q||p)/eta."""
 11    logw = np.log(np.maximum(p, 1e-300)) - eta * c
 12    logw -= np.max(logw)
 13    q = np.exp(logw)
 14    return q / q.sum()
 15
 16
 17def make_operator(lam):
 18    # On the tangent plane, A is skew-symmetric with eigenvalues +/- i*lam.
 19    return lam * np.array([[0., -1., 1.], [1., 0., -1.], [-1., 1., 0.]]) / math.sqrt(3.)
 20
 21
 22def md_step(p, A, eta):
 23    return prox(p, A @ p, eta)
 24
 25
 26def mp_step(p, A, eta):
 27    z = prox(p, A @ p, eta)
 28    return prox(p, A @ z, eta)
 29
 30
 31def radius(p):
 32    return float(np.linalg.norm(p - np.ones(3) / 3))
 33
 34
 35def jacobian_factor(step, u, A, eta, eps=1e-7):
 36    # Two tangent directions; both have the same theoretical factor here.
 37    basis = np.array([[1., -1., 0.], [1., 1., -2.]])
 38    basis /= np.linalg.norm(basis, axis=1)[:, None]
 39    vals = []
 40    for v in basis:
 41        p = u + eps * v
 42        vals.append(radius(step(p, A, eta)) / radius(p))
 43    return float(np.mean(vals))
 44
 45
 46def trajectory(step, p0, A, eta, steps=60):
 47    p = p0.copy(); rs = [radius(p)]
 48    for _ in range(steps):
 49        p = step(p, A, eta); rs.append(radius(p))
 50    return np.asarray(rs)
 51
 52
 53def main():
 54    u = np.ones(3) / 3
 55    p0 = u + np.array([0.035, -0.021, -0.014])
 56    rows = []
 57
 58    # Prediction 1: zero coupling gives exactly identity updates.
 59    A = make_operator(0.)
 60    zero_effect = max(float(np.max(np.abs(md_step(p0, A, .7) - p0))),
 61                      float(np.max(np.abs(mp_step(p0, A, .7) - p0))))
 62
 63    # At uniform p, linearized KL prox is I-(eta/3)A. Thus x=eta*Lambda/3,
 64    # MD factor=sqrt(1+x^2), MP factor=sqrt(1-x^2+x^4), neutral at x=1.
 65    etas = [0.30, 1.00, 2.00]
 66    lambdas = [0.50, 1.00, 2.00, 3.00, 5.00]
 67    for eta in etas:
 68        for lam in lambdas:
 69            x = eta * lam / 3.0
 70            A = make_operator(lam)
 71            measured_md = jacobian_factor(md_step, u, A, eta)
 72            measured_mp = jacobian_factor(mp_step, u, A, eta)
 73            rows.append({
 74                'eta': eta, 'Lambda': lam, 'x': x,
 75                'md_measured': measured_md, 'md_predicted': math.sqrt(1 + x*x),
 76                'mp_measured': measured_mp, 'mp_predicted': math.sqrt(1 - x*x + x**4)
 77            })
 78
 79    # Equal-step coupled-routing proxy at x=2/3: MP should damp the cycle while
 80    # standard one-stage entropic mirror descent amplifies it.
 81    eta, lam = 1.00, 2.00
 82    A = make_operator(lam)
 83    md_r = trajectory(md_step, p0, A, eta)
 84    mp_r = trajectory(mp_step, p0, A, eta)
 85
 86    def vi_residual(p):
 87        c = A @ p
 88        return float(np.max(c - c @ p))
 89
 90    pmd, pmp = p0.copy(), p0.copy(); md_vi, mp_vi = [], []
 91    for _ in range(60):
 92        md_vi.append(vi_residual(pmd)); mp_vi.append(vi_residual(pmp))
 93        pmd = md_step(pmd, A, eta); pmp = mp_step(pmp, A, eta)
 94
 95    tr = [r for r in rows if abs(r['eta'] - 2.0) < 1e-12]
 96    below = [r['x'] for r in tr if r['mp_measured'] < 1.0]
 97    above = [r['x'] for r in tr if r['mp_measured'] > 1.0]
 98    md_err = max(abs(r['md_measured'] - r['md_predicted']) for r in rows)
 99    mp_err = max(abs(r['mp_measured'] - r['mp_predicted']) for r in rows)
100    result = {
101        'seed': SEED,
102        'zero_lambda_max_change': zero_effect,
103        'zero_lambda_prediction_confirmed': bool(zero_effect < 1e-14),
104        'local_sweep': rows,
105        'max_abs_local_prediction_error_md': float(md_err),
106        'max_abs_local_prediction_error_mp': float(mp_err),
107        'transition': {
108            'predicted_x': 1.0,
109            'observed_below_max_x': float(max(below)) if below else None,
110            'observed_above_min_x': float(min(above)) if above else None,
111            'note': 'MP factor is below 1 for 0<x<1 and above 1 for x>1.'
112        },
113        'mini_experiment': {
114            'eta': eta, 'Lambda': lam, 'x': eta * lam / 3.0,
115            'md_radius_initial': float(md_r[0]), 'md_radius_final': float(md_r[-1]),
116            'mp_radius_initial': float(mp_r[0]), 'mp_radius_final': float(mp_r[-1]),
117            'md_vi_initial': float(md_vi[0]), 'md_vi_final': float(md_vi[-1]),
118            'mp_vi_initial': float(mp_vi[0]), 'mp_vi_final': float(mp_vi[-1]),
119            'md_radius_max': float(md_r.max()), 'mp_radius_max': float(mp_r.max())
120        }
121    }
122    Path('results.json').write_text(json.dumps(result, indent=2))
123    print(json.dumps(result, indent=2))
124
125
126if __name__ == '__main__':
127    main()