1import json
  2import math
  3import numpy as np
  4
  5SEED = 1515
  6
  7def invsqrt(A, eps=1e-12):
  8    w, V = np.linalg.eigh((A + A.T) * 0.5)
  9    w = np.maximum(w, eps)
 10    return (V * (1.0 / np.sqrt(w))) @ V.T
 11
 12def metric_from_history(G, beta, lam, rho, alpha):
 13    p, m = G.shape
 14    S = beta * np.zeros((p, p)) + (1-beta) * (G @ G.T) / m + lam * np.eye(p)
 15    R = invsqrt(S)
 16    if alpha == 1.0:
 17        # Projection is mathematically unused in this baseline branch.
 18        P = np.zeros((p, p))
 19        Q = R @ R
 20    else:
 21        K = G.T @ np.linalg.solve(S, G) + rho * np.eye(m)
 22        P = R @ G @ np.linalg.solve(K, G.T) @ R
 23        Q = R @ ((1-alpha) * P + alpha * np.eye(p)) @ R
 24    return S, P, Q
 25
 26def predicted_threshold(H, Q):
 27    # For general noncommuting H,Q this is the exact quadratic GD threshold.
 28    A = (Q @ H + H @ Q) * 0.5
 29    return 2.0 / np.max(np.linalg.eigvalsh(A))
 30
 31def run_fixed(H, G, eta, beta, lam, rho, alpha, steps=100):
 32    _, _, Q = metric_from_history(G, beta, lam, rho, alpha)
 33    theta = np.array([1.0, 1.0])
 34    losses = []
 35    for _ in range(steps):
 36        losses.append(0.5 * theta @ H @ theta)
 37        theta = theta - eta * Q @ (H @ theta)
 38        if not np.all(np.isfinite(theta)) or np.linalg.norm(theta) > 1e8:
 39            return False, losses
 40    return losses[-1] < losses[0] and np.linalg.norm(theta) < 1.0, losses
 41
 42def observed_threshold(H, G, beta, lam, rho, alpha):
 43    pred = predicted_threshold(H, metric_from_history(G,beta,lam,rho,alpha)[2])
 44    # binary search for largest tested stable eta; stability means monotone decay over 300 steps
 45    lo, hi = 0.0, pred * 1.8
 46    for _ in range(30):
 47        mid = (lo + hi) / 2
 48        ok, _ = run_fixed(H,G,mid,beta,lam,rho,alpha,steps=300)
 49        if ok: lo = mid
 50        else: hi = mid
 51    return lo, pred
 52
 53def quadratic_training(H, Gbase, eta, lam, rho, alpha, rng, steps=160, noise=0.25):
 54    p = H.shape[0]
 55    theta = np.array([1.0, 1.0])
 56    hist = []
 57    losses = []
 58    for t in range(steps):
 59        g = H @ theta + noise * rng.normal(size=p)
 60        hist.append(g.copy())
 61        G = np.column_stack(hist[-Gbase:])
 62        _, _, Q = metric_from_history(G, 0.0, lam, rho, alpha)
 63        theta -= eta * Q @ g
 64        losses.append(0.5 * theta @ H @ theta)
 65    return float(np.mean(losses[-30:])), float(np.mean(losses)), losses
 66
 67def main():
 68    H = np.diag([12.0, 1.0])
 69    # The history only identifies the stiff coordinate: this makes the weak-support claim explicit.
 70    G = np.array([[1.0, 0.8, 1.2, 0.9], [0.0, 0.0, 0.0, 0.0]])
 71    beta, rho, alpha = 0.0, 0.15, 0.08
 72
 73    # Prediction 1: increasing ridge shrinks the supported metric eigenvalue and raises the stability limit.
 74    ridge_rows = []
 75    for lam in [0.01, 0.03, 0.1, 0.3, 1.0]:
 76        obs, pred = observed_threshold(H,G,beta,lam,rho,alpha)
 77        Q = metric_from_history(G,beta,lam,rho,alpha)[2]
 78        ridge_rows.append({'lambda':lam, 'q_supported':float(Q[1,1]), 'pred_eta_c':pred, 'observed_eta_c':obs})
 79
 80    # Prediction 2: increasing rho weakens projection-supported directions and raises eta_c.
 81    rho_rows = []
 82    for r in [0.01, 0.03, 0.1, 0.3, 1.0]:
 83        obs, pred = observed_threshold(H,G,beta,0.1,r,alpha)
 84        Q = metric_from_history(G,beta,0.1,r,alpha)[2]
 85        rho_rows.append({'rho':r, 'q_supported':float(Q[0,0]), 'pred_eta_c':pred, 'observed_eta_c':obs})
 86
 87    # Prediction 3: alpha controls unsupported directions: at alpha=0 they receive zero metric.
 88    alpha_rows = []
 89    for a in [0.0, 0.02, 0.1, 0.3, 1.0]:
 90        Q = metric_from_history(G,beta,0.1,rho,a)[2]
 91        alpha_rows.append({'alpha':a, 'q_unsupported':float(Q[1,1]), 'q_supported':float(Q[0,0])})
 92
 93    # Secondary matched-step comparison on noisy quadratic, averaging fixed seeds.
 94    baseline, idea = [], []
 95    eta = 0.055
 96    for seed in range(8):
 97        rb = np.random.default_rng(SEED + seed)
 98        ri = np.random.default_rng(SEED + seed)
 99        baseline.append(quadratic_training(H, 4, eta, 1e-12, 0.0, 1.0, rb)[0])
100        idea.append(quadratic_training(H, 4, eta, 0.1, rho, alpha, ri)[0])
101    result = {
102        'settings': {'H_diag':[12.0,1.0], 'history_rank':1, 'beta':beta, 'rho':rho, 'alpha':alpha},
103        'ridge_sweep':ridge_rows, 'rho_sweep':rho_rows, 'alpha_sweep':alpha_rows,
104        'comparison': {'baseline_mean_final_loss':float(np.mean(baseline)), 'baseline_std':float(np.std(baseline)), 'idea_mean_final_loss':float(np.mean(idea)), 'idea_std':float(np.std(idea)), 'eta':eta},
105        'checks': {
106            'ridge_predicted_monotonic': bool(all(ridge_rows[i]['pred_eta_c'] < ridge_rows[i+1]['pred_eta_c'] for i in range(4))),
107            'ridge_observed_monotonic': bool(all(ridge_rows[i]['observed_eta_c'] < ridge_rows[i+1]['observed_eta_c'] for i in range(4))),
108            'rho_predicted_monotonic': bool(all(rho_rows[i]['pred_eta_c'] < rho_rows[i+1]['pred_eta_c'] for i in range(4))),
109            'rho_observed_monotonic': bool(all(rho_rows[i]['observed_eta_c'] < rho_rows[i+1]['observed_eta_c'] for i in range(4))),
110            'unsupported_zero_at_alpha0': bool(alpha_rows[0]['q_unsupported'] < 1e-10),
111            'unsupported_increases_with_alpha': bool(all(alpha_rows[i]['q_unsupported'] < alpha_rows[i+1]['q_unsupported'] for i in range(4)))
112        }
113    }
114    with open('results.json','w') as f: json.dump(result,f,indent=2)
115    print(json.dumps(result, indent=2))
116
117if __name__ == '__main__': main()