Create 18_j_cell_reformed.py

18_j_cell_reformed.py

@@ -0,0 +1,609 @@
+# ═══════════════════════════════════════════════════════════════════════
+# Cell J''' — per-patch axis-feature classifier
+# ═══════════════════════════════════════════════════════════════════════
+# Same task as J/J'/J''. Same architectures. Only the encode + feature
+# extraction stages change.
+#
+# Key fix from J' and J'':
+#   J'  : encode_axes(images, patch_idx=0) → [B, V, n_axes]
+#         → max-pool over V → [B, n_axes]
+#         Used 1 of 256 patches per tile.
+#
+#   J'' : same as J' but with V-stats instead of max-pool.
+#         Still using 1 of 256 patches per tile.
+#
+#   J''': encode_axes(images)  # no patch_idx → [B, n_patches=256, V, n_axes]
+#         → spatial stats over patches AND value stats over V
+#         Uses ALL 256 patches per tile. 256× more spatial signal.
+#
+# Codebooks calibrated with the new per-patch averaging path
+# (sample_agg='mean', patch_agg='mean') for codebooks that reflect the
+# bank's spatial-mean response.
+
+import json
+import math
+import time
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import geolip_svae.arrays
+from transformers import AutoModel
+
+array_model = globals().get('array_model')
+if array_model is None:
+    array_model = AutoModel.from_pretrained("AbstractPhil/geolip-svae-h2-64")
+    array_model = (array_model.cuda().eval()
+                    if torch.cuda.is_available() else array_model.eval())
+
+DEVICE = next(array_model.parameters()).device
+EXP_DIR = Path("/content/h2_64_exp")
+EXP_DIR.mkdir(parents=True, exist_ok=True)
+TILE_SIZE = 64
+
+NOISE_NAMES = {
+    0: 'gaussian', 1: 'uniform', 2: 'uniform_scaled', 3: 'poisson',
+    4: 'pink', 5: 'brown', 6: 'salt_pepper', 7: 'sparse_impulses',
+    8: 'block_upsampled', 9: 'gradient_gaussian', 10: 'checker',
+    11: 'gauss_uniform_mix', 12: 'four_quadrant',
+    13: 'cauchy', 14: 'exponential', 15: 'laplace',
+}
+
+# ══════════════════════════════════════════════════════════════════════
+# Noise generators — inlined from Cell J so this cell is self-contained
+# ══════════════════════════════════════════════════════════════════════
+
+def _pink(shape, rng):
+    w = torch.randn(shape, generator=rng)
+    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, rng):
+    w = torch.randn(shape, generator=rng)
+    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 gen_noise(noise_type, size, seed):
+    """Pure noise generator. size must be even (some generators use s//2)."""
+    rng_t = torch.Generator().manual_seed(seed)
+    rng_n = np.random.RandomState(seed)
+    s = size
+    if noise_type == 0:
+        img = torch.randn(3, s, s, generator=rng_t)
+    elif noise_type == 1:
+        img = torch.rand(3, s, s, generator=rng_t) * 2 - 1
+    elif noise_type == 2:
+        img = (torch.rand(3, s, s, generator=rng_t) - 0.5) * 4
+    elif noise_type == 3:
+        lam = rng_n.uniform(0.5, 20.0)
+        img = torch.poisson(torch.full((3, s, s), lam), generator=rng_t) / lam - 1.0
+    elif noise_type == 4:
+        img = _pink((3, s, s), rng_t); img = img / (img.std() + 1e-8)
+    elif noise_type == 5:
+        img = _brown((3, s, s), rng_t); img = img / (img.std() + 1e-8)
+    elif noise_type == 6:
+        mask = torch.rand(3, s, s, generator=rng_t) > 0.5
+        img = torch.where(mask, torch.ones(3, s, s) * 2, torch.ones(3, s, s) * -2)
+        img = img + torch.randn(3, s, s, generator=rng_t) * 0.1
+    elif noise_type == 7:
+        mask = torch.rand(3, s, s, generator=rng_t) > 0.9
+        img = torch.randn(3, s, s, generator=rng_t) * mask.float() * 3
+    elif noise_type == 8:
+        block = rng_n.randint(2, 16)
+        small = torch.randn(3, s // block + 1, s // block + 1, generator=rng_t)
+        img = F.interpolate(small.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)
+        angle = rng_n.uniform(0, 2 * math.pi)
+        grad = math.cos(angle) * gx + math.sin(angle) * gy
+        img = (grad.unsqueeze(0).expand(3, -1, -1)
+                + torch.randn(3, s, s, generator=rng_t) * 0.5)
+    elif noise_type == 10:
+        cs = rng_n.randint(2, 16)
+        cy = torch.arange(s) // cs; cx = torch.arange(s) // cs
+        checker = ((cy.unsqueeze(1) + cx.unsqueeze(0)) % 2).float() * 2 - 1
+        img = (checker.unsqueeze(0).expand(3, -1, -1)
+                + torch.randn(3, s, s, generator=rng_t) * 0.3)
+    elif noise_type == 11:
+        a = torch.randn(3, s, s, generator=rng_t)
+        b = torch.rand(3, s, s, generator=rng_t) * 2 - 1
+        alpha = rng_n.uniform(0.2, 0.8)
+        img = alpha * a + (1 - alpha) * b
+    elif noise_type == 12:
+        img = torch.zeros(3, s, s)
+        h2 = s // 2
+        img[:, :h2, :h2] = torch.randn(3, h2, h2, generator=rng_t)
+        img[:, :h2, h2:] = torch.rand(3, h2, h2, generator=rng_t) * 2 - 1
+        img[:, h2:, :h2] = _pink((3, h2, h2), rng_t) / 2
+        sp = torch.where(torch.rand(3, h2, h2, generator=rng_t) > 0.5,
+                          torch.ones(3, h2, h2), -torch.ones(3, h2, h2))
+        img[:, h2:, h2:] = sp
+    elif noise_type == 13:
+        u = torch.rand(3, s, s, generator=rng_t)
+        img = torch.tan(math.pi * (u - 0.5)).clamp(-3, 3)
+    elif noise_type == 14:
+        img = torch.empty(3, s, s).exponential_(1.0, generator=rng_t) - 1.0
+    elif noise_type == 15:
+        u = torch.rand(3, s, s, generator=rng_t) - 0.5
+        img = -torch.sign(u) * torch.log1p(-2 * u.abs())
+    else:
+        raise ValueError(f"Unknown noise_type {noise_type}")
+    return img.clamp(-4, 4).float()
+
+
+def gen_zone_matte(res, n_zones, seed):
+    """Spatially-mixed noise: n_zones grid of different noise types."""
+    assert n_zones in (4, 9, 16), "Use 2×2, 3×3, or 4×4 grids"
+    side = int(math.sqrt(n_zones))
+    cell = res // side
+    assert cell % 2 == 0, f"cell size {cell} must be even for noise generators"
+    rng_n = np.random.RandomState(seed)
+    zone_types = rng_n.choice(16, size=n_zones, replace=False).tolist()
+    img = torch.zeros(3, res, res)
+    zone_map = torch.zeros(res, res, dtype=torch.long)
+    for i in range(side):
+        for j in range(side):
+            zi = i * side + j
+            nt = zone_types[zi]
+            cell_seed = seed * 1000 + zi + 1
+            cell_img = gen_noise(nt, cell, cell_seed)
+            img[:, i*cell:(i+1)*cell, j*cell:(j+1)*cell] = cell_img
+            zone_map[i*cell:(i+1)*cell, j*cell:(j+1)*cell] = zi
+    return img, zone_types, zone_map
+
+SUBSET_BATTERY_IDS = list(range(16)) + [19, 20]
+SUBSET_PHASE = 'best'
+N_BATTERIES_SUB = len(SUBSET_BATTERY_IDS)
+COMP_LABELS = {16: 'zone_4', 17: 'zone_9', 18: 'zone_16'}
+N_CLASSES = 16 + len(COMP_LABELS)
+
+def label_name(n):
+    return NOISE_NAMES.get(n, COMP_LABELS.get(n, f"?{n}"))
+
+print("=" * 78)
+print("PHASE J''' — PER-PATCH AXIS-FEATURE SCANNER")
+print("=" * 78)
+print(f"Subset: {N_BATTERIES_SUB} batteries, phase={SUBSET_PHASE}")
+print(f"Per tile: 256 patches × 32 V × ~27 axes = ~221K activations")
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Codebook calibration — use the new batched + per-patch API
+# ══════════════════════════════════════════════════════════════════════
+
+print(f"\nCalibrating codebooks (batched, per-patch averaging)...")
+t0 = time.time()
+g = torch.Generator().manual_seed(42)
+calib_imgs = torch.randn(512, 3, 64, 64, generator=g)
+
+targets = [(bid, SUBSET_PHASE) for bid in SUBSET_BATTERY_IDS]
+codebooks_dict = array_model.compute_axis_codebooks(
+    targets=targets,
+    calibration_images=calib_imgs,
+    sample_agg='mean',
+    patch_agg='mean',  # NEW: average across 256 patches per image
+    batch_size=64,
+)
+
+codebooks = {bid: codebooks_dict[(bid, SUBSET_PHASE)].to(DEVICE)
+             for bid in SUBSET_BATTERY_IDS}
+for bid in SUBSET_BATTERY_IDS:
+    print(f"  battery {bid:>2} ({label_name(bid):<22}): "
+          f"{codebooks[bid].shape[0]} axes")
+
+MAX_AXES = max(cb.shape[0] for cb in codebooks.values())
+print(f"Calibration time: {time.time() - t0:.1f}s,  MAX_AXES: {MAX_AXES}")
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Per-patch tile scan — uses encode_axes WITHOUT patch_idx
+# ══════════════════════════════════════════════════════════════════════
+
+@torch.no_grad()
+def perpatch_tile_scan(image, codebooks, battery_ids, tile_size=TILE_SIZE):
+    """For each tile, get full per-patch activations against each bank.
+
+    Returns features condensed at scan-time (full tensor too big to store):
+        [n_tiles, n_banks, MAX_AXES, N_STATS]
+    where N_STATS = stats over (n_patches × V) jointly.
+
+    Stats per (bank, axis): max, mean, std, top10_mean, entropy
+    Computed over the joint distribution of activations across both
+    patches AND V rows (because both contribute to "how does this tile
+    align with this axis").
+    """
+    C, H, W = image.shape
+    n_h, n_w = H // tile_size, W // tile_size
+    n_tiles = n_h * n_w
+
+    tiles = image.unfold(1, tile_size, tile_size).unfold(2, tile_size, tile_size)
+    tiles = tiles.permute(1, 2, 0, 3, 4).contiguous().reshape(
+        n_tiles, C, tile_size, tile_size).to(DEVICE)
+
+    n_banks = len(battery_ids)
+    out = torch.zeros(n_tiles, n_banks, MAX_AXES, N_STATS, dtype=torch.float32)
+
+    tile_batch = 32  # smaller because per-patch activations are larger
+    for b_i, bid in enumerate(battery_ids):
+        cb = codebooks[bid]                      # [n_axes_i, D]
+        n_axes_i = cb.shape[0]
+
+        for start in range(0, n_tiles, tile_batch):
+            end = min(start + tile_batch, n_tiles)
+            batch = tiles[start:end]
+            # Per-patch encoding: [B_t, n_patches=256, V=32, n_axes_i]
+            acts = array_model.encode_axes(
+                images=batch, battery_idx=bid,
+                phase=SUBSET_PHASE, codebook=cb,
+            )
+
+            B_t, P, V, n_ax = acts.shape
+
+            # Reshape to [B_t, P*V, n_axes_i] — joint patch+V distribution
+            joint = acts.reshape(B_t, P * V, n_ax).cpu()
+
+            # Stats over the joint patch+V dimension, per axis:
+            mx = joint.max(dim=1).values                           # [B_t, n_ax]
+            mn = joint.mean(dim=1)                                  # [B_t, n_ax]
+            sd = joint.std(dim=1)                                   # [B_t, n_ax]
+            
+            k = min(10, P * V)
+            top_k = joint.topk(k, dim=1).values
+            top10 = top_k.mean(dim=1)                               # [B_t, n_ax]
+
+            # Entropy over softmax(activations across patch×V): low entropy
+            # means a few specific (patch, V-row) positions dominate, high
+            # entropy means uniform alignment across the spatial-row plane.
+            sm = F.softmax(joint, dim=1)                            # [B_t, P*V, n_ax]
+            ent = -(sm * (sm + 1e-12).log()).sum(dim=1)             # [B_t, n_ax]
+            ent = ent / math.log(P * V)                              # normalized
+
+            # Stack [B_t, n_axes_i, N_STATS] and place into output
+            stats = torch.stack([mx, mn, sd, top10, ent], dim=-1)
+            out[start:end, b_i, :n_axes_i, :] = stats
+
+    return out
+
+
+N_STATS = 5  # max, mean, std, top10_mean, entropy
+print(f"\nN_STATS per (bank, axis) per tile: {N_STATS}")
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Build feature bank
+# ══════════════════════════════════════════════════════════════════════
+
+RESOLUTIONS = [256, 512, 1024]
+N_IMAGES_PER_LABEL = 24
+
+print(f"\nBuilding per-patch axis-stats scan bank...")
+scan_bank = {}
+t0 = time.time()
+
+for res in RESOLUTIONS:
+    scan_bank[res] = {}
+    n_tiles = (res // TILE_SIZE) ** 2
+    print(f"\n  res={res} ({n_tiles} tiles per image):")
+
+    for nt in range(16):
+        for img_idx in range(N_IMAGES_PER_LABEL):
+            seed = 1_000_000 + res * 100 + nt * 100 + img_idx
+            img = gen_noise(nt, res, seed)
+            stats = perpatch_tile_scan(img, codebooks, SUBSET_BATTERY_IDS)
+            scan_bank[res][(nt, img_idx)] = stats
+        print(f"    {label_name(nt):<22} done")
+
+    for zone_n, zone_lbl in [(4, 16), (9, 17), (16, 18)]:
+        side = int(math.sqrt(zone_n))
+        if res % side != 0 or (res // side) % 2 != 0:
+            print(f"    {label_name(zone_lbl):<22} SKIP")
+            continue
+        for img_idx in range(N_IMAGES_PER_LABEL):
+            seed = 2_000_000 + res * 100 + zone_n * 10 + img_idx
+            img, _, _ = gen_zone_matte(res, zone_n, seed)
+            stats = perpatch_tile_scan(img, codebooks, SUBSET_BATTERY_IDS)
+            scan_bank[res][(zone_lbl, img_idx)] = stats
+        print(f"    {label_name(zone_lbl):<22} done")
+
+print(f"\nTotal scan time: {time.time() - t0:.1f}s")
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Feature builders — A''' summary, B''' attn-pool over tiles
+# ═════════════════���════════════════════════════════════════════════════
+
+def perpatch_summary_features(stats_scan):
+    """Aggregate over tiles via mean+max for each (bank, axis, stat).
+
+    stats_scan: [n_tiles, n_banks, MAX_AXES, N_STATS]
+    Returns: [n_banks * MAX_AXES * N_STATS * 2] flat
+    """
+    mn = stats_scan.mean(dim=0)
+    mx = stats_scan.max(dim=0).values
+    return torch.stack([mn, mx], dim=-1).flatten()
+
+
+def perpatch_tile_grid_features(stats_scan, max_tiles=16):
+    """Tile-grid for attention pool. Flattens (bank, axis, stat) per tile.
+
+    Returns: [max_tiles, n_banks * MAX_AXES * N_STATS]
+    """
+    n_tiles, n_banks, n_ax, n_st = stats_scan.shape
+    flat = stats_scan.reshape(n_tiles, n_banks * n_ax * n_st)
+    if n_tiles >= max_tiles:
+        idx = torch.randperm(n_tiles)[:max_tiles]
+        return flat[idx]
+    pad = torch.zeros(max_tiles - n_tiles, n_banks * n_ax * n_st)
+    return torch.cat([flat, pad], dim=0)
+
+
+print(f"\nBuilding feature tensors per resolution...")
+features_A_by_res = {}
+features_B_by_res = {}
+labels_by_res = {}
+MAX_TILES = 16
+
+for res in RESOLUTIONS:
+    feat_A, feat_B, labs = [], [], []
+    for (lbl, img_idx), stats in scan_bank[res].items():
+        feat_A.append(perpatch_summary_features(stats))
+        feat_B.append(perpatch_tile_grid_features(stats, max_tiles=MAX_TILES))
+        labs.append(lbl)
+    features_A_by_res[res] = torch.stack(feat_A)
+    features_B_by_res[res] = torch.stack(feat_B)
+    labels_by_res[res] = torch.tensor(labs)
+    print(f"  res={res}: {features_A_by_res[res].shape[0]} samples, "
+          f"A''' feat {features_A_by_res[res].shape[1]}-dim, "
+          f"B''' feat {tuple(features_B_by_res[res].shape[1:])}")
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Classifiers (same as J/J'/J'' for fair comparison)
+# ══════════════════════════════════════════════════════════════════════
+
+class SummaryMLP(nn.Module):
+    def __init__(self, in_dim, hidden=128, n_classes=N_CLASSES):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, n_classes),
+        )
+    def forward(self, x): return self.net(x)
+
+
+class AttentionPoolMLP(nn.Module):
+    def __init__(self, n_features, hidden=128, n_classes=N_CLASSES):
+        super().__init__()
+        self.attn_scorer = nn.Linear(n_features, 1)
+        self.classifier = nn.Sequential(
+            nn.Linear(n_features, hidden),
+            nn.ReLU(),
+            nn.Linear(hidden, n_classes),
+        )
+    def forward(self, x):
+        scores = self.attn_scorer(x).squeeze(-1)
+        weights = torch.softmax(scores, dim=1)
+        pooled = (x * weights.unsqueeze(-1)).sum(dim=1)
+        return self.classifier(pooled)
+
+
+def train_classifier(model_cls, train_x, train_y, test_x, test_y,
+                      in_spec, n_epochs=200, lr=1e-2):
+    torch.manual_seed(42)
+    clf = model_cls(in_spec)
+    optimizer = torch.optim.Adam(clf.parameters(), lr=lr)
+    batch_size = 128
+    train_hist, test_hist = [], []
+    for epoch in range(n_epochs):
+        perm = torch.randperm(train_x.shape[0])
+        clf.train()
+        for i in range(0, train_x.shape[0], batch_size):
+            idx = perm[i:i + batch_size]
+            loss = F.cross_entropy(clf(train_x[idx]), train_y[idx])
+            optimizer.zero_grad(); loss.backward(); optimizer.step()
+        clf.eval()
+        with torch.no_grad():
+            train_acc = (clf(train_x).argmax(dim=1) == train_y).float().mean().item()
+            test_acc = (clf(test_x).argmax(dim=1) == test_y).float().mean().item()
+        train_hist.append(train_acc); test_hist.append(test_acc)
+
+    clf.eval()
+    with torch.no_grad():
+        preds = clf(test_x).argmax(dim=1)
+    classes = torch.unique(test_y).tolist()
+    per_class = {c: ((preds == test_y) & (test_y == c)).sum().item() /
+                    max(1, (test_y == c).sum().item())
+                 for c in classes}
+    return test_hist[-1], per_class, train_hist, test_hist
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Train + compare against all priors
+# ══════════════════════════════════════════════════════════════════════
+
+ref_paths = {
+    'cell_j':    EXP_DIR / "results_expJ.json",
+    'cell_jp':   EXP_DIR / "results_expJ_axes.json",
+    'cell_jpp':  EXP_DIR / "results_expJ_vstats.json",
+}
+refs = {}
+for k, p in ref_paths.items():
+    if p.exists():
+        with open(p) as f:
+            r = json.load(f)
+        refs[k] = {int(res): {'A': v['accuracy_A'], 'B': v['accuracy_B']}
+                   for res, v in r['per_resolution'].items()}
+        print(f"\nLoaded {k}: {p}")
+
+results = {}
+for res in RESOLUTIONS:
+    print(f"\n{'─' * 78}")
+    print(f"Resolution {res}×{res}")
+    print(f"{'─' * 78}")
+
+    n_items = features_A_by_res[res].shape[0]
+    rng = np.random.RandomState(42)
+    indices = rng.permutation(n_items)
+    n_train = int(n_items * 0.8)
+    train_idx, test_idx = indices[:n_train], indices[n_train:]
+    labels = labels_by_res[res]
+
+    # A'''
+    xA = features_A_by_res[res]
+    mA, sA = xA[train_idx].mean(dim=0), xA[train_idx].std(dim=0).clamp(min=1e-8)
+    xA = (xA - mA) / sA
+    accA, per_class_A, tA, vA = train_classifier(
+        SummaryMLP, xA[train_idx], labels[train_idx],
+        xA[test_idx], labels[test_idx], in_spec=xA.shape[1])
+
+    # B'''
+    xB = features_B_by_res[res]
+    flat = xB[train_idx].reshape(-1, xB.shape[-1])
+    mB, sB = flat.mean(dim=0), flat.std(dim=0).clamp(min=1e-8)
+    xB = (xB - mB) / sB
+    accB, per_class_B, tB, vB = train_classifier(
+        AttentionPoolMLP, xB[train_idx], labels[train_idx],
+        xB[test_idx], labels[test_idx], in_spec=xB.shape[-1])
+
+    print(f"  A''' (per-patch summary):  test={accA:.1%}")
+    if 'cell_j' in refs:
+        d = accA - refs['cell_j'][res]['A']
+        print(f"    vs Cell J A (MSE):       {refs['cell_j'][res]['A']:.1%}   Δ {d:+.1%}")
+    if 'cell_jp' in refs:
+        d = accA - refs['cell_jp'][res]['A']
+        print(f"    vs Cell J' A (max-axes): {refs['cell_jp'][res]['A']:.1%}   Δ {d:+.1%}")
+    if 'cell_jpp' in refs:
+        d = accA - refs['cell_jpp'][res]['A']
+        print(f"    vs Cell J'' A (V-stats): {refs['cell_jpp'][res]['A']:.1%}   Δ {d:+.1%}")
+
+    print(f"\n  B''' (per-patch attn):     test={accB:.1%}")
+    if 'cell_j' in refs:
+        d = accB - refs['cell_j'][res]['B']
+        print(f"    vs Cell J B (MSE):       {refs['cell_j'][res]['B']:.1%}   Δ {d:+.1%}")
+    if 'cell_jp' in refs:
+        d = accB - refs['cell_jp'][res]['B']
+        print(f"    vs Cell J' B (max-axes): {refs['cell_jp'][res]['B']:.1%}   Δ {d:+.1%}")
+    if 'cell_jpp' in refs:
+        d = accB - refs['cell_jpp'][res]['B']
+        print(f"    vs Cell J'' B (V-stats): {refs['cell_jpp'][res]['B']:.1%}   Δ {d:+.1%}")
+
+    print(f"\n  {'Class':<22} {'A':>9} {'B':>9}  {'Δ(B-A)':>9}")
+    for c in sorted(per_class_A.keys()):
+        a = per_class_A[c]; b = per_class_B.get(c, 0.0)
+        sym = '+' if b > a + 0.01 else '-' if b < a - 0.01 else ' '
+        print(f"  {label_name(c):<22} {a:>9.1%} {b:>9.1%}  {sym}{abs(b-a):>8.1%}")
+
+    results[res] = {
+        'accuracy_A': accA, 'accuracy_B': accB,
+        'per_class_A': {label_name(c): per_class_A[c] for c in per_class_A},
+        'per_class_B': {label_name(c): per_class_B.get(c, 0.0) for c in per_class_A},
+        'train_curve_A': tA, 'test_curve_A': vA,
+        'train_curve_B': tB, 'test_curve_B': vB,
+    }
+
+
+# ══════════════════════════════════════════════════════════════════════
+# Plots + verdict
+# ══════════════════════════════════════════════════════════════════════
+
+fig, axes = plt.subplots(1, 2, figsize=(20, 6))
+for idx, (clf_label, key_curve, key_acc) in enumerate(
+    [("A''' per-patch summary MLP", 'test_curve_A', 'accuracy_A'),
+     ("B''' per-patch attn-pool MLP", 'test_curve_B', 'accuracy_B')]
+):
+    ax = axes[idx]
+    for res in RESOLUTIONS:
+        ax.plot(results[res][key_curve],
+                label=f'{res} ({results[res][key_acc]:.1%})',
+                linewidth=1.5, alpha=0.85)
+    ax.axhline(1 / N_CLASSES, color='gray', linestyle='--', linewidth=1,
+                label=f'Random ({1/N_CLASSES:.1%})')
+    ax.set_xlabel('Epoch'); ax.set_ylabel('Test accuracy')
+    ax.set_title(clf_label)
+    ax.legend(loc='lower right'); ax.grid(linestyle=':', alpha=0.5)
+    ax.set_ylim(0, 1.05)
+plt.tight_layout()
+plt.savefig(EXP_DIR / 'expJ_perpatch_curves.png', dpi=120, bbox_inches='tight')
+plt.show()
+
+
+print(f"\n{'=' * 78}")
+print(f"PHASE J''' VERDICT — per-patch axis features")
+print(f"{'=' * 78}")
+
+if 'cell_j' in refs:
+    print(f"\n{'Res':<6} | {'MSE':>7} {'maxax':>7} {'vstat':>7} {'perpatch':>9}    "
+          f"| {'MSE':>7} {'maxax':>7} {'vstat':>7} {'perpatch':>9}")
+    print(f"       | {'A':>7} {'A':>7} {'A':>7} {'A':>9}    "
+          f"| {'B':>7} {'B':>7} {'B':>7} {'B':>9}")
+    print("-" * 100)
+    for res in RESOLUTIONS:
+        ja = refs['cell_j'][res]['A']; jb = refs['cell_j'][res]['B']
+        pa = refs.get('cell_jp', {}).get(res, {}).get('A', float('nan'))
+        pb = refs.get('cell_jp', {}).get(res, {}).get('B', float('nan'))
+        va = refs.get('cell_jpp', {}).get(res, {}).get('A', float('nan'))
+        vb = refs.get('cell_jpp', {}).get(res, {}).get('B', float('nan'))
+        ka = results[res]['accuracy_A']; kb = results[res]['accuracy_B']
+        print(f"{str(res):<6} | {ja:>6.1%} {pa:>6.1%} {va:>6.1%} {ka:>8.1%}    "
+              f"| {jb:>6.1%} {pb:>6.1%} {vb:>6.1%} {kb:>8.1%}")
+
+    avg_dA_mse = np.mean([results[r]['accuracy_A'] - refs['cell_j'][r]['A']
+                           for r in RESOLUTIONS])
+    avg_dB_mse = np.mean([results[r]['accuracy_B'] - refs['cell_j'][r]['B']
+                           for r in RESOLUTIONS])
+
+    print(f"\nMean delta vs Cell J (MSE baseline):")
+    print(f"  A (summary):  {avg_dA_mse:+.1%}")
+    print(f"  B (attn):     {avg_dB_mse:+.1%}")
+
+    print()
+    if avg_dA_mse > 0.03 and avg_dB_mse > 0.03:
+        print("✓ PER-PATCH AXIS FEATURES BEAT MSE on both classifiers.")
+        print("  The 256× spatial signal was the missing piece.")
+    elif avg_dA_mse > 0.03 or avg_dB_mse > 0.03:
+        print("~ PER-PATCH AXIS FEATURES BEAT MSE on one classifier.")
+        print("  Mixed result — investigate per-class.")
+    elif abs(avg_dA_mse) < 0.03 and abs(avg_dB_mse) < 0.03:
+        print("= PER-PATCH AXIS FEATURES MATCH MSE.")
+        print("  Comparable performance, axis pipeline now competitive.")
+    else:
+        print("✗ PER-PATCH AXIS FEATURES UNDERPERFORM MSE.")
+        print("  Even with 256× more spatial data, axes lose to MSE here.")
+        print("  Reconstruction error remains the structurally optimal signal")
+        print("  for noise discrimination given how the banks were trained.")
+
+with open(EXP_DIR / 'results_expJ_perpatch.json', 'w') as f:
+    json.dump({
+        'subset_battery_ids': SUBSET_BATTERY_IDS,
+        'subset_phase': SUBSET_PHASE,
+        'n_classes': N_CLASSES,
+        'n_stats': N_STATS,
+        'stat_names': ['max', 'mean', 'std', 'top10_mean', 'entropy'],
+        'codebook_sizes': {bid: codebooks[bid].shape[0]
+                            for bid in SUBSET_BATTERY_IDS},
+        'codebook_calibration': 'mean+mean (per-patch averaging)',
+        'max_axes_padded_to': MAX_AXES,
+        'per_resolution': {
+            str(res): {
+                'accuracy_A': results[res]['accuracy_A'],
+                'accuracy_B': results[res]['accuracy_B'],
+                'per_class_A': results[res]['per_class_A'],
+                'per_class_B': results[res]['per_class_B'],
+            }
+            for res in RESOLUTIONS
+        },
+    }, f, indent=2, default=str)
+print(f"\nSaved results_expJ_perpatch.json")