1 files changed, 67 insertions, 0 deletions
diff --git a/ep_run/anderson_control.py b/ep_run/anderson_control.py
new file mode 100644
index 0000000..10147c3
--- /dev/null
+++ b/ep_run/anderson_control.py
@@ -0,0 +1,67 @@
+import torch, pickle, math, numpy as np
+from pathlib import Path
+from scipy.sparse.linalg import LinearOperator, eigs
+import lt_ep_train as L
+from lt_ep_train import EQBlock
+L.DD=Path('data/tinystories_bpe'); L.vocab=pickle.load(open(L.DD/'meta.pkl','rb'))['vocab_size']
+dev='cuda'; B=2; T=256; eps=0.1
+torch.manual_seed(1234); idx,y=L.get_batch('val',B,T); idx=idx.to(dev) if hasattr(idx,'to') else idx
+def load_blk(path):
+    ck=torch.load(path,map_location=dev)
+    blk=EQBlock(512,16,256,256,s=1.0,c=1.0,attn_mode='thick'); blk.qknorm=True
+    with torch.no_grad():
+        for p,w in zip(blk.allp,ck['allp']): p.copy_(w.to(dev))
+    return blk, ck.get('best','?')
+def gmap(blk,xin,z):
+    with torch.no_grad(): return z+eps*blk.force(z,xin).detach()
+def plain(blk,xin,z0,steps=300):
+    z=z0.clone()
+    for _ in range(steps): z=gmap(blk,xin,z)
+    return ((gmap(blk,xin,z)-z).norm()/(z.norm()+1e-9)).item()
+def anderson(blk,xin,z0,m=6,max_iter=400,tol=1e-7,lam=1e-4):
+    Bs,d=z0.shape[0],z0[0].numel()
+    X=torch.zeros(Bs,m,d,device=dev); Fb=torch.zeros(Bs,m,d,device=dev)
+    X[:,0]=z0.reshape(Bs,d); Fb[:,0]=gmap(blk,xin,z0).reshape(Bs,d)
+    X[:,1]=Fb[:,0]; Fb[:,1]=gmap(blk,xin,X[:,1].view_as(z0)).reshape(Bs,d)
+    Hm=torch.zeros(Bs,m+1,m+1,device=dev); Hm[:,0,1:]=1; Hm[:,1:,0]=1
+    yv=torch.zeros(Bs,m+1,1,device=dev); yv[:,0]=1
+    r=1.0; best_r=9.0; z_best=z0.clone()
+    for k in range(2,max_iter):
+        n=min(k,m); Gm=Fb[:,:n]-X[:,:n]
+        Hm[:,1:n+1,1:n+1]=torch.bmm(Gm,Gm.transpose(1,2))+lam*torch.eye(n,device=dev)[None]
+        alpha=torch.linalg.solve(Hm[:,:n+1,:n+1],yv[:,:n+1])[:,1:n+1,0]
+        X[:,k%m]=torch.bmm(alpha[:,None],Fb[:,:n])[:,0]
+        Fb[:,k%m]=gmap(blk,xin,X[:,k%m].view_as(z0)).reshape(Bs,d)
+        r=((Fb[:,k%m]-X[:,k%m]).norm()/(Fb[:,k%m].norm()+1e-9)).item()
+        if r<best_r: best_r=r; z_best=X[:,k%m].view_as(z0).clone()
+        if r<tol or not math.isfinite(r): break
+    return best_r,k+1,z_best
+def eig_at(blk,xin,z,k=6,hrel=1e-3):
+    z=z.detach(); F0=blk.force(z,xin).detach(); zn=z.norm().item(); h=hrel*zn; shp=z.shape; N=z.numel()
+    def Jv(vt):
+        nv=vt.norm().item()
+        if nv<1e-20: return torch.zeros_like(vt)
+        return (blk.force(z+h*(vt/nv),xin).detach()-F0)/h*nv
+    def matvec(v):
+        vt=torch.from_numpy(np.ascontiguousarray(v).astype('float32')).reshape(shp).to(dev)
+        return (vt+eps*Jv(vt)).double().cpu().numpy().reshape(-1)
+    op=LinearOperator((N,N),matvec=matvec,dtype='float64')
+    vals=eigs(op,k=k,which='LM',return_eigenvectors=False,maxiter=4000,tol=1e-5)
+    return F0.norm().item()/(zn+1e-9), sorted(vals,key=lambda x:-abs(x))
+for tag,path in [("s2000 (healthy)","runs/redx_traj/s2000.pt"),
+                 ("s3200 (blew@2.74)","runs/redx_traj/s3200.pt")]:
+    blk,best=load_blk(path); xin=blk.embed(idx).detach()
+    with torch.no_grad():
+        pr=plain(blk,xin,xin.clone(),300)
+        ar,ak,zst=anderson(blk,xin,xin.clone(),max_iter=400)
+    print(f"=== {tag} best={best} ===")
+    print(f"  plain relax(300) res={pr:.3e} | Anderson best_res={ar:.3e} in {ak} iters")
+    if ar<2e-3:
+        g,vals=eig_at(blk,xin,zst)
+        print(f"  -> Anderson FOUND a root (force g={g:.4f}); eigenvalues of M=I+eps*J at the root:")
+        for lam in vals[:4]:
+            mu=(lam-1)/eps
+            print(f"     |lam|={abs(lam):.5f} mu={mu.real:+.4f}{mu.imag:+.4f}j [{'UNSTABLE' if abs(lam)>1+1e-4 else 'STABLE'}{' rot' if abs(lam.imag)>1e-3 else ' real'}]")
+    else:
+        print(f"  -> Anderson did NOT converge to a root (best_res={ar:.3e}) => no reachable fixed point")
+print("=== DONE === key: s3200 root + ReMu<0 = Euler-artifact(integration fixes it); root + ReMu>0 = unstable fixed pt(true instab); no root = true instab")