1import json, math, os
  2import numpy as np
  3
  4SEED = 3057
  5np.random.seed(SEED)
  6
  7
  8def continuous_roots(lam, kappa, beta):
  9    return np.roots([1.0, lam + kappa + beta, beta * lam])
 10
 11
 12def fm_matrix(lam, kappa, beta, eta):
 13    rho = math.exp(-beta * eta)
 14    # theta' = (1-eta*(lambda+kappa))*theta + eta*kappa*m
 15    a = 1.0 - eta * (lam + kappa)
 16    b = eta * kappa
 17    # m' = rho*m + (1-rho)*theta'
 18    return np.array([[a, b], [(1-rho)*a, rho + (1-rho)*b]])
 19
 20
 21def spectral_radius(A):
 22    return float(np.max(np.abs(np.linalg.eigvals(A))))
 23
 24
 25def fm_rollout(lams, kappa, beta, eta, steps, x0):
 26    x = np.asarray(x0, dtype=float).copy()
 27    m = x.copy()
 28    xs, losses, forces = [], [], []
 29    rho = math.exp(-beta * eta)
 30    for _ in range(steps):
 31        g = lams * x
 32        force = kappa * (m - x)
 33        x = x - eta * g + eta * force
 34        m = rho * m + (1-rho) * x
 35        xs.append(x.copy()); losses.append(float(0.5*np.sum(lams*x*x)))
 36        forces.append(float(np.linalg.norm(force)))
 37    return np.asarray(xs), np.asarray(losses), np.asarray(forces)
 38
 39
 40def sgd_rollout(lams, eta, steps, x0):
 41    x = np.asarray(x0, dtype=float).copy(); xs=[]; losses=[]
 42    for _ in range(steps):
 43        x = x - eta*lams*x
 44        xs.append(x.copy()); losses.append(float(0.5*np.sum(lams*x*x)))
 45    return np.asarray(xs), np.asarray(losses)
 46
 47
 48def momentum_rollout(lams, eta, momentum, steps, x0):
 49    x=np.asarray(x0,dtype=float).copy(); v=np.zeros_like(x); xs=[]; losses=[]
 50    for _ in range(steps):
 51        v = momentum*v + lams*x
 52        x = x - eta*v
 53        xs.append(x.copy()); losses.append(float(0.5*np.sum(lams*x*x)))
 54    return np.asarray(xs), np.asarray(losses)
 55
 56
 57def main():
 58    # A known two-mode quadratic, plus a single-mode check for exact roots.
 59    lam, kappa, beta = 3.0, 2.0, 5.0
 60    roots = continuous_roots(lam, kappa, beta)
 61    root_residuals = [abs(r*r + (lam+kappa+beta)*r + beta*lam) for r in roots]
 62
 63    # Predicted discrete boundary is spectral radius == 1; compare with rollout.
 64    grid = np.linspace(0.005, 1.2, 240)
 65    radii = np.array([spectral_radius(fm_matrix(lam,kappa,beta,e)) for e in grid])
 66    stable_pred = grid[radii < 1.0]
 67    predicted_eta_max = float(stable_pred[-1]) if len(stable_pred) else None
 68    # empirical stability criterion: bounded and decreasing from x=1 over 500 steps
 69    empirical = []
 70    for eta in grid:
 71        _, ls, _ = fm_rollout(np.array([lam]),kappa,beta,float(eta),500,[1.0])
 72        empirical.append(bool(np.all(np.isfinite(ls)) and np.max(ls) < 1e8 and ls[-1] < ls[0]))
 73    empirical = np.asarray(empirical)
 74    empirical_eta_max = float(grid[np.where(empirical)[0][-1]]) if np.any(empirical) else None
 75
 76    # Matched tiny optimization benchmark: ill-conditioned quadratic, equal 250 steps.
 77    lams=np.array([1.0, 20.0]); x0=np.array([1.0,1.0]); steps=250
 78    configs = {
 79      'sgd': (lambda eta: sgd_rollout(lams,eta,steps,x0), [0.01,0.02,0.04,0.06,0.08]),
 80      'momentum': (lambda eta: momentum_rollout(lams,eta,0.9,steps,x0), [0.005,0.01,0.02,0.03,0.04]),
 81      'fading_memory': (lambda eta: fm_rollout(lams,2.0,5.0,eta,steps,x0), [0.005,0.01,0.02,0.03,0.04,0.05])
 82    }
 83    bench={}
 84    for name,(runner,etas) in configs.items():
 85        rows=[]
 86        for eta in etas:
 87            out=runner(eta)
 88            ls=out[1]
 89            finite=bool(np.all(np.isfinite(ls)) and np.max(ls)<1e12)
 90            rows.append({'eta':eta,'final_loss':float(ls[-1]) if finite else 1e12,
 91                         'min_loss':float(np.min(ls)) if finite else 1e12,
 92                         'stable':finite,'loss_at_50':float(ls[49]) if finite else 1e12})
 93        bench[name]=rows
 94
 95    # Compare oscillation proxy on the fastest stable choices: sign changes of stiff mode.
 96    osc={}
 97    for name in ['sgd','momentum','fading_memory']:
 98        stable_rows=[r for r in bench[name] if r['stable']]
 99        best=min(stable_rows,key=lambda r:r['final_loss'])
100        eta=best['eta']
101        if name=='sgd': traj=sgd_rollout(lams,eta,steps,x0)[0]
102        elif name=='momentum': traj=momentum_rollout(lams,eta,0.9,steps,x0)[0]
103        else: traj=fm_rollout(lams,2.0,5.0,eta,steps,x0)[0]
104        signs=np.sign(traj[:,1]); changes=int(np.sum(signs[1:]*signs[:-1]<0))
105        osc[name]={'selected_eta':eta,'final_loss':best['final_loss'],'stiff_mode_sign_changes':changes}
106
107    result={'seed':SEED,'continuous_roots':[[float(r.real),float(r.imag)] for r in roots],
108      'root_residuals':root_residuals,'discrete_stability':{
109        'lambda':lam,'kappa':kappa,'beta':beta,'predicted_eta_max':predicted_eta_max,
110        'empirical_eta_max':empirical_eta_max,'relative_gap':abs(predicted_eta_max-empirical_eta_max)/predicted_eta_max},
111      'benchmark':bench,'oscillation_comparison':osc}
112    with open('results.json','w') as f: json.dump(result,f,indent=2)
113    print(json.dumps(result,indent=2))
114
115if __name__=='__main__': main()