1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4
  5SEED=7
  6np.random.seed(SEED); random.seed(SEED)
  7
  8# ---------- Stage 1: transfer-function sanity check ----------
  9def lockin_demo():
 10    # Stable diagonal A, driven by sinusoidal input. Exact H is known.
 11    lambdas=np.array([0.7, 2.0, 5.0])
 12    B=np.array([1.0, .4, .2]); C=np.array([.8, -.3, .5])
 13    w=1.3; eps=.03; dt=.01; periods=25
 14    T=2*np.pi/w*periods; n=int(T/dt); t=np.arange(n)*dt
 15    x=np.zeros(3); ys=[]
 16    for ti in t:
 17        u=eps*np.sin(w*ti)
 18        x += dt*(-lambdas*x+B*u)
 19        ys.append(C@x)
 20    # discard first periods as transient and correlate on integer-period tail
 21    cut=n//5; tt=t[cut:]; yy=np.asarray(ys[cut:])
 22    # For y=Re(H eps exp(i wt)), 2 mean(y exp(-iwt))/eps estimates H.
 23    Hhat=1j*(2*np.mean(yy*np.exp(-1j*w*tt))/eps)
 24    Hexact=np.sum(C*B/(lambdas+1j*w))
 25    rel=abs(Hhat-Hexact)/abs(Hexact)
 26    # discrete-gradient stability transition is independently checked
 27    eta_grid=np.linspace(.02, .55, 107)
 28    stable=[abs(1-eta*max(lambdas))<1 for eta in eta_grid]
 29    empirical=max(eta for eta,s in zip(eta_grid,stable) if s)
 30    boundary=2/max(lambdas)
 31    # Direct quadratic mode test: z_{t+1}=(1-eta*lambda)z_t, with lambda_max=5.
 32    # Below 2/lambda_max it decays; above it grows (alternating when eta*lambda>1).
 33    quad=[]
 34    for eta in [0.30, 0.39, 0.41, 0.50]:
 35        z=np.array([1.0, 0.3, -0.2]); norms=[]
 36        for _ in range(80):
 37            z=(1-eta*lambdas)*z; norms.append(float(np.linalg.norm(z)))
 38        quad.append({'eta':eta,'final_norm':norms[-1], 'decays':bool(norms[-1]<norms[0])})
 39    return {'H_est': [float(Hhat.real),float(Hhat.imag)], 'H_exact':[float(Hexact.real),float(Hexact.imag)],
 40            'relative_transfer_error':float(rel), 'theory_eta_critical':boundary,
 41            'empirical_eta_critical_grid':float(empirical), 'quadratic_mode_sweep':quad, 'stability_claim_observed':bool(empirical<boundary and abs(empirical-boundary)<.02),
 42            'phase_convention': 'corrected sine-input lock-in by multiplying raw estimate by i'}
 43
 44# ---------- Stage 2: tiny real training comparison ----------
 45def run_training():
 46    import torch
 47    from sklearn.datasets import load_digits
 48    from sklearn.model_selection import train_test_split
 49    from sklearn.preprocessing import StandardScaler
 50    torch.manual_seed(SEED); np.random.seed(SEED)
 51    X,y=load_digits(return_X_y=True)
 52    X=StandardScaler().fit_transform(X).astype('float32')
 53    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
 54    dev='cuda' if torch.cuda.is_available() else 'cpu'
 55    try:
 56        device=torch.device(dev)
 57        if dev=='cuda': torch.zeros(1,device=device)
 58    except Exception:
 59        device=torch.device('cpu')
 60    xt=torch.tensor(Xtr,device=device); yt=torch.tensor(ytr,device=device)
 61    xev=torch.tensor(Xte,device=device); yev=torch.tensor(yte,device=device)
 62    def make():
 63        torch.manual_seed(SEED)
 64        return torch.nn.Sequential(torch.nn.Linear(64,48),torch.nn.Tanh(),torch.nn.Linear(48,10)).to(device)
 65    def loss_acc(model, x, y):
 66        with torch.no_grad():
 67            z=model(x); return float(torch.nn.functional.cross_entropy(z,y)), float((z.argmax(1)==y).float().mean())
 68    def hessian_top(model, xb, yb, iters=4):
 69        # Hessian-vector power iteration: an inexpensive empirical local curvature estimate.
 70        params=[p for p in model.parameters() if p.requires_grad]
 71        out=model(xb); l=torch.nn.functional.cross_entropy(out,yb)
 72        g=torch.autograd.grad(l,params,create_graph=True)
 73        v=[torch.randn_like(p) for p in params]; norm=torch.sqrt(sum((q*q).sum() for q in v)); v=[q/norm for q in v]
 74        val=0.
 75        for _ in range(iters):
 76            gv=sum((a*b).sum() for a,b in zip(g,v))
 77            hv=torch.autograd.grad(gv,params,retain_graph=True)
 78            norm=torch.sqrt(sum((q*q).sum() for q in hv)).detach().clamp_min(1e-8)
 79            v=[q/norm for q in hv]
 80            val=float(sum((a*b).sum() for a,b in zip(v,hv)).detach())
 81        return max(val,1e-3)
 82    results={}
 83    steps=700; bs=96; K=35
 84    for mode in ['fixed','cosine','impedance']:
 85        model=make(); opt=torch.optim.SGD(model.parameters(),lr=.08)
 86        eta=.08; max_loss=0.; t0=time.time(); hist=[]
 87        for step in range(steps):
 88            idx=torch.randint(0,len(xt),(bs,),device=device); xb,yb=xt[idx],yt[idx]
 89            opt.zero_grad(set_to_none=True); l=torch.nn.functional.cross_entropy(model(xb),yb); l.backward(); opt.step()
 90            if mode=='cosine':
 91                eta=.08*.5*(1+math.cos(math.pi*step/steps));
 92                for g in opt.param_groups:g['lr']=eta
 93            elif mode=='impedance' and step%K==0:
 94                # Probe/estimate the local dominant relaxation pole (curvature lambda_max),
 95                # then apply eta=gamma*(2/lambda_max), clipped to avoid abrupt jumps.
 96                lam=hessian_top(model,xt[:bs],yt[:bs]); target=.7*2/lam
 97                eta=float(np.clip(min(eta*1.1,target),.005,.35))
 98                for g in opt.param_groups:g['lr']=eta
 99            if step%K==0:
100                q,_=loss_acc(model,xt[:min(256,len(xt))],yt[:min(256,len(yt))]); max_loss=max(max_loss,q); hist.append((step,eta,q))
101        tr,ta=loss_acc(model,xt,yt); te,ea=loss_acc(model,xev,yev)
102        results[mode]={'train_loss':tr,'test_loss':te,'test_accuracy':ea,'max_sampled_loss':max_loss,'final_lr':eta,'history':hist,'seconds':time.time()-t0}
103    # fixed-rate sweep exposes the alternating/divergent side of the claimed boundary.
104    sweep=[]
105    for eta0 in [.03,.08,.16,.30,.45,.60]:
106        m=make(); o=torch.optim.SGD(m.parameters(),lr=eta0); losses=[]
107        for step in range(250):
108            idx=torch.randint(0,len(xt),(bs,),device=device); o.zero_grad(set_to_none=True)
109            ll=torch.nn.functional.cross_entropy(m(xt[idx]),yt[idx]); ll.backward(); o.step()
110            if step%25==0: losses.append(loss_acc(m,xt[:256],yt[:256])[0])
111        sweep.append({'eta':eta0,'losses':losses,'finite':bool(np.isfinite(losses).all()),'final_loss':losses[-1]})
112    return {'device':str(device),'comparisons':results,'fixed_rate_sweep':sweep}
113
114if __name__=='__main__':
115    out={'seed':SEED,'lockin':lockin_demo()}
116    try: out['training']=run_training()
117    except Exception as e: out['training_error']=repr(e)
118    Path('results.json').write_text(json.dumps(out,indent=2))
119    print(json.dumps(out,indent=2))