Create trainer.py

v18_johanna_curriculum/trainer.py

@@ -0,0 +1,631 @@
+"""
+Johanna-Tiny Full Battery Diagnostic
+======================================
+Comprehensive analysis of the curriculum-trained 16-type noise model.
+
+Tests:
+  1. Per-type MSE (100 samples each, full eval)
+  2. Per-type byte accuracy (discrete reconstruction precision)
+  3. Geometric fingerprint per noise type (S₀, ratio, erank, CV)
+  4. Cross-type omega token similarity (cosine distance matrix)
+  5. Spectrum analysis per type (which modes carry which distributions)
+  6. Reconstruction visualization grid (all 16 types)
+  7. Zero-shot transfer: real images through noise-trained model
+  8. Zero-shot transfer: text bytes through noise-trained model
+  9. Piecemeal 256→64: can tiny do tiled reconstruction?
+  10. Noise-to-noise: encode type A, does it look like type A?
+  11. Effective capacity: what percentage of the signal survives?
+  12. Alpha profile: what did the cross-attention learn?
+"""
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+import math
+import time
+import numpy as np
+import json
+from collections import defaultdict
+
+# ── Load model ───────────────────────────────────────────────────
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Option 1: Load from local checkpoint
+CHECKPOINT = '/content/checkpoints/best.pt'
+# Option 2: Load from HuggingFace
+HF_CHECKPOINT = 'AbstractPhil/geolip-SVAE'
+HF_FILE = 'v18_johanna_curriculum/checkpoints/epoch_0300.pt'
+
+
+def load_model():
+    """Load model from local or HF checkpoint."""
+    from huggingface_hub import hf_hub_download
+
+    # Try local first
+    if os.path.exists(CHECKPOINT):
+        path = CHECKPOINT
+        print(f"  Loading local: {path}")
+    else:
+        path = hf_hub_download(repo_id=HF_CHECKPOINT, filename=HF_FILE, repo_type="model")
+        print(f"  Loading HF: {HF_FILE}")
+
+    ckpt = torch.load(path, map_location='cpu', weights_only=False)
+    cfg = ckpt['config']
+    print(f"  Epoch: {ckpt.get('epoch')}, MSE: {ckpt.get('test_mse', '?')}")
+    print(f"  Config: {cfg}")
+
+    # Build model inline (same architecture)
+    from types import SimpleNamespace
+
+    class BoundarySmooth(nn.Module):
+        def __init__(self, channels=3, mid=16):
+            super().__init__()
+            self.net = nn.Sequential(nn.Conv2d(channels, mid, 3, padding=1), nn.GELU(),
+                                      nn.Conv2d(mid, channels, 3, padding=1))
+            nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)
+        def forward(self, x): return x + self.net(x)
+
+    class SpectralCrossAttention(nn.Module):
+        def __init__(self, D, n_heads=4, max_alpha=0.2, alpha_init=-2.0):
+            super().__init__()
+            self.n_heads = n_heads; self.head_dim = D // n_heads
+            self.max_alpha = max_alpha
+            self.qkv = nn.Linear(D, 3*D); self.out_proj = nn.Linear(D, D)
+            self.norm = nn.LayerNorm(D); self.scale = self.head_dim**-0.5
+            self.alpha_logits = nn.Parameter(torch.full((D,), alpha_init))
+        @property
+        def alpha(self): return self.max_alpha * torch.sigmoid(self.alpha_logits)
+        def forward(self, S):
+            B, N, D = S.shape; S_n = self.norm(S)
+            qkv = self.qkv(S_n).reshape(B,N,3,self.n_heads,self.head_dim).permute(2,0,3,1,4)
+            q, k, v = qkv[0], qkv[1], qkv[2]
+            out = (((q @ k.transpose(-2,-1))*self.scale).softmax(-1) @ v).transpose(1,2).reshape(B,N,D)
+            return S * (1.0 + self.alpha.unsqueeze(0).unsqueeze(0) * torch.tanh(self.out_proj(out)))
+
+    class PatchSVAE(nn.Module):
+        def __init__(self, V=256, D=16, ps=16, hidden=768, depth=4, n_cross=2):
+            super().__init__()
+            self.matrix_v, self.D, self.patch_size = V, D, ps
+            self.patch_dim = 3*ps*ps; self.mat_dim = V*D
+            self.enc_in = nn.Linear(self.patch_dim, hidden)
+            self.enc_blocks = nn.ModuleList([nn.Sequential(
+                nn.LayerNorm(hidden), nn.Linear(hidden, hidden),
+                nn.GELU(), nn.Linear(hidden, hidden)) for _ in range(depth)])
+            self.enc_out = nn.Linear(hidden, self.mat_dim)
+            self.dec_in = nn.Linear(self.mat_dim, hidden)
+            self.dec_blocks = nn.ModuleList([nn.Sequential(
+                nn.LayerNorm(hidden), nn.Linear(hidden, hidden),
+                nn.GELU(), nn.Linear(hidden, hidden)) for _ in range(depth)])
+            self.dec_out = nn.Linear(hidden, self.patch_dim)
+            nn.init.orthogonal_(self.enc_out.weight)
+            self.cross_attn = nn.ModuleList([
+                SpectralCrossAttention(D, n_heads=min(4,D)) for _ in range(n_cross)])
+            self.boundary_smooth = BoundarySmooth(channels=3, mid=16)
+
+        def _svd(self, A):
+            orig = A.dtype
+            with torch.amp.autocast('cuda', enabled=False):
+                A_d = A.double()
+                G = torch.bmm(A_d.transpose(1,2), A_d)
+                G.diagonal(dim1=-2, dim2=-1).add_(1e-12)
+                eig, V = torch.linalg.eigh(G)
+                eig = eig.flip(-1); V = V.flip(-1)
+                S = torch.sqrt(eig.clamp(min=1e-24))
+                U = torch.bmm(A_d, V) / S.unsqueeze(1).clamp(min=1e-16)
+                Vh = V.transpose(-2,-1).contiguous()
+            return U.to(orig), S.to(orig), Vh.to(orig)
+
+        def encode_patches(self, patches):
+            B, N, _ = patches.shape
+            h = F.gelu(self.enc_in(patches.reshape(B*N,-1)))
+            for block in self.enc_blocks: h = h + block(h)
+            M = F.normalize(self.enc_out(h).reshape(B*N, self.matrix_v, self.D), dim=-1)
+            U, S, Vt = self._svd(M)
+            U = U.reshape(B,N,self.matrix_v,self.D); S = S.reshape(B,N,self.D)
+            Vt = Vt.reshape(B,N,self.D,self.D); M = M.reshape(B,N,self.matrix_v,self.D)
+            S_c = S
+            for layer in self.cross_attn: S_c = layer(S_c)
+            return {'U':U, 'S_orig':S, 'S':S_c, 'Vt':Vt, 'M':M}
+
+        def decode_patches(self, U, S, Vt):
+            B, N, V, D = U.shape
+            M_hat = torch.bmm(U.reshape(B*N,V,D)*S.reshape(B*N,D).unsqueeze(1), Vt.reshape(B*N,D,D))
+            h = F.gelu(self.dec_in(M_hat.reshape(B*N,-1)))
+            for block in self.dec_blocks: h = h + block(h)
+            return self.dec_out(h).reshape(B, N, -1)
+
+        def forward(self, images):
+            B, C, H, W = images.shape
+            ps = self.patch_size
+            gh, gw = H//ps, W//ps
+            p = images.reshape(B,C,gh,ps,gw,ps).permute(0,2,4,1,3,5).reshape(B,gh*gw,C*ps*ps)
+            svd = self.encode_patches(p)
+            dec = self.decode_patches(svd['U'], svd['S'], svd['Vt'])
+            dec = dec.reshape(B,gh,gw,3,ps,ps).permute(0,3,1,4,2,5).reshape(B,3,gh*ps,gw*ps)
+            return {'recon': self.boundary_smooth(dec), 'svd': svd, 'gh': gh, 'gw': gw}
+
+        @staticmethod
+        def effective_rank(S):
+            p = S / (S.sum(-1, keepdim=True)+1e-8); p = p.clamp(min=1e-8)
+            return (-(p * p.log()).sum(-1)).exp()
+
+    model = PatchSVAE(V=cfg['V'], D=cfg['D'], ps=cfg['patch_size'],
+                       hidden=cfg['hidden'], depth=cfg['depth'],
+                       n_cross=cfg['n_cross_layers'])
+    model.load_state_dict(ckpt['model_state_dict'], strict=True)
+    model = model.to(DEVICE).eval()
+    print(f"  Loaded {sum(p.numel() for p in model.parameters()):,} params")
+    return model, cfg
+
+
+# ── Noise Generators ─────────────────────────────────────────────
+
+NOISE_NAMES = {
+    0: 'gaussian', 1: 'uniform', 2: 'uniform_scaled', 3: 'poisson',
+    4: 'pink', 5: 'brown', 6: 'salt_pepper', 7: 'sparse',
+    8: 'block', 9: 'gradient', 10: 'checkerboard', 11: 'mixed',
+    12: 'structural', 13: 'cauchy', 14: 'exponential', 15: 'laplace',
+}
+
+def _pink(shape):
+    w = torch.randn(shape); S = torch.fft.rfft2(w)
+    h, ww = shape[-2], shape[-1]
+    fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, ww//2+1)
+    fx = torch.fft.rfftfreq(ww).unsqueeze(0).expand(h, -1)
+    return torch.fft.irfft2(S / torch.sqrt(fx**2 + fy**2).clamp(min=1e-8), s=(h, ww))
+
+def _brown(shape):
+    w = torch.randn(shape); S = torch.fft.rfft2(w)
+    h, ww = shape[-2], shape[-1]
+    fy = torch.fft.fftfreq(h).unsqueeze(-1).expand(-1, ww//2+1)
+    fx = torch.fft.rfftfreq(ww).unsqueeze(0).expand(h, -1)
+    return torch.fft.irfft2(S / (fx**2 + fy**2).clamp(min=1e-8), s=(h, ww))
+
+def generate_noise(noise_type, n, s=64):
+    """Generate n samples of a given noise type."""
+    imgs = []
+    rng = np.random.RandomState(42)
+    for _ in range(n):
+        if noise_type == 0: img = torch.randn(3,s,s)
+        elif noise_type == 1: img = torch.rand(3,s,s)*2-1
+        elif noise_type == 2: img = (torch.rand(3,s,s)-0.5)*4
+        elif noise_type == 3:
+            lam = rng.uniform(0.5, 20.0)
+            img = torch.poisson(torch.full((3,s,s), lam))/lam - 1.0
+        elif noise_type == 4: img = _pink((3,s,s)); img = img/(img.std()+1e-8)
+        elif noise_type == 5: img = _brown((3,s,s)); img = img/(img.std()+1e-8)
+        elif noise_type == 6:
+            img = torch.where(torch.rand(3,s,s)>0.5, torch.ones(3,s,s)*2, -torch.ones(3,s,s)*2)
+            img = img + torch.randn(3,s,s)*0.1
+        elif noise_type == 7: img = torch.randn(3,s,s)*(torch.rand(3,s,s)>0.9).float()*3
+        elif noise_type == 8:
+            b = rng.randint(2,16); sm = torch.randn(3,s//b+1,s//b+1)
+            img = F.interpolate(sm.unsqueeze(0), size=s, mode='nearest').squeeze(0)
+        elif noise_type == 9:
+            gy = torch.linspace(-2,2,s).unsqueeze(1).expand(s,s)
+            gx = torch.linspace(-2,2,s).unsqueeze(0).expand(s,s)
+            a = rng.uniform(0, 2*math.pi)
+            img = (math.cos(a)*gx + math.sin(a)*gy).unsqueeze(0).expand(3,-1,-1) + torch.randn(3,s,s)*0.5
+        elif noise_type == 10:
+            cs = rng.randint(2,16); cy = torch.arange(s)//cs; cx = torch.arange(s)//cs
+            img = ((cy.unsqueeze(1)+cx.unsqueeze(0))%2).float().unsqueeze(0).expand(3,-1,-1)*2-1 + torch.randn(3,s,s)*0.3
+        elif noise_type == 11:
+            alpha = rng.uniform(0.2, 0.8)
+            img = alpha*torch.randn(3,s,s) + (1-alpha)*(torch.rand(3,s,s)*2-1)
+        elif noise_type == 12:
+            img = torch.zeros(3,s,s); h2 = s//2
+            img[:,:h2,:h2] = torch.randn(3,h2,h2)
+            img[:,:h2,h2:] = torch.rand(3,h2,h2)*2-1
+            img[:,h2:,:h2] = _pink((3,h2,h2))/2
+            img[:,h2:,h2:] = torch.where(torch.rand(3,h2,h2)>0.5, torch.ones(3,h2,h2), -torch.ones(3,h2,h2))
+        elif noise_type == 13: img = torch.tan(math.pi*(torch.rand(3,s,s)-0.5)).clamp(-3,3)
+        elif noise_type == 14: img = torch.empty(3,s,s).exponential_(1.0)-1.0
+        elif noise_type == 15:
+            u = torch.rand(3,s,s)-0.5; img = -torch.sign(u)*torch.log1p(-2*u.abs())
+        else: img = torch.randn(3,s,s)
+        imgs.append(img.clamp(-4,4))
+    return torch.stack(imgs)
+
+
+# ════════════════════════════════════════════════════════════════
+# DIAGNOSTIC TESTS
+# ════════════════════════════════════════════════════════════════
+
+def test_1_per_type_mse(model, n=100, s=64):
+    """Per-type reconstruction MSE."""
+    print(f"\n{'='*70}")
+    print("TEST 1: Per-Type Reconstruction MSE (100 samples each)")
+    print(f"{'='*70}")
+    results = {}
+    model.eval()
+    with torch.no_grad():
+        for t in range(16):
+            imgs = generate_noise(t, n, s).to(DEVICE)
+            out = model(imgs)
+            mse = F.mse_loss(out['recon'], imgs, reduction='none').mean(dim=(1,2,3))
+            results[NOISE_NAMES[t]] = {
+                'mean': mse.mean().item(),
+                'std': mse.std().item(),
+                'min': mse.min().item(),
+                'max': mse.max().item(),
+            }
+            print(f"  {NOISE_NAMES[t]:18s}: {mse.mean():.6f} ± {mse.std():.6f}  "
+                  f"[{mse.min():.6f} — {mse.max():.6f}]")
+    return results
+
+
+def test_2_byte_accuracy(model, n=100, s=64):
+    """Byte-level reconstruction accuracy per type."""
+    print(f"\n{'='*70}")
+    print("TEST 2: Byte-Level Accuracy (quantized to 256 levels)")
+    print(f"{'='*70}")
+    results = {}
+    model.eval()
+    with torch.no_grad():
+        for t in range(16):
+            imgs = generate_noise(t, n, s).to(DEVICE)
+            out = model(imgs)
+            # Quantize to 256 levels
+            orig_q = ((imgs + 4) / 8 * 255).round().clamp(0, 255).long()
+            recon_q = ((out['recon'] + 4) / 8 * 255).round().clamp(0, 255).long()
+            acc = (orig_q == recon_q).float().mean().item()
+            # Within-1 accuracy
+            acc1 = ((orig_q - recon_q).abs() <= 1).float().mean().item()
+            results[NOISE_NAMES[t]] = {'exact': acc, 'within_1': acc1}
+            print(f"  {NOISE_NAMES[t]:18s}: exact={acc*100:5.1f}%  ±1={acc1*100:5.1f}%")
+    return results
+
+
+def test_3_geometric_fingerprint(model, n=64, s=64):
+    """Geometric properties per noise type."""
+    print(f"\n{'='*70}")
+    print("TEST 3: Geometric Fingerprint Per Type")
+    print(f"{'='*70}")
+    D = model.D
+    results = {}
+    model.eval()
+    with torch.no_grad():
+        for t in range(16):
+            imgs = generate_noise(t, n, s).to(DEVICE)
+            out = model(imgs)
+            S = out['svd']['S']  # (B, N, D)
+            S_mean = S.mean(dim=(0, 1))
+            ratio = (S_mean[0] / (S_mean[-1] + 1e-8)).item()
+            erank = model.effective_rank(S.reshape(-1, D)).mean().item()
+            s0 = S_mean[0].item()
+            sd = S_mean[-1].item()
+            results[NOISE_NAMES[t]] = {'S0': s0, 'SD': sd, 'ratio': ratio, 'erank': erank}
+            print(f"  {NOISE_NAMES[t]:18s}: S₀={s0:.3f}  SD={sd:.3f}  "
+                  f"ratio={ratio:.2f}  erank={erank:.2f}")
+    return results
+
+
+def test_4_omega_similarity(model, n=32, s=64):
+    """Cross-type omega token cosine similarity matrix."""
+    print(f"\n{'='*70}")
+    print("TEST 4: Cross-Type Omega Token Similarity")
+    print(f"{'='*70}")
+    D = model.D
+    type_centroids = {}
+    model.eval()
+    with torch.no_grad():
+        for t in range(16):
+            imgs = generate_noise(t, n, s).to(DEVICE)
+            out = model(imgs)
+            # Average omega token per type: (D,)
+            omega = out['svd']['S'].mean(dim=(0, 1))
+            type_centroids[t] = omega
+
+    # Cosine similarity matrix
+    keys = sorted(type_centroids.keys())
+    centroids = torch.stack([type_centroids[k] for k in keys])
+    centroids_norm = F.normalize(centroids, dim=-1)
+    sim_matrix = centroids_norm @ centroids_norm.T
+
+    # Print matrix
+    header = "         " + " ".join([f"{NOISE_NAMES[k][:5]:>5s}" for k in keys])
+    print(f"  {header}")
+    for i, ki in enumerate(keys):
+        row = f"  {NOISE_NAMES[ki]:8s}"
+        for j, kj in enumerate(keys):
+            v = sim_matrix[i, j].item()
+            row += f" {v:5.2f}"
+        print(row)
+    return sim_matrix.cpu()
+
+
+def test_5_spectrum_per_type(model, n=64, s=64):
+    """Singular value spectrum analysis per type."""
+    print(f"\n{'='*70}")
+    print("TEST 5: Spectrum Profile Per Type")
+    print(f"{'='*70}")
+    D = model.D
+    results = {}
+    model.eval()
+    with torch.no_grad():
+        for t in range(16):
+            imgs = generate_noise(t, n, s).to(DEVICE)
+            out = model(imgs)
+            S_mean = out['svd']['S'].mean(dim=(0, 1))
+            total = (S_mean**2).sum()
+            cum = 0
+            spectrum = []
+            for d in range(D):
+                e = (S_mean[d]**2).item()
+                cum += e
+                spectrum.append({'value': S_mean[d].item(), 'energy_pct': cum/total.item()*100})
+            results[NOISE_NAMES[t]] = spectrum
+
+    # Print top-3 and bottom-3 modes per type
+    for t in range(16):
+        name = NOISE_NAMES[t]
+        sp = results[name]
+        top = f"S0={sp[0]['value']:.3f}({sp[0]['energy_pct']:.1f}%)"
+        mid = f"S7={sp[7]['value']:.3f}({sp[7]['energy_pct']:.1f}%)"
+        bot = f"S15={sp[15]['value']:.3f}(100%)"
+        print(f"  {name:18s}: {top}  {mid}  {bot}")
+    return results
+
+
+def test_6_reconstruction_grid(model, s=64):
+    """Visual reconstruction grid — all 16 types."""
+    print(f"\n{'='*70}")
+    print("TEST 6: Reconstruction Grid (saved to johanna_diagnostic_grid.png)")
+    print(f"{'='*70}")
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+
+    model.eval()
+    fig, axes = plt.subplots(16, 3, figsize=(9, 48))
+
+    with torch.no_grad():
+        for t in range(16):
+            img = generate_noise(t, 1, s).to(DEVICE)
+            out = model(img)
+            recon = out['recon']
+            mse = F.mse_loss(recon, img).item()
+
+            orig_np = img[0].cpu().clamp(-3, 3).permute(1, 2, 0).numpy()
+            recon_np = recon[0].cpu().clamp(-3, 3).permute(1, 2, 0).numpy()
+            diff_np = (img[0] - recon[0]).abs().cpu().clamp(0, 2).permute(1, 2, 0).numpy()
+
+            # Normalize for display
+            for arr in [orig_np, recon_np]:
+                arr -= arr.min(); arr /= (arr.max() + 1e-8)
+            diff_np /= (diff_np.max() + 1e-8)
+
+            axes[t, 0].imshow(orig_np); axes[t, 0].set_ylabel(NOISE_NAMES[t], fontsize=8)
+            axes[t, 1].imshow(recon_np)
+            axes[t, 2].imshow(diff_np)
+            for j in range(3):
+                axes[t, j].axis('off')
+
+    axes[0, 0].set_title('Original', fontsize=9)
+    axes[0, 1].set_title('Recon', fontsize=9)
+    axes[0, 2].set_title('|Error|', fontsize=9)
+    plt.tight_layout()
+    plt.savefig('johanna_diagnostic_grid.png', dpi=150, bbox_inches='tight')
+    print("  Saved: johanna_diagnostic_grid.png")
+    plt.close()
+
+
+def test_7_real_images(model, s=64):
+    """Zero-shot: real images through noise-trained model."""
+    print(f"\n{'='*70}")
+    print("TEST 7: Zero-Shot Real Image Reconstruction")
+    print(f"{'='*70}")
+    from datasets import load_dataset
+
+    ds = load_dataset('zh-plus/tiny-imagenet', split='valid', streaming=True)
+    transform = T.Compose([T.ToTensor(), T.Normalize((0.4802,0.4481,0.3975),(0.2770,0.2691,0.2821))])
+
+    imgs = []
+    for i, sample in enumerate(ds):
+        img = sample['image'].convert('RGB')
+        imgs.append(transform(img))
+        if i >= 99:
+            break
+
+    batch = torch.stack(imgs).to(DEVICE)
+    model.eval()
+    with torch.no_grad():
+        out = model(batch)
+        mse = F.mse_loss(out['recon'], batch, reduction='none').mean(dim=(1,2,3))
+
+    print(f"  TinyImageNet (100 images, {s}×{s}):")
+    print(f"    Mean MSE: {mse.mean():.6f}")
+    print(f"    Std:      {mse.std():.6f}")
+    print(f"    Min/Max:  {mse.min():.6f} / {mse.max():.6f}")
+    print(f"    Fidelity: {(1 - mse.mean())*100:.3f}%")
+    return {'mean': mse.mean().item(), 'std': mse.std().item()}
+
+
+def test_8_text_bytes(model, s=64):
+    """Zero-shot: text through noise-trained model."""
+    print(f"\n{'='*70}")
+    print("TEST 8: Zero-Shot Text Byte Reconstruction")
+    print(f"{'='*70}")
+
+    texts = [
+        "Hello, world! This is a test of the Johanna geometric encoder.",
+        "The quick brown fox jumps over the lazy dog. 0123456789 ABCDEF",
+        "import torch; model = PatchSVAE(); output = model(x)",
+        "E = mc² — Albert Einstein, theoretical physicist, 1905",
+        "To be, or not to be, that is the question. — Shakespeare",
+    ]
+
+    n_bytes = 3 * s * s
+    model.eval()
+
+    for text in texts:
+        raw = text.encode('utf-8')
+        actual_len = min(len(raw), n_bytes)
+        if len(raw) < n_bytes:
+            raw = raw + b'\x00' * (n_bytes - len(raw))
+        else:
+            raw = raw[:n_bytes]
+
+        arr = np.frombuffer(raw, dtype=np.uint8).copy()
+        tensor = torch.from_numpy(arr).float()
+        tensor = (tensor / 127.5) - 1.0
+        tensor = tensor.reshape(1, 3, s, s).to(DEVICE)
+
+        with torch.no_grad():
+            out = model(tensor)
+            recon = out['recon']
+            mse = F.mse_loss(recon, tensor).item()
+
+        recon_bytes = ((recon.squeeze(0).cpu().flatten() + 1.0) * 127.5).round().clamp(0, 255).byte().numpy()
+        recovered = recon_bytes[:actual_len].tobytes().decode('utf-8', errors='replace')
+
+        orig_b = ((tensor.squeeze(0).cpu().flatten() + 1.0) * 127.5).round().clamp(0, 255).byte()
+        recon_b = ((recon.squeeze(0).cpu().flatten() + 1.0) * 127.5).round().clamp(0, 255).byte()
+        exact_acc = (orig_b[:actual_len] == recon_b[:actual_len]).float().mean().item()
+
+        print(f"\n  Input:    '{text[:60]}'")
+        print(f"  Output:   '{recovered[:60]}'")
+        print(f"  MSE:      {mse:.6f}")
+        print(f"  Byte acc: {exact_acc*100:.1f}%")
+
+
+def test_9_piecemeal(model, s=64):
+    """Piecemeal: tile 256×256 noise into 64×64 tiles."""
+    print(f"\n{'='*70}")
+    print(f"TEST 9: Piecemeal 256→{s} Tiled Reconstruction")
+    print(f"{'='*70}")
+    model.eval()
+
+    results = {}
+    with torch.no_grad():
+        for t in [0, 4, 6, 13]:  # Gaussian, Pink, Salt-pepper, Cauchy
+            img_256 = generate_noise(t, 1, 256).squeeze(0)  # (3, 256, 256)
+            tiles = []
+            gh, gw = 256 // s, 256 // s
+            for gy in range(gh):
+                for gx in range(gw):
+                    tile = img_256[:, gy*s:(gy+1)*s, gx*s:(gx+1)*s]
+                    tiles.append(tile)
+            tile_batch = torch.stack(tiles).to(DEVICE)
+            out = model(tile_batch)
+            recon_tiles = out['recon'].cpu()
+
+            # Stitch
+            recon_full = torch.zeros(3, 256, 256)
+            idx = 0
+            for gy in range(gh):
+                for gx in range(gw):
+                    recon_full[:, gy*s:(gy+1)*s, gx*s:(gx+1)*s] = recon_tiles[idx]
+                    idx += 1
+
+            mse = F.mse_loss(recon_full, img_256).item()
+            results[NOISE_NAMES[t]] = mse
+            n_tiles = gh * gw
+            print(f"  {NOISE_NAMES[t]:18s}: {n_tiles} tiles, MSE={mse:.6f}")
+    return results
+
+
+def test_10_signal_survival(model, n=100, s=64):
+    """What percentage of the original signal energy survives reconstruction?"""
+    print(f"\n{'='*70}")
+    print("TEST 10: Signal Energy Survival Rate")
+    print(f"{'='*70}")
+    model.eval()
+    with torch.no_grad():
+        for t in range(16):
+            imgs = generate_noise(t, n, s).to(DEVICE)
+            out = model(imgs)
+            recon = out['recon']
+            orig_energy = (imgs**2).mean().item()
+            recon_energy = (recon**2).mean().item()
+            error_energy = ((imgs - recon)**2).mean().item()
+            survival = recon_energy / (orig_energy + 1e-8) * 100
+            snr = 10 * math.log10(orig_energy / (error_energy + 1e-8))
+            print(f"  {NOISE_NAMES[t]:18s}: survival={survival:6.1f}%  SNR={snr:5.1f}dB  "
+                  f"orig_E={orig_energy:.3f}  recon_E={recon_energy:.3f}")
+
+
+def test_11_alpha_profile(model):
+    """Cross-attention alpha analysis."""
+    print(f"\n{'='*70}")
+    print("TEST 11: Cross-Attention Alpha Profile")
+    print(f"{'='*70}")
+    for li, layer in enumerate(model.cross_attn):
+        alpha = layer.alpha.detach().cpu()
+        print(f"\n  Layer {li}: mean={alpha.mean():.4f}  max={alpha.max():.4f}  "
+              f"min={alpha.min():.4f}  std={alpha.std():.6f}")
+        bar_scale = 50 / (alpha.max().item() + 1e-8)
+        for d in range(len(alpha)):
+            bar = "█" * int(alpha[d].item() * bar_scale)
+            print(f"    α[{d:2d}]: {alpha[d]:.5f}  {bar}")
+
+
+def test_12_compression_ratio(model, s=64):
+    """Actual compression metrics."""
+    print(f"\n{'='*70}")
+    print("TEST 12: Compression Metrics")
+    print(f"{'='*70}")
+    D = model.D
+    ps = model.patch_size
+    n_patches = (s // ps) ** 2
+    input_values = 3 * s * s
+    latent_values = D * n_patches
+    ratio = input_values / latent_values
+    print(f"  Input:   {s}×{s}×3 = {input_values:,} values")
+    print(f"  Latent:  {D}×{n_patches} = {latent_values:,} values (omega tokens)")
+    print(f"  Ratio:   {ratio:.1f}:1 compression")
+    print(f"  Patches: {n_patches} of {ps}×{ps}")
+    print(f"  Omega shape: ({D}, {s//ps}, {s//ps})")
+
+    # Bits per value at different quantization levels
+    for bits in [8, 16, 32]:
+        input_bytes = input_values * (bits // 8)
+        latent_bytes = latent_values * (bits // 8)
+        print(f"  At {bits}-bit: input={input_bytes/1024:.1f}KB  latent={latent_bytes/1024:.1f}KB  "
+              f"ratio={input_bytes/latent_bytes:.1f}:1")
+
+
+# ═════════════════════════════════════════════════════════════��══
+# RUN ALL
+# ════════════════════════════════════════════════════════════════
+
+def run_all():
+    print("=" * 70)
+    print("JOHANNA-TINY FULL BATTERY DIAGNOSTIC")
+    print("=" * 70)
+
+    model, cfg = load_model()
+    s = cfg.get('img_size', 64)
+
+    all_results = {}
+    all_results['config'] = cfg
+
+    all_results['per_type_mse'] = test_1_per_type_mse(model, n=100, s=s)
+    all_results['byte_accuracy'] = test_2_byte_accuracy(model, n=100, s=s)
+    all_results['geometry'] = test_3_geometric_fingerprint(model, n=64, s=s)
+    sim_matrix = test_4_omega_similarity(model, n=32, s=s)
+    all_results['spectrum'] = test_5_spectrum_per_type(model, n=64, s=s)
+    test_6_reconstruction_grid(model, s=s)
+    all_results['real_images'] = test_7_real_images(model, s=s)
+    test_8_text_bytes(model, s=s)
+    all_results['piecemeal'] = test_9_piecemeal(model, s=s)
+    test_10_signal_survival(model, n=100, s=s)
+    test_11_alpha_profile(model)
+    test_12_compression_ratio(model, s=s)
+
+    # Save results
+    out_path = 'johanna_diagnostic_results.json'
+    with open(out_path, 'w') as f:
+        json.dump(all_results, f, indent=2, default=str)
+    print(f"\n  Results saved: {out_path}")
+
+    print(f"\n{'='*70}")
+    print("DIAGNOSTIC COMPLETE")
+    print(f"{'='*70}")
+
+
+if __name__ == "__main__":
+    run_all()