Create analysis_v2.py

analysis_v2.py

@@ -0,0 +1,424 @@
+"""
+Constellation Diffusion — Analysis
+=====================================
+Paste after training. Uses `model` and `bn` from memory.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+import os
+from torchvision import datasets, transforms
+from torchvision.utils import save_image, make_grid
+
+DEVICE = "cuda"
+os.makedirs("analysis_cd", exist_ok=True)
+
+def compute_cv(points, n_samples=1500, n_points=5):
+    N = points.shape[0]
+    if N < n_points: return float('nan')
+    points = F.normalize(points.to(DEVICE).float(), dim=-1)
+    vols = []
+    for _ in range(n_samples):
+        idx = torch.randperm(min(N, 5000), device=DEVICE)[:n_points]
+        pts = points[idx].unsqueeze(0)
+        gram = torch.bmm(pts, pts.transpose(1, 2))
+        norms = torch.diagonal(gram, dim1=1, dim2=2)
+        d2 = norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * gram
+        d2 = F.relu(d2)
+        cm = torch.zeros(1, 6, 6, device=DEVICE, dtype=torch.float32)
+        cm[:, 0, 1:] = 1; cm[:, 1:, 0] = 1; cm[:, 1:, 1:] = d2
+        v2 = -torch.linalg.det(cm) / 9216
+        if v2[0].item() > 1e-20:
+            vols.append(v2[0].sqrt().cpu())
+    if len(vols) < 50: return float('nan')
+    vt = torch.stack(vols)
+    return (vt.std() / (vt.mean() + 1e-8)).item()
+
+def eff_dim(x):
+    x_c = x - x.mean(0, keepdim=True)
+    n = min(512, x.shape[0])
+    _, S, _ = torch.linalg.svd(x_c[:n].float(), full_matrices=False)
+    p = S / S.sum()
+    return p.pow(2).sum().reciprocal().item()
+
+CLASS_NAMES = ['plane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
+
+model.eval()
+bn = model.bottleneck
+
+print("=" * 80)
+print("CONSTELLATION DIFFUSION — PURE BOTTLENECK ANALYSIS")
+n_params = sum(p.numel() for p in model.parameters())
+n_bn = sum(p.numel() for p in bn.parameters())
+print(f"  Total: {n_params:,}  Bottleneck: {n_bn:,} ({100*n_bn/n_params:.1f}%)")
+print(f"  Compression: {bn.spatial_dim} → {bn.n_patches * bn.n_anchors * bn.n_phases} "
+      f"({bn.spatial_dim / (bn.n_patches * bn.n_anchors * bn.n_phases):.1f}×)")
+print("=" * 80)
+
+# Test data
+transform = transforms.Compose([
+    transforms.ToTensor(), transforms.Normalize((0.5,)*3, (0.5,)*3)])
+test_ds = datasets.CIFAR10('./data', train=False, download=True, transform=transform)
+test_loader = torch.utils.data.DataLoader(test_ds, batch_size=256, shuffle=False)
+
+
+# Helper: run encoder to get sphere embeddings
+@torch.no_grad()
+def get_sphere_embeddings(images, labels, t_val=0.0):
+    """Run encoder + projection, return patches on S^15 and tri profiles."""
+    B = images.shape[0]
+    t = torch.full((B,), t_val, device=DEVICE)
+    eps = torch.randn_like(images)
+    t_b = t.view(B, 1, 1, 1)
+    x_t = (1 - t_b) * images + t_b * eps
+
+    cond = model.time_emb(t) + model.class_emb(labels)
+    h = model.in_conv(x_t)
+    for i in range(len(model.ch_mults)):
+        for block in model.enc[i]:
+            if isinstance(block, nn.Sequential):
+                h = block[0](h); h = block[1](h, cond)
+            else:
+                h = block(h, cond)
+        if i < len(model.enc_down):
+            h = model.enc_down[i](h)
+
+    h_flat = h.reshape(B, -1)
+    emb = bn.proj_in(h_flat)
+    patches = emb.reshape(B, bn.n_patches, bn.patch_dim)
+    patches_n = F.normalize(patches, dim=-1)
+    tri = bn.triangulate(patches_n)
+    return patches_n, tri, h_flat
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 1: DRIFT & ANCHOR DIAGNOSTICS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 1: Drift & Anchor Diagnostics")
+print(f"{'━'*80}")
+
+with torch.no_grad():
+    drift = bn.drift().detach()
+    home = F.normalize(bn.home, dim=-1).detach()
+    curr = F.normalize(bn.anchors, dim=-1).detach()
+    P, A, d = home.shape
+
+print(f"  Drift: mean={drift.mean():.6f} rad ({math.degrees(drift.mean().item()):.2f}°)")
+print(f"         max={drift.max():.6f} rad ({math.degrees(drift.max().item()):.2f}°)")
+print(f"  Near 0.29154: {(drift - 0.29154).abs().lt(0.05).float().mean().item():.1%}")
+print(f"  Near 0.29154 (±0.03): {(drift - 0.29154).abs().lt(0.03).float().mean().item():.1%}")
+
+# Drift distribution
+all_d = drift.flatten().cpu().numpy()
+print(f"\n  Drift distribution ({len(all_d)} anchors):")
+bins = [0.0, 0.05, 0.10, 0.15, 0.20, 0.25, 0.29154, 0.35, 0.40, 0.50]
+hist, _ = np.histogram(all_d, bins=bins)
+for i in range(len(bins)-1):
+    bar = "█" * (hist[i] // 2 + (1 if hist[i] > 0 else 0))
+    label = " ◄ BINDING" if bins[i+1] == 0.29154 else ""
+    print(f"    {bins[i]:.3f}-{bins[i+1]:.3f}: {hist[i]:3d} {bar}{label}")
+
+# Per-patch summary
+print(f"\n  Per-patch drift summary:")
+for p in range(P):
+    d_mean = drift[p].mean().item()
+    d_max = drift[p].max().item()
+    n_near = (drift[p] - 0.29154).abs().lt(0.05).sum().item()
+    flags = []
+    if abs(d_mean - 0.29154) < 0.05: flags.append("MEAN≈0.29")
+    if abs(d_max - 0.29154) < 0.05: flags.append("MAX≈0.29")
+    if d_max > 0.29154: flags.append("CROSSED")
+    flag_str = " ◄ " + ", ".join(flags) if flags else ""
+    print(f"    P{p:2d}: mean={d_mean:.4f} max={d_max:.4f} near={n_near}/{A}{flag_str}")
+
+# Anchor spread
+print(f"\n  Anchor effective dimensionality:")
+for p in range(P):
+    _, S, _ = torch.linalg.svd(curr[p].float(), full_matrices=False)
+    pr = S / S.sum()
+    ed = pr.pow(2).sum().reciprocal().item()
+    print(f"    P{p:2d}: {ed:.1f} / {A}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 2: SPHERE GEOMETRY — CV ON S^15
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 2: Sphere Geometry — per-patch CV across timesteps")
+print(f"{'━'*80}")
+
+images_t, labels_t = next(iter(test_loader))
+images_t, labels_t = images_t.to(DEVICE), labels_t.to(DEVICE)
+
+# Per-patch CV at t=0
+patches_n, tri, _ = get_sphere_embeddings(images_t, labels_t, 0.0)
+print(f"\n  Per-patch CV at t=0.0 (natural S^15 = 0.20):")
+for p in range(P):
+    cv_p = compute_cv(patches_n[:, p, :], 1000)
+    print(f"    P{p:2d}: CV={cv_p:.4f}")
+
+# Across timesteps
+print(f"\n  {'t':>6} {'CV_sphere':>10} {'CV_tri':>10} {'eff_d_sph':>10} {'eff_d_tri':>10}")
+for t_val in [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]:
+    pn, tr, _ = get_sphere_embeddings(images_t, labels_t, t_val)
+    sph_flat = pn.reshape(pn.shape[0], -1)
+    cv_s = compute_cv(sph_flat, 1000)
+    cv_t = compute_cv(tr, 1000)
+    ed_s = eff_dim(sph_flat)
+    ed_t = eff_dim(tr)
+    print(f"  {t_val:>6.2f} {cv_s:>10.4f} {cv_t:>10.4f} {ed_s:>10.1f} {ed_t:>10.1f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 3: PER-CLASS ANCHOR ROUTING — ALL PATCHES
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 3: Per-Class Anchor Routing")
+print(f"{'━'*80}")
+
+class_nearest = {c: [] for c in range(10)}
+anchors_n = F.normalize(bn.anchors.detach(), dim=-1)
+
+for imgs_b, labs_b in test_loader:
+    imgs_b, labs_b = imgs_b.to(DEVICE), labs_b.to(DEVICE)
+    pn, _, _ = get_sphere_embeddings(imgs_b, labs_b, 0.0)
+    cos = torch.einsum('bpd,pad->bpa', pn, anchors_n)
+    nearest = cos.argmax(dim=-1).cpu()
+    for i in range(imgs_b.shape[0]):
+        class_nearest[labs_b[i].item()].append(nearest[i])
+    if sum(len(v) for v in class_nearest.values()) > 8000:
+        break
+
+# Show top 4 patches
+for p_idx in range(min(4, P)):
+    print(f"\n  Patch {p_idx}:")
+    print(f"  {'class':>8}", end="")
+    for a in range(A):
+        print(f" {a:>4}", end="")
+    print("   entropy")
+
+    for c in range(10):
+        if not class_nearest[c]: continue
+        nearest_all = torch.stack(class_nearest[c])
+        counts = torch.bincount(nearest_all[:, p_idx], minlength=A).float()
+        counts = counts / counts.sum()
+        entropy = -(counts * (counts + 1e-8).log()).sum().item()
+
+        row = f"  {CLASS_NAMES[c]:>8}"
+        for a in range(A):
+            pct = counts[a].item()
+            if pct > 0.15: row += f" {pct:>3.0%}█"
+            elif pct > 0.08: row += f" {pct:>3.0%}░"
+            elif pct > 0.02: row += f" {pct:>3.0%} "
+            else: row += f"    ."
+        row += f"   {entropy:.2f}"
+        print(row)
+
+# Global utilization
+all_nearest = torch.cat([torch.stack(v) for v in class_nearest.values() if v])
+unique_per_patch = []
+for p_idx in range(P):
+    unique_per_patch.append(all_nearest[:, p_idx].unique().numel())
+print(f"\n  Unique anchors per patch: {unique_per_patch}")
+print(f"  Mean utilization: {np.mean(unique_per_patch):.1f}/{A} "
+      f"({100*np.mean(unique_per_patch)/A:.0f}%)")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 4: RECONSTRUCTION FIDELITY — THROUGH THE BOTTLENECK
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 4: Reconstruction Fidelity — what survives 768 dims?")
+print(f"{'━'*80}")
+
+print(f"  {'t':>6} {'input_norm':>12} {'output_norm':>12} {'cos_sim':>10} "
+      f"{'rel_error':>10} {'mse':>10}")
+
+for t_val in [0.0, 0.25, 0.5, 0.75, 1.0]:
+    B = images_t.shape[0]
+    t = torch.full((B,), t_val, device=DEVICE)
+    eps = torch.randn_like(images_t)
+    t_b = t.view(B, 1, 1, 1)
+    x_t = (1 - t_b) * images_t + t_b * eps
+    cond = model.time_emb(t) + model.class_emb(labels_t)
+
+    with torch.no_grad():
+        # Run encoder
+        h = model.in_conv(x_t)
+        for i in range(len(model.ch_mults)):
+            for block in model.enc[i]:
+                if isinstance(block, nn.Sequential):
+                    h = block[0](h); h = block[1](h, cond)
+                else: h = block(h, cond)
+            if i < len(model.enc_down): h = model.enc_down[i](h)
+
+        h_flat = h.reshape(B, -1)
+        h_reconstructed = bn(h_flat, cond)
+
+        in_norm = h_flat.norm(dim=-1).mean().item()
+        out_norm = h_reconstructed.norm(dim=-1).mean().item()
+        cos = F.cosine_similarity(h_flat, h_reconstructed).mean().item()
+        rel_err = (h_flat - h_reconstructed).norm(dim=-1).mean().item() / (in_norm + 1e-8)
+        mse = F.mse_loss(h_flat, h_reconstructed).item()
+
+    print(f"  {t_val:>6.2f} {in_norm:>12.2f} {out_norm:>12.2f} {cos:>10.6f} "
+          f"{rel_err:>10.4f} {mse:>10.2f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 5: GENERATION QUALITY
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 5: Generation Quality — per class")
+print(f"{'━'*80}")
+
+print(f"  {'class':>8} {'intra_cos':>10} {'std':>8} {'CV':>8} {'norm':>8}")
+
+all_gen = []
+for c in range(10):
+    with torch.no_grad():
+        imgs, _ = sample(model, 64, 50, cls=c)
+    imgs = (imgs + 1) / 2
+    all_gen.append(imgs)
+
+    flat = imgs.reshape(64, -1)
+    flat_n = F.normalize(flat, dim=-1)
+    sim = flat_n @ flat_n.T
+    mask = ~torch.eye(64, device=DEVICE, dtype=torch.bool)
+    print(f"  {CLASS_NAMES[c]:>8} {sim[mask].mean().item():>10.4f} "
+          f"{sim[mask].std().item():>8.4f} {compute_cv(flat, 500):>8.4f} "
+          f"{flat.norm(dim=-1).mean().item():>8.2f}")
+
+    save_image(make_grid(imgs[:16], nrow=4), f"analysis_cd/class_{CLASS_NAMES[c]}.png")
+
+all_grid = torch.cat([g[:4] for g in all_gen])
+save_image(make_grid(all_grid, nrow=10), "analysis_cd/all_classes.png")
+print(f"  ✓ Saved to analysis_cd/")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 6: VELOCITY FIELD
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 6: Velocity Field Quality")
+print(f"{'━'*80}")
+
+print(f"  {'t':>6} {'v_norm':>10} {'v·target':>10} {'mse':>10}")
+
+for t_val in [0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95]:
+    B = 128
+    imgs_v = images_t[:B]
+    labs_v = labels_t[:B]
+    t = torch.full((B,), t_val, device=DEVICE)
+    eps = torch.randn_like(imgs_v)
+    t_b = t.view(B, 1, 1, 1)
+    x_t = (1 - t_b) * imgs_v + t_b * eps
+    v_target = eps - imgs_v
+
+    with torch.no_grad():
+        v_pred = model(x_t, t, labs_v)
+    v_cos = F.cosine_similarity(
+        v_pred.reshape(B, -1), v_target.reshape(B, -1)).mean().item()
+    mse = F.mse_loss(v_pred, v_target).item()
+    v_norm = v_pred.reshape(B, -1).norm(dim=-1).mean().item()
+    print(f"  {t_val:>6.2f} {v_norm:>10.2f} {v_cos:>10.4f} {mse:>10.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 7: ODE TRAJECTORY
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 7: ODE Trajectory — geometry through generation")
+print(f"{'━'*80}")
+
+B_traj = 256
+x = torch.randn(B_traj, 3, 32, 32, device=DEVICE)
+labs_traj = torch.randint(0, 10, (B_traj,), device=DEVICE)
+dt = 1.0 / 50
+
+print(f"  {'step':>6} {'t':>6} {'norm':>10} {'std':>10} {'CV':>8}")
+for step in range(50):
+    t = torch.full((B_traj,), 1.0 - step * dt, device=DEVICE)
+    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        v = model(x, t, labs_traj)
+    x = x - v.float() * dt
+    if step in [0, 1, 5, 10, 20, 30, 40, 49]:
+        xf = x.reshape(B_traj, -1)
+        print(f"  {step:>6} {1.0-step*dt:>6.2f} {xf.norm(dim=-1).mean().item():>10.2f} "
+              f"{x.std().item():>10.4f} {compute_cv(xf, 500):>8.4f}")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 8: INTER vs INTRA CLASS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 8: Class Separation")
+print(f"{'━'*80}")
+
+intra, inter = [], []
+for c in range(10):
+    f = F.normalize(all_gen[c].reshape(64, -1), dim=-1)
+    s = f @ f.T
+    m = ~torch.eye(64, device=DEVICE, dtype=torch.bool)
+    intra.append(s[m].mean().item())
+
+for i in range(10):
+    for j in range(i+1, 10):
+        fi = F.normalize(all_gen[i].reshape(64, -1), dim=-1)
+        fj = F.normalize(all_gen[j].reshape(64, -1), dim=-1)
+        inter.append((fi @ fj.T).mean().item())
+
+print(f"  Intra-class cos: {np.mean(intra):.4f} ± {np.std(intra):.4f}")
+print(f"  Inter-class cos: {np.mean(inter):.4f} ± {np.std(inter):.4f}")
+print(f"  Separation ratio: {np.mean(intra) / (np.mean(inter) + 1e-8):.3f}×")
+
+
+# ══════════════════════════════════════════════════════════════════
+# TEST 9: COMPARISON WITH PREVIOUS VERSIONS
+# ══════════════════════════════════════════════════════════════════
+
+print(f"\n{'━'*80}")
+print("TEST 9: Comparison Summary")
+print(f"{'━'*80}")
+
+print(f"""
+  {'':>25} {'Regulator':>12} {'Skip BN':>12} {'Pure BN':>12}
+  {'':>25} {'(v1)':>12} {'(v2)':>12} {'(v3)':>12}
+  {'─'*73}
+  {'Relay/BN params':>25} {'76K':>12} {'281M':>12} {f'{n_bn:,}':>12}
+  {'Total params':>25} {'6.1M':>12} {'287M':>12} {f'{n_params:,}':>12}
+  {'Best loss':>25} {'0.1900':>12} {'0.1757':>12} {f'{best_loss:.4f}':>12}
+  {'Constellation signal':>25} {'6%':>12} {'88%':>12} {'100%':>12}
+  {'Skip params':>25} {'0':>12} {'268M':>12} {'0':>12}
+  {'Anchor routing':>25} {'2 active':>12} {'class-spec':>12} {'(see T3)':>12}
+""")
+
+# Final drift
+with torch.no_grad():
+    drift = bn.drift().detach()
+    near = (drift - 0.29154).abs().lt(0.05).float().mean().item()
+    near_tight = (drift - 0.29154).abs().lt(0.03).float().mean().item()
+    crossed = (drift > 0.29154).float().mean().item()
+
+print(f"  Final drift stats:")
+print(f"    Mean:         {drift.mean():.6f} rad ({math.degrees(drift.mean().item()):.2f}°)")
+print(f"    Max:          {drift.max():.6f} rad ({math.degrees(drift.max().item()):.2f}°)")
+print(f"    Near 0.29154: {near:.1%} (±0.05)  {near_tight:.1%} (±0.03)")
+print(f"    Crossed 0.29: {crossed:.1%}")
+
+
+print(f"\n{'='*80}")
+print("ANALYSIS COMPLETE")
+print(f"{'='*80}")