Create tests.py

tests.py

@@ -0,0 +1,92 @@
+"""
+Flow ensemble — smoke test + diagnostics.
+"""
+import torch
+import torch.nn as nn
+import sys, time
+
+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}")
+
+# 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)),
+]
+
+print(f"\n  {'Flow':<22} {'Params':>8} {'Out shape':>14} {'Fwd (ms)':>10} {'Conf μ':>8}")
+print(f"  {'─'*22} {'─'*8} {'─'*14} {'─'*10} {'─'*8}")
+
+live_flows = []
+for name, ctor in flows_cfg:
+    try:
+        flow = ctor().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}")
+        live_flows.append(flow)
+    except Exception as e:
+        print(f"  {name:<22} FAILED: {str(e)[:40]}")
+
+# Test ensemble
+print(f"\n  Ensemble tests:")
+for fusion in ['weighted', 'gated', 'residual']:
+    try:
+        ens = FlowEnsemble(live_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}")
+    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)