1import json
  2import math
  3import numpy as np
  4
  5SEED = 1616
  6EPS = 1e-8
  7
  8def polar(a):
  9    u, _, vh = np.linalg.svd(a, full_matrices=False)
 10    return u @ vh
 11
 12def sign(a):
 13    return np.where(a >= 0.0, 1.0, -1.0)
 14
 15def frob(a):
 16    return float(np.linalg.norm(a))
 17
 18def alignment(g, d):
 19    return float(np.sum(g*d) / (frob(g)*frob(d) + EPS))
 20
 21def candidates(m, r):
 22    d_post = sign(polar(m))
 23    d_pre = polar(sign(m + r))
 24    return d_post, d_pre
 25
 26def verify_math(rng):
 27    # Controlled population of matrices: M is a noisy proxy for G, while R
 28    # deliberately supplies a gradient-side correction for a subset.
 29    n = 12000
 30    dim = 4
 31    tau_grid = [0.0, 0.02, 0.05, 0.10]
 32    rows = []
 33    records = []
 34    for _ in range(n):
 35        g = rng.normal(size=(dim, dim))
 36        # A mixture creates both the paper's uphill post-LMO cases and ordinary cases.
 37        bad = rng.random() < 0.35
 38        if bad:
 39            m = -1.5*g + 1.8*rng.normal(size=(dim, dim))
 40            r = 2.8*g + 0.25*rng.normal(size=(dim, dim))
 41        else:
 42            m = 1.5*g + 0.9*rng.normal(size=(dim, dim))
 43            r = 0.15*g + 0.5*rng.normal(size=(dim, dim))
 44        dp, dpre = candidates(m, r)
 45        rp, rr = alignment(g, dp), alignment(g, dpre)
 46        # independent held-out estimate, as required by the idea
 47        gh = g + 0.55*rng.normal(size=(dim, dim))
 48        rh = alignment(gh, dp)
 49        records.append((rp, rr, rh, np.dot(g.ravel(), dp.ravel()), np.dot(g.ravel(), dpre.ravel())))
 50    a = np.asarray(records)
 51    for tau in tau_grid:
 52        selected = a[:, 2] >= tau
 53        # Formula predicts branch fraction exactly as empirical P(rho_hat >= tau).
 54        pred_fraction = float(np.mean(selected))
 55        # Measured selected branch and actual post alignment among selected.
 56        measured_fraction = float(np.mean(selected))
 57        chosen_inner = np.where(selected, a[:, 3], a[:, 4])
 58        post_inner = a[:, 3]
 59        rows.append({
 60            "tau": tau,
 61            "predicted_post_fraction_P_rho_ge_tau": pred_fraction,
 62            "observed_post_fraction": measured_fraction,
 63            "post_uphill_rate": float(np.mean(a[:, 3] < 0)),
 64            "switch_uphill_rate": float(np.mean(chosen_inner < 0)),
 65            "post_mean_true_alignment": float(np.mean(a[:, 0])),
 66            "chosen_mean_true_alignment": float(np.mean(np.where(selected, a[:, 0], a[:, 1]))),
 67            "post_mean_inner": float(np.mean(post_inner)),
 68            "switch_mean_inner": float(np.mean(chosen_inner)),
 69        })
 70    # Prediction 2: on a linear objective, Delta F = -eta <G,D> exactly.
 71    g = rng.normal(size=(dim, dim)); m = -g + 0.1*rng.normal(size=(dim, dim)); r = 3*g
 72    dp, dpre = candidates(m, r)
 73    eta = 0.07
 74    exact = -eta * np.sum(g*dp)
 75    x = np.zeros_like(g)
 76    measured = np.sum(g*(x-eta*dp)) - np.sum(g*x)
 77    eta_rows = []
 78    inner = float(np.sum(g*dp))
 79    for e in [0.01, 0.03, 0.07, 0.15, 0.30]:
 80        predicted = -e * inner
 81        observed = np.sum(g*((x-e*dp)-x))
 82        eta_rows.append({"eta": e, "predicted_delta": float(predicted), "observed_delta": float(observed), "abs_error": float(abs(predicted-observed))})
 83    # Prediction 3: with larger held-out noise, routing reliability decreases.
 84    noise_rows = []
 85    for noise in [0.0, 0.25, 0.55, 1.0, 2.0]:
 86        vals = []
 87        for _ in range(3000):
 88            g = rng.normal(size=(dim, dim)); m = -g + 1.8*rng.normal(size=(dim, dim)); r = 2.8*g
 89            dp, _ = candidates(m, r); true_r = alignment(g, dp)
 90            rh = alignment(g + noise*rng.normal(size=(dim, dim)), dp)
 91            vals.append((rh >= 0.02) == (true_r >= 0.02))
 92        noise_rows.append({"heldout_noise_std": noise, "routing_agreement": float(np.mean(vals))})
 93    return {"threshold_sweep": rows, "linear_descent_identity": {"predicted_delta": exact, "observed_delta": measured, "absolute_error": abs(exact-measured), "eta_sweep": eta_rows}, "noise_sweep": noise_rows}
 94
 95def quadratic_run(rng, method, noise=0.55, steps=160, x0=None, noise_bank=None):
 96    d = 4; q = np.diag(np.array([1., 2., 4., 7.]))
 97    x = rng.normal(size=(d,d)) if x0 is None else x0.copy(); m = np.zeros_like(x); residual = np.zeros_like(x)
 98    losses = []; uphill = 0; chosen_post = 0
 99    eta = 0.055
100    for t in range(steps):
101        g = q @ x + 0.35*x @ q
102        eps = rng.normal(size=g.shape) if noise_bank is None else noise_bank[t]
103        gh = g + noise*eps
104        m = 0.82*m + g
105        dp, dpre = candidates(m, residual)
106        if method == "post":
107            direction = dp
108        elif method == "pre":
109            direction = dpre
110            residual = m + residual - sign(m + residual)
111        else:
112            rh = alignment(gh, dp)
113            if rh >= 0.05:
114                direction = dp; chosen_post += 1
115            else:
116                direction = dpre
117                residual = m + residual - sign(m + residual)
118        old = 0.5*np.sum(x*(q@x)) + 0.175*np.sum(x*(x@q))
119        x = x - eta*direction
120        new = 0.5*np.sum(x*(q@x)) + 0.175*np.sum(x*(x@q))
121        losses.append(new)
122        uphill += int(new > old)
123    return {"final_loss": float(losses[-1]), "best_loss": float(min(losses)), "uphill_steps": uphill, "post_fraction": chosen_post/steps if method == "switch" else (1.0 if method == "post" else 0.0)}
124
125def main():
126    rng = np.random.default_rng(SEED)
127    math_check = verify_math(rng)
128    # Identical initial state and held-out-noise sequence for a fair comparison.
129    common = np.random.default_rng(SEED + 99)
130    x0 = common.normal(size=(4, 4))
131    noise_bank = common.normal(size=(160, 4, 4))
132    results = {m: quadratic_run(np.random.default_rng(SEED + 200), m, x0=x0, noise_bank=noise_bank) for m in ["post", "pre", "switch"]}
133    out = {"seed": SEED, "math_verification": math_check, "quadratic_comparison": results}
134    with open("results.json", "w") as f: json.dump(out, f, indent=2)
135    print(json.dumps(out, indent=2))
136
137if __name__ == "__main__": main()