1 files changed, 47 insertions, 0 deletions
diff --git a/ep_run/dip_screen.py b/ep_run/dip_screen.py
new file mode 100644
index 0000000..5af4c14
--- /dev/null
+++ b/ep_run/dip_screen.py
@@ -0,0 +1,47 @@
+"""ARPACK dip-screening of a dip-farm trajectory (the governor's dataset): for every frozen snapshot
+>= --from_step, top-3 leading eigenvalues of the forward map at the 400-deep state. Output: one line
+per snapshot -> dips (all |lam|<1) marked DIP. Usage: dip_screen.py --seed N [--from_step 800]"""
+import argparse, numpy as np, torch, scipy.sparse.linalg as sla
+from pathlib import Path
+from torch.autograd.functional import jvp
+import lt_ep_train as L
+
+ap = argparse.ArgumentParser()
+ap.add_argument('--seed', type=int, required=True)
+ap.add_argument('--from_step', type=int, default=800)
+cfg = ap.parse_args()
+EPS, B, C = 0.1, 6, 1.0
+traj = Path(f'runs/dipfarm_s{cfg.seed}_traj')
+steps = sorted(int(p.stem[1:]) for p in traj.glob('s*.pt') if int(p.stem[1:]) >= cfg.from_step)
+print(f"[seed {cfg.seed}] screening {len(steps)} snapshots: {steps}", flush=True)
+
+for s in steps:
+    torch.manual_seed(0)
+    blk = L.EQBlock(512, 16, 256, 256, c=C, attn_mode='thick'); blk.qknorm = True
+    ck = torch.load(traj / f'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=1500, tol=2e-4),
+                      key=lambda x: -abs(x))
+        stable = all(abs(l) < 1.0 for l in vals)
+        out = "  ".join(f"|l|={abs(l):.5f}" for l in vals)
+        print(f"s{cfg.seed}/{s:<5} {out}  {'DIP' if stable else ''}", flush=True)
+    except Exception as e:
+        print(f"s{cfg.seed}/{s:<5} ARPACK-fail {type(e).__name__}", flush=True)
+print(f"DIP_SCREEN_S{cfg.seed}_DONE", flush=True)