1import json, math, random
 2from pathlib import Path
 3import numpy as np
 4
 5SEED = 2313
 6
 7def set_seed(seed=SEED):
 8    random.seed(seed); np.random.seed(seed)
 9
10def metrics(x, dt=1.0):
11    x=np.asarray(x,float); z=x-x.mean(); v=float(np.mean(z*z))
12    rho=float(np.sum(z[1:]*z[:-1])/max(np.sum(z*z),1e-15))
13    return v,rho,float(-math.log(max(rho,1e-3))/dt)
14
15def ar1_sweep():
16    rng=np.random.default_rng(SEED); a_values=np.array([.2,.4,.6,.75,.85,.92,.96,.98])
17    sigma=.5; n=120000; burn=1000; rows=[]
18    for a in a_values:
19        e=rng.normal(0,sigma,n+burn); x=np.zeros(n+burn)
20        for i in range(1,len(x)): x[i]=a*x[i-1]+e[i]
21        v,rho,r=metrics(x[burn:]); rows.append({'a':float(a),'rho_hat':rho,'rho_pred':float(a),
22          'var_hat':v,'var_pred':float(sigma*sigma/(1-a*a)),'recovery_hat':r,
23          'recovery_pred':float(-math.log(a))})
24    return {'rows':rows,'rho_rmse':float(np.sqrt(np.mean([(q['rho_hat']-q['rho_pred'])**2 for q in rows]))),
25      'variance_mean_relative_error':float(np.mean([abs(q['var_hat']-q['var_pred'])/q['var_pred'] for q in rows])),
26      'recovery_rmse':float(np.sqrt(np.mean([(q['recovery_hat']-q['recovery_pred'])**2 for q in rows])))}
27
28def delayed_quadratic(scheduler, seed=SEED, steps=3000, delay=8, lr=.035, momentum=.85):
29    A=np.diag([1.,10.8]); x=np.array([2.,2.]); velocity=np.zeros(2); hist=[]
30    lr_now=lr; beta_now=momentum; triggered=False; trigger_step=None; ema=None; vals=[]
31    base_stats=None; cooldown=0; losses=[]
32    for t in range(steps):
33        loss=.5*float(x@A@x); g=A@x
34        stale=g if delay==0 or t<delay else hist[t-delay]
35        velocity=beta_now*velocity+stale; x=x-lr_now*velocity; hist.append(g.copy()); losses.append(loss)
36        # Fixed random parameter projection is an alternative scalar observable.
37        obs=float(np.log1p(loss)); ema=obs if ema is None else .97*ema+.03*obs; vals.append(obs-ema)
38        if len(vals)>=60 and t%10==0:
39            v,rho,rate=metrics(vals[-60:])
40            if base_stats is None and t==59: base_stats=(max(v,1e-10),rho,rate)
41            if scheduler and not triggered and base_stats is not None and t>=70:
42                bv,br,bc=base_stats
43                if rho>br+.10 or v>2*bv or rate<.65*bc:
44                    lr_now*=.5; beta_now=min(beta_now,.90); triggered=True; trigger_step=t
45            if triggered and cooldown==0 and rate>.9*base_stats[2]:
46                lr_now=lr; beta_now=momentum; cooldown=100
47        cooldown=max(0,cooldown-1)
48        if not np.all(np.isfinite(x)) or np.linalg.norm(x)>1e8:
49            return {'diverged':True,'divergence_step':t,'final_loss':'inf','trigger_step':trigger_step,'min_loss':float(np.min(losses))}
50    return {'diverged':False,'divergence_step':None,'final_loss':float(losses[-1]),'trigger_step':trigger_step,'min_loss':float(np.min(losses))}
51
52def training_sweep():
53    rows=[]
54    for d in [0,2,4,6,8,10,12]:
55        rows.append({'delay':d,'baseline':delayed_quadratic(False,delay=d),'scheduler':delayed_quadratic(True,delay=d)})
56    return rows
57
58def main():
59    set_seed(); out={'seed':SEED,'ar1':ar1_sweep(),'delayed_training':training_sweep()}
60    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
61if __name__=='__main__': main()