1import json
 2import random
 3from pathlib import Path
 4import numpy as np
 5
 6
 7def cubic_check(r=0.2, L=1.0, omega=1.0, rho=0.8):
 8    a1 = 2*r/L
 9    a2 = (r/L)**2 + omega**2
10    kappa = 1.5*omega/L
11    boundary = a1*a2/kappa
12    rows = []
13    for frac in [0.2, 0.8, 0.99, 1.0, 1.01, 1.5]:
14        g = frac*boundary
15        roots = np.roots([1., a1, a2, kappa*g])
16        rows.append({'fraction': frac, 'chi': float(kappa*g/(a1*a2)),
17                     'max_real_root': float(np.max(roots.real)),
18                     'roots': [[float(z.real), float(z.imag)] for z in roots]})
19    return {'a1': a1, 'a2': a2, 'kappa': kappa,
20            'boundary_g': boundary, 'cap_g': rho*boundary, 'rows': rows,
21            'crossing_observed': bool(rows[2]['max_real_root'] < 0 and rows[4]['max_real_root'] > 0)}
22
23
24def run_momentum(lr, momentum, H, steps=250, cap=False, rho=0.8):
25    # Two-dimensional quadratic, with one deliberately high-curvature direction.
26    x = np.array([1., 1.])
27    v = np.zeros(2)
28    losses, gains, chis = [], [], []
29    # Use the cubic coefficients as the specified local stability model.
30    r, L, omega = .2, 1., 1.
31    a1, a2, kappa = 2*r/L, (r/L)**2 + omega**2, 1.5*omega/L
32    gmax = rho*a1*a2/kappa
33    effective_lr = min(lr, gmax/max(H)) if cap else lr
34    for _ in range(steps):
35        grad = H*x
36        g = effective_lr*max(H)
37        chi = kappa*g/(a1*a2)
38        v = momentum*v + grad
39        x = x - effective_lr*v
40        losses.append(0.5*float(x@(H*x))); gains.append(g); chis.append(chi)
41        if not np.all(np.isfinite(x)) or losses[-1] > 1e100:
42            break
43    tail = losses[-50:]
44    return {'final_loss': losses[-1] if losses else float('inf'),
45            'min_loss': min(losses) if losses else float('inf'),
46            'loss_std_tail': float(np.std(tail)) if tail else float('inf'),
47            'steps': len(losses), 'effective_lr': effective_lr,
48            'max_chi': max(chis) if chis else float('inf'),
49            'diverged': len(losses) < steps or not np.isfinite(losses[-1])}
50
51
52def main():
53    random.seed(7); np.random.seed(7)
54    check = cubic_check()
55    # Sweep a range around the ordinary quadratic momentum stability boundary.
56    H = np.array([1., 20.])
57    settings = []
58    for lr in [.12, .2, .4, .8, 1.2]:
59        settings.append({'lr': lr,
60                         'baseline': run_momentum(lr, .9, H, cap=False),
61                         'routh_cap': run_momentum(lr, .9, H, cap=True)})
62    out = {'cubic_check': check, 'quadratic_experiment': settings}
63    Path('results.json').write_text(json.dumps(out, indent=2))
64    print(json.dumps(out, indent=2))
65
66if __name__ == '__main__':
67    main()