1"""Cycle-aware heavy-ball safeguard: math check and tiny reproducible experiments."""
 2import json, math, random
 3import numpy as np
 4
 5
 6def spectral_radius(lam, lr, beta):
 7    A = np.array([[1 + beta - lr * lam, -beta], [1., 0.]])
 8    return float(np.max(np.abs(np.linalg.eigvals(A))))
 9
10
11def diagnostics(d_t, d_prev, d_lag2, eps=1e-12):
12    a, b, c = (np.asarray(x, float) for x in (d_t, d_prev, d_lag2))
13    corr = float(np.dot(a, c) / (np.linalg.norm(a)*np.linalg.norm(c) + eps))
14    ratio = float(np.linalg.norm(a) / (np.linalg.norm(b) + eps))
15    return corr, ratio
16
17
18class CycleAwareHB:
19    """Numpy HB optimizer with fixed random projection and periodic safeguard."""
20    def __init__(self, shape, lr, beta, projection_dim=32, interval=20, hold=100, seed=0):
21        self.lr0, self.beta0 = lr, beta
22        self.lr, self.beta = lr, beta
23        self.interval, self.hold = interval, hold
24        self.rng = np.random.default_rng(seed)
25        n = int(np.prod(shape))
26        self.P = self.rng.normal(0, 1/math.sqrt(projection_dim), (projection_dim, n))
27        self.v = np.zeros(n); self.prev_d = None; self.prev_prev_d = None
28        self.cooldown = 0; self.triggers = 0; self.history = []
29
30    def step(self, x, grad, curvature):
31        old = x.copy(); self.v = self.beta*self.v + grad
32        x = x - self.lr*self.v
33        d = self.P @ (x-old)
34        # Preserve d_{t-1} and d_{t-2} before installing current d_t.
35        d_prev, d_lag2 = self.prev_d, self.prev_prev_d
36        self.prev_prev_d, self.prev_d = d_prev, d
37        step_no = len(self.history) + 1
38        trig = False; c = float('nan'); r = float('nan'); rho = spectral_radius(curvature, self.lr, self.beta)
39        if step_no % self.interval == 0 and d_lag2 is not None and d_prev is not None:
40            c, r = diagnostics(d, d_prev, d_lag2)
41            trig = c > .8 and abs(r-1) < .15 and rho >= .98
42            if trig:
43                self.beta = min(.5, .7*self.beta); self.lr *= .8
44                self.v.fill(0); self.cooldown = self.hold; self.triggers += 1
45        if self.cooldown > 0: self.cooldown -= 1
46        elif not trig and (step_no % self.interval == 0) and spectral_radius(curvature, self.lr0, self.beta0) < .95:
47            # Never recover into the originally diagnosed unstable regime.
48            self.lr += .05*(self.lr0-self.lr); self.beta += .05*(self.beta0-self.beta)
49        self.history.append((c,r,rho,trig,self.lr,self.beta))
50        return x
51
52
53def math_check():
54    # For constant curvature the measured linear transition must equal the formula.
55    lr, beta, lam = .039, .95, 100.
56    A = np.array([[1+beta-lr*lam, -beta], [1., 0.]])
57    z = np.array([.37, -.21]); observed = A @ z
58    # direct recurrence gives the same next state
59    x, old = z
60    direct = np.array([(1+beta)*x - beta*old - lr*lam*x, x])
61    cyc = np.array([[1.,0.],[-.5,.8660254],[-.5,-.8660254], [1.,0.]])
62    cycle_corr, cycle_ratio = diagnostics(cyc[3]-cyc[2], cyc[2]-cyc[1], cyc[1]-cyc[0])
63    two_corr, two_ratio = diagnostics(np.array([1.,0.]), np.array([-1.,0.]), np.array([1.,0.]))
64    return {"matrix_error": float(np.max(np.abs(observed-direct))),
65            "period3_lag2_corr": cycle_corr, "period3_ratio": cycle_ratio,
66            "period2_lag2_corr": two_corr, "period2_ratio": two_ratio,
67            "rho_near_unit": spectral_radius(lam,lr,beta),
68            "rho_after_reduction": spectral_radius(lam,.8*lr,.7*beta)}
69
70
71def quadratic_run(guarded, seed=7, steps=400):
72    rng=np.random.default_rng(seed); n=24
73    eig=np.geomspace(1,100,n); x=rng.normal(size=n); loss=[]
74    if guarded: opt=CycleAwareHB((n,),.039,.95,seed=seed)
75    v=np.zeros(n)
76    for t in range(steps):
77        g=eig*x
78        if guarded: x=opt.step(x,g,float(eig.max()))
79        else:
80            v=.95*v+g; x=x-.039*v
81        loss.append(float(.5*np.dot(eig*x,x)))
82        if not np.isfinite(loss[-1]): break
83    return {"final_loss":loss[-1], "best_loss":min(loss), "triggers":opt.triggers if guarded else 0,
84            "loss_at_100":loss[min(99,len(loss)-1)], "loss_at_400":loss[-1]}
85
86
87def main():
88    np.random.seed(0); random.seed(0)
89    out={"math":math_check(), "baseline":quadratic_run(False), "safeguard":quadratic_run(True)}
90    print(json.dumps(out, indent=2))
91
92if __name__ == '__main__': main()