Update tests.py

tests.py

@@ -1,92 +1,362 @@
 """
-Flow ensemble — smoke test + diagnostics.
+Flow Ensemble — Expanded Test Suite.
+
+Assumes geolip-core is installed (Colab with repo loaded).
+Tests: smoke, linalg integration, multi-scale, ensemble fusion,
+       gradient health, ablation, compile compatibility, memory.
 """
 import torch
 import torch.nn as nn
-import sys, time
+import torch.nn.functional as F
+import sys, time, gc
+
+# ── Verify geolip_core.linalg is available ──
+try:
+    import geolip_core.linalg as LA
+    HAS_GEOLIP_LINALG = True
+    print(f"geolip_core.linalg: available")
+    LA.backend.status()
+except ImportError:
+    import torch.linalg as LA
+    HAS_GEOLIP_LINALG = False
+    print("geolip_core.linalg: NOT available, using torch.linalg fallback")
 
-sys.path.insert(0, '.')
-from flows import (
-    QuaternionFlow, QuaternionLiteFlow, VelocityFlow,
-    MagnitudeFlow, OrbitalFlow, AlignmentFlow, FlowEnsemble,
-)
 
 dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-B, n, k, d = 32, 128, 64, 256
 
-anchors = torch.randn(B, k, d, device=dev)
-queries = torch.randn(B, n, d, device=dev)
 
-print("=" * 68)
-print("  Flow Ensemble — Smoke Test")
-print("=" * 68)
-print(f"  B={B} n={n} k={k} d={d} device={dev}")
+def sync():
+    if dev.type == 'cuda':
+        torch.cuda.synchronize()
+
+def time_fn(fn, warmup=5, runs=50):
+    for _ in range(warmup): fn()
+    sync()
+    t0 = time.perf_counter()
+    for _ in range(runs): fn()
+    sync()
+    return (time.perf_counter() - t0) / runs * 1000
+
+def fmt(ms):
+    if ms < 1: return f"{ms*1000:.0f}us"
+    return f"{ms:.2f}ms"
+
+def make_data(B, n, k, d):
+    anchors = F.normalize(torch.randn(B, k, d, device=dev), dim=-1)
+    queries = F.normalize(torch.randn(B, n, d, device=dev), dim=-1)
+    return anchors, queries
+
+
+# ═══════════════════════════════════════════════════════════════════
+print("=" * 72)
+print("  Flow Ensemble — Expanded Test Suite")
+print("=" * 72)
+print(f"  device={dev}  geolip_core.linalg={HAS_GEOLIP_LINALG}")
+if dev.type == 'cuda':
+    print(f"  GPU: {torch.cuda.get_device_name()}")
+print()
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 1. SMOKE TEST — all flows, all shapes
+# ═══════════════════════════════════════════════════════════════════
+print(f"{'='*72}\n  1. SMOKE TEST\n{'='*72}")
+
+B, n, k, d = 16, 64, 32, 128
+anchors, queries = make_data(B, n, k, d)
 
-# Test each flow independently
 flows_cfg = [
-    ('QuaternionFlow',     lambda: QuaternionFlow(d, k, n_heads=4)),
-    ('QuaternionLiteFlow', lambda: QuaternionLiteFlow(d, k)),
-    ('VelocityFlow',       lambda: VelocityFlow(d, k)),
-    ('MagnitudeFlow',      lambda: MagnitudeFlow(d, k)),
-    ('OrbitalFlow',        lambda: OrbitalFlow(d, k)),
-    ('AlignmentFlow',      lambda: AlignmentFlow(d, k)),
+    ('QuaternionFlow',     lambda d,k: QuaternionFlow(d, k, n_heads=4)),
+    ('QuaternionLiteFlow', lambda d,k: QuaternionLiteFlow(d, k)),
+    ('VelocityFlow',       lambda d,k: VelocityFlow(d, k)),
+    ('MagnitudeFlow',      lambda d,k: MagnitudeFlow(d, k)),
+    ('OrbitalFlow',        lambda d,k: OrbitalFlow(d, k)),
+    ('AlignmentFlow',      lambda d,k: AlignmentFlow(d, k)),
 ]
 
-print(f"\n  {'Flow':<22} {'Params':>8} {'Out shape':>14} {'Fwd (ms)':>10} {'Conf μ':>8}")
-print(f"  {'─'*22} {'─'*8} {'─'*14} {'─'*10} {'─'*8}")
+print(f"\n  {'Flow':<22} {'Params':>8} {'Shape':>14} {'Time':>10} {'Conf':>8} {'Res norm':>10}")
+print(f"  {'─'*22} {'─'*8} {'─'*14} {'─'*10} {'─'*8} {'─'*10}")
 
 live_flows = []
+flow_ctors = []
 for name, ctor in flows_cfg:
     try:
-        flow = ctor().to(dev)
+        flow = ctor(d, k).to(dev)
         params = sum(p.numel() for p in flow.parameters())
-
-        # Warmup
-        for _ in range(3):
-            flow(anchors, queries)
-        if dev.type == 'cuda':
-            torch.cuda.synchronize()
-
-        # Time
-        t0 = time.perf_counter()
-        N_runs = 50
-        for _ in range(N_runs):
-            pred, conf = flow(anchors, queries)
-        if dev.type == 'cuda':
-            torch.cuda.synchronize()
-        elapsed = (time.perf_counter() - t0) / N_runs * 1000
-
-        print(f"  {name:<22} {params:>8,} {str(tuple(pred.shape)):>14} {elapsed:>9.2f} {conf.mean().item():>8.3f}")
+        pred, conf = flow(anchors, queries)
+        ms = time_fn(lambda: flow(anchors, queries))
+        res = (pred - queries).norm(dim=-1).mean().item()
+        shape_str = str(tuple(pred.shape))
+        print(f"  {name:<22} {params:>8,} {shape_str:>14} {fmt(ms):>10} {conf.mean().item():>8.3f} {res:>10.3f}")
         live_flows.append(flow)
+        flow_ctors.append((name, ctor))
     except Exception as e:
-        print(f"  {name:<22} FAILED: {str(e)[:40]}")
+        print(f"  {name:<22} FAILED: {str(e)[:50]}")
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 2. LINALG INTEGRATION
+# ═══════════════════════════════════════════════════════════════════
+print(f"\n{'='*72}\n  2. LINALG INTEGRATION\n{'='*72}")
+
+if HAS_GEOLIP_LINALG:
+    print(f"\n  Testing eigh dispatch in MagnitudeFlow and OrbitalFlow...")
+    for FlowCls in [MagnitudeFlow, OrbitalFlow]:
+        flow = FlowCls(d, k).to(dev)
+        pred, conf = flow(anchors, queries)
+        ok = torch.isfinite(pred).all().item() and torch.isfinite(conf).all().item()
+        print(f"  {flow.name:<18} finite={ok}  conf={conf.mean():.3f}")
+
+    oflow = OrbitalFlow(d, k).to(dev)
+    a_geom = oflow.anchor_proj(anchors)
+    G = torch.bmm(a_geom.transpose(-2, -1), a_geom)
+    vals, vecs = LA.eigh(G)
+    print(f"\n  Gram eigenspectrum: shape={tuple(vals.shape)} "
+          f"range=[{vals.min().item():.4f}, {vals.max().item():.4f}]")
+    print(f"  Eigenvector orth err: {(torch.bmm(vecs.mT, vecs) - torch.eye(oflow.geom_dim, device=dev)).abs().max().item():.2e}")
+else:
+    print("  Skipped — geolip_core.linalg not available")
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 3. MULTI-SCALE
+# ═══════════════════════════════════════════════════════════════════
+print(f"\n{'='*72}\n  3. MULTI-SCALE\n{'='*72}")
+
+configs = [
+    (4,  16,  8,   64,  'tiny'),
+    (16, 64,  32,  128, 'small'),
+    (32, 128, 64,  256, 'medium'),
+    (64, 256, 128, 256, 'large'),
+    (8,  512, 256, 512, 'wide'),
+]
+
+print(f"\n  OrbitalFlow across scales:")
+print(f"  {'Config':<10} {'B':>4} {'n':>5} {'k':>5} {'d':>5} {'Time':>10} {'OK':>4}")
+print(f"  {'─'*10} {'─'*4} {'─'*5} {'─'*5} {'─'*5} {'─'*10} {'─'*4}")
+
+for B_, n_, k_, d_, label in configs:
+    try:
+        of = OrbitalFlow(d_, k_).to(dev)
+        a, q = make_data(B_, n_, k_, d_)
+        pred, conf = of(a, q)
+        ms = time_fn(lambda: of(a, q), warmup=3, runs=20)
+        ok = torch.isfinite(pred).all().item()
+        print(f"  {label:<10} {B_:>4} {n_:>5} {k_:>5} {d_:>5} {fmt(ms):>10} {'OK' if ok else 'NO':>4}")
+        del of, a, q
+    except Exception as e:
+        print(f"  {label:<10} {B_:>4} {n_:>5} {k_:>5} {d_:>5} FAILED: {str(e)[:30]}")
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 4. ENSEMBLE FUSION MODES
+# ═══════════════════════════════════════════════════════════════════
+print(f"\n{'='*72}\n  4. ENSEMBLE FUSION\n{'='*72}")
+
+B, n, k, d = 16, 64, 32, 128
+anchors, queries = make_data(B, n, k, d)
 
-# Test ensemble
-print(f"\n  Ensemble tests:")
 for fusion in ['weighted', 'gated', 'residual']:
+    ens = FlowEnsemble(live_flows, d, fusion=fusion).to(dev)
+    out = ens(anchors, queries)
+    ms = time_fn(lambda: ens(anchors, queries), warmup=3, runs=20)
+
+    preds = [flow(anchors, queries)[0] for flow in ens.flows]
+    cos_sims = []
+    for i in range(len(preds)):
+        for j in range(i+1, len(preds)):
+            cs = F.cosine_similarity(preds[i].flatten(1), preds[j].flatten(1), dim=-1).mean().item()
+            cos_sims.append(cs)
+    avg_sim = sum(cos_sims) / max(len(cos_sims), 1)
+
+    print(f"\n  {fusion}:  time={fmt(ms)}  norm={out.norm(dim=-1).mean():.3f}  diversity={1-avg_sim:.3f}")
+    diag = ens.flow_diagnostics(anchors, queries)
+    for fname, stats in diag.items():
+        print(f"    {fname:<18} conf={stats['confidence_mean']:.3f}±{stats['confidence_std']:.3f}  "
+              f"res={stats['residual_norm']:.3f}")
+    del ens
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 5. GRADIENT HEALTH
+# ═══════════════════════════════════════════════════════════════════
+print(f"\n{'='*72}\n  5. GRADIENT HEALTH\n{'='*72}")
+
+B, n, k, d = 16, 64, 32, 128
+anchors, queries = make_data(B, n, k, d)
+
+losses = {
+    'mse': (lambda o,q: (o - q).pow(2).mean()),
+    'cosine': (lambda o,q: (1 - F.cosine_similarity(o, q, dim=-1)).mean()),
+    'norm': (lambda o,q: o.norm(dim=-1).mean()),
+}
+
+print(f"\n  {'Flow':<18} {'Loss':<10} {'Grad norm':>12} {'Status':>8}")
+print(f"  {'─'*18} {'─'*10} {'─'*12} {'─'*8}")
+
+for loss_name, loss_fn in losses.items():
+    # Fresh flows for each loss — avoids in-place grad corruption across losses
+    try:
+        test_flows_grad = [ctor(d, k).to(dev) for _, ctor in flow_ctors]
+        ens_g = FlowEnsemble(test_flows_grad, d, fusion='residual').to(dev)
+        ens_g.zero_grad()
+        anchors_g = anchors.detach().clone().requires_grad_(True)
+        queries_g = queries.detach().clone().requires_grad_(True)
+        out = ens_g(anchors_g, queries_g)
+        loss = loss_fn(out, queries_g.detach())
+        loss.backward()
+
+        for flow in ens_g.flows:
+            grads = [p.grad for p in flow.parameters() if p.grad is not None]
+            if grads:
+                gn = torch.cat([g.flatten() for g in grads]).norm().item()
+                status = "OK" if 1e-8 < gn < 1e4 else "WARN"
+                print(f"  {flow.name:<18} {loss_name:<10} {gn:>12.2e} {status:>8}")
+            else:
+                print(f"  {flow.name:<18} {loss_name:<10} {'no grads':>12} {'WARN':>8}")
+        del ens_g, test_flows_grad
+    except RuntimeError as e:
+        if 'inplace' in str(e).lower() or 'in-place' in str(e).lower() or 'modified by' in str(e):
+            print(f"  {'*':>18} {loss_name:<10} {'IN-PLACE ERR':>12} {'NOTE':>8}")
+            print(f"    FL eigh deflation uses indexed assignment — needs .clone() fix")
+        else:
+            print(f"  {'*':>18} {loss_name:<10} {'ERROR':>12}")
+            print(f"    {str(e)[:60]}")
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 6. ABLATION — solo vs pairs vs full ensemble
+# ═══════════════════════════════════════════════════════════════════
+print(f"\n{'='*72}\n  6. ABLATION (100 training steps, rotation target)\n{'='*72}")
+
+B, n, k, d = 32, 128, 64, 256
+anchors, queries = make_data(B, n, k, d)
+R = torch.linalg.qr(torch.randn(d, d, device=dev)).Q.unsqueeze(0)
+target = torch.bmm(queries, R.expand(B, -1, -1))
+
+def eval_quality(model, anchors, queries, target, steps=100, lr=1e-3):
+    opt = torch.optim.Adam(model.parameters(), lr=lr)
+    for _ in range(steps):
+        opt.zero_grad()
+        pred = model(anchors, queries) if isinstance(model, FlowEnsemble) else model(anchors, queries)[0]
+        loss = (pred - target).pow(2).mean()
+        loss.backward()
+        opt.step()
+    with torch.no_grad():
+        pred = model(anchors, queries) if isinstance(model, FlowEnsemble) else model(anchors, queries)[0]
+        return (pred - target).pow(2).mean().item()
+
+print(f"\n  {'Configuration':<35} {'MSE':>10} {'Params':>10}")
+print(f"  {'─'*35} {'─'*10} {'─'*10}")
+
+for name, ctor in flow_ctors:
+    try:
+        flow = ctor(d, k).to(dev)
+        params = sum(p.numel() for p in flow.parameters())
+        mse = eval_quality(flow, anchors, queries, target)
+        print(f"  {name:<35} {mse:>10.4f} {params:>10,}")
+        del flow
+    except Exception as e:
+        print(f"  {name:<35} FAILED: {str(e)[:30]}")
+
+pairs = [
+    ('Quat + Orbital', [0, 4]),
+    ('Velocity + Magnitude', [2, 3]),
+    ('Orbital + Alignment', [4, 5]),
+    ('Velocity + Orbital', [2, 4]),
+]
+for pair_name, indices in pairs:
+    try:
+        pair_flows = [flow_ctors[i][1](d, k).to(dev) for i in indices if i < len(flow_ctors)]
+        if len(pair_flows) >= 2:
+            ens = FlowEnsemble(pair_flows, d, fusion='weighted').to(dev)
+            params = sum(p.numel() for p in ens.parameters())
+            mse = eval_quality(ens, anchors, queries, target)
+            print(f"  {pair_name:<35} {mse:>10.4f} {params:>10,}")
+            del ens, pair_flows
+    except Exception as e:
+        print(f"  {pair_name:<35} FAILED: {str(e)[:30]}")
+
+for fusion in ['weighted', 'residual']:
     try:
-        ens = FlowEnsemble(live_flows, d, fusion=fusion).to(dev)
+        all_flows = [ctor(d, k).to(dev) for _, ctor in flow_ctors]
+        ens = FlowEnsemble(all_flows, d, fusion=fusion).to(dev)
         params = sum(p.numel() for p in ens.parameters())
-        out = ens(anchors, queries)
-        print(f"  {fusion:<12} params={params:>10,}  out={tuple(out.shape)}  norm={out.norm(dim=-1).mean():.3f}")
-
-        # Diagnostics
-        diag = ens.flow_diagnostics(anchors, queries)
-        for fname, stats in diag.items():
-            print(f"    {fname:<18} conf={stats['confidence_mean']:.3f}±{stats['confidence_std']:.3f}  "
-                  f"residual={stats['residual_norm']:.3f}  temp={stats['temperature']:.3f}")
+        mse = eval_quality(ens, anchors, queries, target)
+        print(f"  {'Full (' + fusion + ')':<35} {mse:>10.4f} {params:>10,}")
+        del ens, all_flows
     except Exception as e:
-        print(f"  {fusion:<12} FAILED: {str(e)[:50]}")
-
-# Gradient flow test
-print(f"\n  Gradient flow:")
-ens = FlowEnsemble(live_flows, d, fusion='weighted').to(dev)
-out = ens(anchors, queries)
-loss = out.sum()
-loss.backward()
-for flow in ens.flows:
-    grads = [p.grad is not None for p in flow.parameters()]
-    pct = sum(grads) / max(len(grads), 1) * 100
-    print(f"  {flow.name:<18} {pct:.0f}% params have gradients")
-
-print("=" * 68)
+        print(f"  {'Full (' + fusion + ')':<35} FAILED: {str(e)[:30]}")
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 7. COMPILE COMPATIBILITY
+# ═══════════════════════════════════════════════════════════════════
+print(f"\n{'='*72}\n  7. COMPILE COMPATIBILITY\n{'='*72}")
+
+B, n, k, d = 8, 32, 16, 64
+anchors, queries = make_data(B, n, k, d)
+
+print(f"\n  {'Flow':<22} {'fullgraph':>12} {'Raw':>10} {'Compiled':>12}")
+print(f"  {'─'*22} {'─'*12} {'─'*10} {'─'*12}")
+
+for name, ctor in flow_ctors:
+    try:
+        flow = ctor(d, k).to(dev)
+        t_raw = time_fn(lambda: flow(anchors, queries), warmup=3, runs=30)
+        try:
+            compiled = torch.compile(flow, fullgraph=True)
+            compiled(anchors, queries); sync()
+            t_comp = time_fn(lambda: compiled(anchors, queries), warmup=3, runs=30)
+            status = "OK"
+        except Exception as e:
+            t_comp = -1
+            status = str(e)[:12]
+        t_str = fmt(t_comp) if t_comp > 0 else "N/A"
+        print(f"  {name:<22} {status:>12} {fmt(t_raw):>10} {t_str:>12}")
+        del flow
+    except Exception as e:
+        print(f"  {name:<22} FAILED: {str(e)[:40]}")
+
+
+# ═══════════════════════════════════════════════════════════════════
+# 8. MEMORY
+# ═══════════════════════════════════════════════════════════════════
+if dev.type == 'cuda':
+    print(f"\n{'='*72}\n  8. MEMORY (B=32, n=128, k=64, d=256)\n{'='*72}")
+
+    B, n, k, d = 32, 128, 64, 256
+    anchors, queries = make_data(B, n, k, d)
+
+    print(f"\n  {'Flow':<22} {'Peak MB':>10}")
+    print(f"  {'─'*22} {'─'*10}")
+
+    for name, ctor in flow_ctors:
+        try:
+            flow = ctor(d, k).to(dev)
+            torch.cuda.empty_cache(); gc.collect()
+            torch.cuda.reset_peak_memory_stats()
+            base = torch.cuda.memory_allocated()
+            pred, conf = flow(anchors, queries); sync()
+            peak = (torch.cuda.max_memory_allocated() - base) / 1024**2
+            print(f"  {name:<22} {peak:>9.1f}")
+            del flow, pred, conf
+        except Exception as e:
+            print(f"  {name:<22} FAILED: {str(e)[:30]}")
+
+    try:
+        all_flows = [ctor(d, k).to(dev) for _, ctor in flow_ctors]
+        ens = FlowEnsemble(all_flows, d, fusion='weighted').to(dev)
+        torch.cuda.empty_cache(); gc.collect()
+        torch.cuda.reset_peak_memory_stats()
+        base = torch.cuda.memory_allocated()
+        out = ens(anchors, queries); sync()
+        peak = (torch.cuda.max_memory_allocated() - base) / 1024**2
+        print(f"  {'Full ensemble':<22} {peak:>9.1f}")
+        del ens, all_flows
+    except Exception as e:
+        print(f"  {'Full ensemble':<22} FAILED: {str(e)[:30]}")
+
+print(f"\n{'='*72}")
+print(f"  Done.")
+print(f"{'='*72}")