Create eigen_barrage_parallel_refinement_testing.py

eigen_barrage_parallel_refinement_testing.py

@@ -0,0 +1,174 @@
+"""
+Diagnostic: what exactly breaks in parallel root-finding?
+
+Test 1: Pure parallel Laguerre (no Aberth, no clamp, no damp)
+Test 2: Parallel Laguerre + Aberth
+Test 3: Sequential Laguerre + deflation (baseline)
+
+Prints per-iteration convergence to identify exactly where it goes wrong.
+"""
+import math, torch
+
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.set_float32_matmul_precision('highest')
+
+dev = torch.device('cuda')
+B = 512; N = 6
+torch.manual_seed(42)
+A = (lambda R: (R+R.mT)/2)(torch.randn(B, N, N, device=dev))
+rv, rV = torch.linalg.eigh(A)
+
+# FL Phase 1 — get characteristic polynomial
+sc = (torch.linalg.norm(A.reshape(B,-1), dim=-1) / math.sqrt(N)).clamp(min=1e-12)
+As = A / sc[:, None, None]; Ad = As.double()
+I_d = torch.eye(N, device=dev, dtype=torch.float64).unsqueeze(0).expand(B,-1,-1)
+c = torch.zeros(B, N+1, device=dev, dtype=torch.float64); c[:, N] = 1.0
+Mk = torch.zeros(B, N, N, device=dev, dtype=torch.float64)
+for k in range(1, N+1):
+    Mk = torch.bmm(Ad, Mk) + c[:, N-k+1, None, None] * I_d
+    c[:, N-k] = -(Ad * Mk).sum((-2,-1)) / k
+
+# True roots (scaled)
+true_roots = (rv / sc.unsqueeze(-1)).double().sort(dim=-1).values
+
+# Init from diagonal
+z_init = Ad.diagonal(dim1=-2, dim2=-1).sort(dim=-1).values
+pert = torch.linspace(-1e-3, 1e-3, N, device=dev, dtype=torch.float64).unsqueeze(0)
+z_init = z_init + pert
+
+def horner_pd(c, z):
+    """Evaluate p(z), p'(z), p''(z)/2 via Horner. c: [B,n+1], z: [B,n]"""
+    B, n_roots = z.shape
+    n = c.shape[1] - 1
+    pv = c[:, n:n+1].expand(B, n_roots)
+    dp = torch.zeros_like(pv)
+    d2 = torch.zeros_like(pv)
+    for j in range(n-1, -1, -1):
+        d2 = d2 * z + dp
+        dp = dp * z + pv
+        pv = pv * z + c[:, j:j+1]
+    return pv, dp, d2
+
+def laguerre_step(c, z, n):
+    pv, dp, d2 = horner_pd(c, z)
+    ok = pv.abs() > 1e-30
+    ps = torch.where(ok, pv, torch.ones_like(pv))
+    G = torch.where(ok, dp / ps, torch.zeros_like(dp))
+    H = G * G - torch.where(ok, 2.0 * d2 / ps, torch.zeros_like(d2))
+    disc = ((n-1.0) * (n * H - G * G)).clamp(min=0.0)
+    sq = torch.sqrt(disc)
+    gp = G + sq; gm = G - sq
+    den = torch.where(gp.abs() >= gm.abs(), gp, gm)
+    return torch.where(den.abs() > 1e-20, float(n) / den, torch.zeros_like(den))
+
+mask_eye = torch.eye(N, device=dev, dtype=torch.bool).unsqueeze(0)
+
+def aberth_correction(z):
+    diffs = z.unsqueeze(-1) - z.unsqueeze(-2)
+    diffs_safe = diffs.masked_fill(mask_eye, 1.0)
+    return (1.0 / diffs_safe).masked_fill(mask_eye, 0.0).sum(-1)
+
+def report(label, z, iteration):
+    err = (z.sort(dim=-1).values - true_roots).abs().max().item()
+    # Check for duplicates: min gap between sorted roots
+    zs = z.sort(dim=-1).values
+    min_gap = (zs[:, 1:] - zs[:, :-1]).min().item()
+    # p(z) residual
+    pv, _, _ = horner_pd(c, z)
+    p_res = pv.abs().max().item()
+    print(f"  {label:>5} it={iteration:>2}  max_err={err:.2e}  min_gap={min_gap:.2e}  |p(z)|={p_res:.2e}")
+
+print("="*78)
+print("  Diagnostic: Parallel Root-Finding")
+print("="*78)
+print(f"  B={B} N={N}")
+print(f"  True eigenvalue range: [{true_roots.min().item():.3f}, {true_roots.max().item():.3f}]")
+print(f"  Diagonal init range:   [{z_init.min().item():.3f}, {z_init.max().item():.3f}]")
+
+# ═══ Test 1: Pure parallel Laguerre (no Aberth) ═══
+print(f"\n  --- Test 1: Pure Laguerre (no Aberth) ---")
+z = z_init.clone()
+for it in range(20):
+    step = laguerre_step(c, z, N)
+    z = z - step
+    if it < 5 or it % 5 == 4:
+        report("PurL", z, it)
+
+# ═══ Test 2: Laguerre + Aberth (full strength) ═══
+print(f"\n  --- Test 2: Laguerre + Aberth (full) ---")
+z = z_init.clone()
+for it in range(20):
+    step = laguerre_step(c, z, N)
+    corr = aberth_correction(z)
+    denom = 1.0 - step * corr
+    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))
+    full_step = torch.where(denom.abs() > 1e-20, step / denom_safe, step)
+    z = z - full_step
+    if it < 5 or it % 5 == 4:
+        report("LA-F", z, it)
+
+# ═══ Test 3: Laguerre + weak Aberth (0.1× correction) ═══
+print(f"\n  --- Test 3: Laguerre + weak Aberth (0.1x) ---")
+z = z_init.clone()
+for it in range(20):
+    step = laguerre_step(c, z, N)
+    corr = aberth_correction(z)
+    denom = 1.0 - 0.1 * step * corr
+    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))
+    full_step = torch.where(denom.abs() > 1e-20, step / denom_safe, step)
+    z = z - full_step
+    if it < 5 or it % 5 == 4:
+        report("LA.1", z, it)
+
+# ═══ Test 4: Pure Laguerre + post-sort each iteration ═══
+print(f"\n  --- Test 4: Pure Laguerre + re-sort ---")
+z = z_init.clone()
+for it in range(20):
+    step = laguerre_step(c, z, N)
+    z = z - step
+    z = z.sort(dim=-1).values  # keep sorted
+    if it < 5 or it % 5 == 4:
+        report("PL+S", z, it)
+
+# ═══ Test 5: Laguerre + Aberth + damped ramp ═══
+print(f"\n  --- Test 5: Laguerre + Aberth damped (0.1 → 1.0) ---")
+z = z_init.clone()
+for it in range(20):
+    step = laguerre_step(c, z, N)
+    corr = aberth_correction(z)
+    alpha = min(1.0, 0.1 + 0.1 * it)
+    denom = 1.0 - alpha * step * corr
+    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))
+    full_step = torch.where(denom.abs() > 1e-20, step / denom_safe, step)
+    z = z - full_step
+    z = z.sort(dim=-1).values
+    if it < 5 or it % 5 == 4:
+        report("LADa", z, it)
+
+# ═══ Test 6: Newton + Aberth (original Aberth-Ehrlich) ═══
+print(f"\n  --- Test 6: Newton + Aberth ---")
+z = z_init.clone()
+for it in range(20):
+    pv, dp, _ = horner_pd(c, z)
+    ok = dp.abs() > 1e-30
+    w = torch.where(ok, pv / dp, torch.zeros_like(pv))
+    corr = aberth_correction(z)
+    denom = 1.0 - w * corr
+    denom_safe = torch.where(denom.abs() > 1e-20, denom, torch.ones_like(denom))
+    full_step = torch.where(denom.abs() > 1e-20, w / denom_safe, w)
+    z = z - full_step
+    if it < 5 or it % 5 == 4:
+        report("NwAb", z, it)
+
+# ═══ Test 7: Pure Newton (no Aberth) ═══
+print(f"\n  --- Test 7: Pure Newton ---")
+z = z_init.clone()
+for it in range(20):
+    pv, dp, _ = horner_pd(c, z)
+    ok = dp.abs() > 1e-30
+    w = torch.where(ok, pv / dp, torch.zeros_like(pv))
+    z = z - w
+    if it < 5 or it % 5 == 4:
+        report("PurN", z, it)
+
+print("="*78)