summaryrefslogtreecommitdiff
path: root/ep_run/eig_traj3.py
diff options
context:
space:
mode:
Diffstat (limited to 'ep_run/eig_traj3.py')
-rw-r--r--ep_run/eig_traj3.py39
1 files changed, 39 insertions, 0 deletions
diff --git a/ep_run/eig_traj3.py b/ep_run/eig_traj3.py
new file mode 100644
index 0000000..1d864b6
--- /dev/null
+++ b/ep_run/eig_traj3.py
@@ -0,0 +1,39 @@
+"""E1b: ARPACK gold-standard re-measurement of the redx trajectory at 6 key ckpts (lead_rho's cold 2-D
+iteration under-reads near-unity clusters by ~0.02-0.03, so absolute mu from eig_traj.py is suspect).
+Top-3 |lam| of the forward map M = I + eps*J_F at the 400-step deep state, same seed-42 batch."""
+import numpy as np, torch, scipy.sparse.linalg as sla
+from torch.autograd.functional import jvp
+import lt_ep_train as L
+
+EPS, B, C = 0.1, 6, 1.0
+KEY = [2100, 2200, 2300]
+
+for s in KEY:
+ torch.manual_seed(0)
+ blk = L.EQBlock(512, 16, 256, 256, c=C, attn_mode='thick'); blk.qknorm = True
+ ck = torch.load(f'runs/redx_traj/s{s}.pt', map_location=L.dev)
+ with torch.no_grad():
+ for p, w in zip(blk.allp, ck['allp']):
+ p.copy_(w.to(L.dev))
+ torch.manual_seed(42)
+ idx, _ = L.get_batch('train', B, 256)
+ xin = blk.embed(idx).detach()
+ z = L.relax(blk, xin.clone(), xin, 400, EPS)
+ sh, n = z.shape, z.numel()
+ kk = 1.0 - EPS * (1.0 + C)
+
+ def mv(x, z=z, sh=sh):
+ v = torch.from_numpy(np.asarray(x, dtype=np.float32)).to(L.dev).view(sh)
+ with torch.no_grad():
+ Mv = kk * v + EPS * jvp(blk.nc_force, z, v)[1]
+ return Mv.reshape(-1).double().cpu().numpy()
+
+ A = sla.LinearOperator((n, n), matvec=mv, dtype=np.float64)
+ try:
+ vals = sorted(sla.eigs(A, k=3, which='LM', return_eigenvectors=False, maxiter=2000, tol=1e-4),
+ key=lambda x: -abs(x))
+ out = " ".join(f"|l|={abs(l):.5f}(mu={(l.real-1)/EPS:+.4f}{l.imag/EPS:+.3f}j)" for l in vals)
+ except Exception as e:
+ out = f"ARPACK-fail {type(e).__name__}"
+ print(f"s{s:<5} {out}", flush=True)
+print("EIG_TRAJ2_DONE", flush=True)