1import json, math, time
  2from dataclasses import dataclass
  3import numpy as np
  4
  5@dataclass
  6class Controller:
  7    eta0: float
  8    B0: int
  9    rho0: float
 10    kind: str = "polyak"
 11    safety: float = 0.5
 12    beta: float = 1.5
 13
 14    def eta(self, B, rho):
 15        if self.kind == "nesterov":
 16            scale = self.safety * (B / self.B0) ** self.beta * (1-rho)/(1-self.rho0)
 17        else:
 18            scale = self.safety * (B / self.B0) * (1-rho)/(1-self.rho0)
 19        return self.eta0 * min(1.0, scale)
 20
 21def spectral_radius(kind, eta, rho, lam=1.0):
 22    if kind == "polyak":
 23        A = np.array([[1-eta*lam, -eta*rho], [lam, rho]], float)
 24    elif kind == "nesterov":
 25        A = np.array([[1-eta*lam, -eta*rho*(1-eta*lam)],
 26                      [lam, rho*(1-eta*lam)]], float)
 27    else:
 28        return abs(1-eta*lam)
 29    return float(np.max(np.abs(np.linalg.eigvals(A))))
 30
 31def critical_eta(kind, rho, lam=1.0):
 32    lo, hi = 0.0, 1.0 / lam
 33    while spectral_radius(kind, hi, rho, lam) < 1.0:
 34        hi *= 2
 35    for _ in range(70):
 36        mid = (lo + hi) / 2
 37        if spectral_radius(kind, mid, rho, lam) < 1.0:
 38            lo = mid
 39        else:
 40            hi = mid
 41    return lo
 42
 43def stochastic_run(kind, B, eta, rho, seed, steps=400, d=8, noise=1.0):
 44    rng = np.random.default_rng(seed)
 45    x = rng.normal(size=d)
 46    v = np.zeros(d)
 47    losses = []
 48    updates = []
 49    for _ in range(steps):
 50        if kind == "nesterov":
 51            grad = x - eta * rho * v + rng.normal(size=d) * noise / math.sqrt(B)
 52        else:
 53            grad = x + rng.normal(size=d) * noise / math.sqrt(B)
 54        old = x.copy()
 55        v = rho * v + grad
 56        x = x - eta * v
 57        losses.append(0.5 * float(np.dot(x, x)))
 58        updates.append(float(np.linalg.norm(x-old)))
 59        if not np.isfinite(losses[-1]) or losses[-1] > 1e12:
 60            return {"stable": False, "final_loss": float("inf"), "median_loss": float("inf"),
 61                    "update": float("inf"), "steps": len(losses)}
 62    return {"stable": bool(np.isfinite(losses).all()),
 63            "final_loss": losses[-1], "median_loss": float(np.median(losses[-50:])),
 64            "update": float(np.median(updates[-50:])), "steps": steps}
 65
 66def scan_stability(kind, B, rho, eta_grid, repeats=3):
 67    out=[]
 68    for eta in eta_grid:
 69        runs=[stochastic_run(kind,B,eta,rho,100+r,steps=300) for r in range(repeats)]
 70        out.append((eta, all(r["stable"] for r in runs)))
 71    stable=[e for e, ok in out if ok]
 72    return max(stable) if stable else 0.0
 73
 74def main():
 75    rho0=0.9; B0=1; eta0=0.05; beta=1.5
 76    # Core math: compare critical-rate ratios to the predicted capped scaling.
 77    rho_values=[0.0,0.5,0.9,0.95]
 78    math_rows=[]
 79    for kind in ["polyak","nesterov"]:
 80        base=critical_eta(kind,rho0)
 81        for rho in rho_values:
 82            c=critical_eta(kind,rho)
 83            predicted=min(1.0, (1-rho)/(1-rho0))
 84            math_rows.append({"kind":kind,"rho":rho,"critical":c,
 85                              "ratio":c/base,"predicted_uncapped_ratio":predicted})
 86    # Batch controller experiment with fixed number of samples (steps shrink as B grows).
 87    rows=[]
 88    for kind in ["polyak","nesterov"]:
 89        ctrl=Controller(eta0,B0,rho0,kind,0.5,beta)
 90        for B in [1,2,4,8,16]:
 91            eta=ctrl.eta(B,rho0)
 92            fixed=stochastic_run(kind,B,eta0,rho0,123,steps=max(20,800//B))
 93            controlled=stochastic_run(kind,B,eta,rho0,123,steps=max(20,800//B))
 94            rows.append({"kind":kind,"B":B,"eta_controller":eta,
 95                         "fixed_stable":fixed["stable"],"controller_stable":controlled["stable"],
 96                         "fixed_final_loss":fixed["final_loss"],"controller_final_loss":controlled["final_loss"],
 97                         "controller_update":controlled["update"]})
 98    # Empirical maximum stable rates at each batch, plus theoretical deterministic boundary.
 99    scan=[]
100    for kind in ["polyak","nesterov"]:
101        for B in [1,2,4,8,16]:
102            grid=np.geomspace(0.002,2.0,45)
103            empirical=scan_stability(kind,B,rho0,grid)
104            scan.append({"kind":kind,"B":B,"empirical_max_eta":empirical,
105                         "deterministic_quadratic_eta":critical_eta(kind,rho0)})
106    result={"math_check":math_rows,"controller_runs":rows,"stability_scan":scan,
107            "config":{"eta0":eta0,"rho0":rho0,"beta":beta,"safety":0.5}}
108    with open("results.json","w") as f: json.dump(result,f,indent=2)
109    print(json.dumps(result,indent=2))
110
111if __name__ == "__main__":
112    main()