Create svae_cadence.py

svae_cadence.py

@@ -0,0 +1,517 @@
+"""
+SVAE v2 Conduit Trainer — Prototype
+=====================================
+Train PatchSVAEv2 from random init on noise.
+The decoder MUST reconstruct from decomposed spectral + conduit bundles.
+No M_hat shortcut. Every conduit element is load-bearing.
+
+Readouts per epoch:
+  Standard: MSE, S profile, erank, s_delta, CV
+  Conduit:  friction stats, settle distribution, char_coeff profile,
+            per-mode reconstruction contribution
+
+Should converge rapidly or fail — we'll know within 10 epochs.
+
+Usage:
+    python train_v2_conduit.py
+"""
+
+import os
+import math
+import time
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+try:
+    from google.colab import userdata
+    os.environ["HF_TOKEN"] = userdata.get('HF_TOKEN')
+    from huggingface_hub import login
+    login(token=os.environ["HF_TOKEN"])
+except Exception:
+    pass
+
+from geolip_svae.model import cv_of, extract_patches, stitch_patches
+from geolip_svae.model_v2 import PatchSVAEv2
+
+
+# ═══════════════════════════════════════════════════════════════
+# CONFIG
+# ═══════════════════════════════════════════════════════════════
+
+HF_REPO = 'AbstractPhil/geolip-SVAE'
+VERSION = 'v51_conduit_proto'
+LOCAL_DIR = f'/content/{VERSION}_checkpoints'
+LOG_PATH = os.path.join(LOCAL_DIR, 'training_log.json')
+
+CFG = dict(
+    # Architecture (inherited from Fresnel v50)
+    V=48, D=4, ps=4, hidden=384, depth=4, n_cross=2,
+
+    # Training
+    img_size=64,
+    batch_size=256,
+    lr=3e-4,
+    epochs=50,
+    ds_size=1280000,
+    val_size=10000,
+
+    # CV soft hand
+    target_cv=0.125,
+    cv_weight=0.3,
+    boost=0.5,
+    sigma=0.15,
+
+    # Checkpointing
+    save_every=5,
+    val_per_type_every=5,
+)
+
+
+# ═══════════════════════════════════════════════════════════════
+# NOISE DATASET (16 types, same as Freckles)
+# ═══════════════════════════════════════════════════════════════
+
+NOISE_NAMES = {
+    0: 'gaussian', 1: 'uniform', 2: 'uniform_sc', 3: 'poisson',
+    4: 'pink', 5: 'brown', 6: 'salt_pepper', 7: 'sparse',
+    8: 'block', 9: 'gradient', 10: 'checker', 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 _gen_noise(noise_type, s, rng):
+    if noise_type == 0: return torch.randn(3, s, s)
+    elif noise_type == 1: return torch.rand(3, s, s) * 2 - 1
+    elif noise_type == 2: return (torch.rand(3, s, s) - 0.5) * 4
+    elif noise_type == 3:
+        lam = rng.uniform(0.5, 20.0)
+        return torch.poisson(torch.full((3, s, s), lam)) / lam - 1.0
+    elif noise_type == 4:
+        img = _pink((3, s, s)); return img / (img.std() + 1e-8)
+    elif noise_type == 5:
+        img = _brown((3, s, s)); return img / (img.std() + 1e-8)
+    elif noise_type == 6:
+        return torch.where(torch.rand(3, s, s) > 0.5,
+                          torch.ones(3, s, s) * 2, -torch.ones(3, s, s) * 2) + torch.randn(3, s, s) * 0.1
+    elif noise_type == 7:
+        return torch.randn(3, s, s) * (torch.rand(3, s, s) > 0.9).float() * 3
+    elif noise_type == 8:
+        b = rng.randint(2, max(3, s // 2))
+        sm = torch.randn(3, s // b + 1, s // b + 1)
+        return 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)
+        return (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, max(3, s // 2))
+        cy = torch.arange(s) // cs; cx = torch.arange(s) // cs
+        return ((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)
+        return 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))
+        return img
+    elif noise_type == 13:
+        return torch.tan(math.pi * (torch.rand(3, s, s) - 0.5)).clamp(-3, 3)
+    elif noise_type == 14:
+        return torch.empty(3, s, s).exponential_(1.0) - 1.0
+    elif noise_type == 15:
+        u = torch.rand(3, s, s) - 0.5
+        return -torch.sign(u) * torch.log1p(-2 * u.abs())
+    return torch.randn(3, s, s)
+
+
+class OmegaNoiseDataset(torch.utils.data.Dataset):
+    def __init__(self, size=1280000, img_size=64):
+        self.size = size
+        self.img_size = img_size
+        self._rng = np.random.RandomState(42)
+        self._call_count = 0
+    def __len__(self):
+        return self.size
+    def __getitem__(self, idx):
+        self._call_count += 1
+        if self._call_count % 1000 == 0:
+            self._rng = np.random.RandomState(int.from_bytes(os.urandom(4), 'big'))
+            torch.manual_seed(int.from_bytes(os.urandom(4), 'big'))
+        noise_type = idx % 16
+        img = _gen_noise(noise_type, self.img_size, self._rng).clamp(-4, 4)
+        return img.float(), noise_type
+
+
+def eval_per_type(model, img_size, device, n_per=32):
+    rng = np.random.RandomState(99)
+    model.eval()
+    results = {}
+    with torch.no_grad():
+        for t in range(16):
+            imgs = torch.stack([_gen_noise(t, img_size, rng).clamp(-4, 4)
+                                for _ in range(n_per)]).to(device)
+            out = model(imgs)
+            results[t] = F.mse_loss(out['recon'], imgs).item()
+    return results
+
+
+# ═══════════════════════════════════════════════════════════════
+# CONDUIT READOUTS
+# ═══════════════════════════════════════════════════════════════
+
+def conduit_readout(model, images, device):
+    """Extract and summarize conduit telemetry from a batch."""
+    model.eval()
+    with torch.no_grad():
+        out = model(images.to(device))
+        packet = model.last_conduit_packet
+
+        S = out['svd']['S_orig']   # (B, N, D)
+        B, N, D = S.shape
+
+        friction = packet.friction.reshape(B, N, D)
+        settle = packet.settle.reshape(B, N, D)
+        char_coeffs = packet.char_coeffs.reshape(B, N, D)
+        ext_order = packet.extraction_order.reshape(B, N, D)
+        refine_res = packet.refinement_residual.reshape(B, N)
+
+        # Log-friction for readable stats
+        log_fric = torch.log1p(friction)
+
+        stats = {
+            'S_mean': S.mean(dim=(0, 1)).cpu().tolist(),
+            'S_std': S.std(dim=(0, 1)).cpu().tolist(),
+            'friction_mean': friction.mean().item(),
+            'friction_max': friction.max().item(),
+            'friction_std': friction.std().item(),
+            'log_fric_mean': log_fric.mean(dim=(0, 1)).cpu().tolist(),
+            'log_fric_std': log_fric.std(dim=(0, 1)).cpu().tolist(),
+            'settle_mean': settle.mean(dim=(0, 1)).cpu().tolist(),
+            'settle_frac_gt2': (settle > 2).float().mean().item(),
+            'char_coeffs_mean': char_coeffs.mean(dim=(0, 1)).cpu().tolist(),
+            'refine_res_mean': refine_res.mean().item(),
+            'refine_res_max': refine_res.max().item(),
+        }
+
+        # Per-mode friction spatial CV
+        for d in range(D):
+            per_img = friction[:, :, d].reshape(B, -1)
+            cvs = per_img.std(dim=1) / (per_img.mean(dim=1) + 1e-8)
+            stats[f'friction_spatial_cv_mode{d}'] = cvs.mean().item()
+
+    return stats
+
+
+def print_conduit_readout(stats, D=4):
+    """Pretty-print conduit telemetry."""
+    print(f"    S:         [{', '.join(f'{v:.3f}' for v in stats['S_mean'])}]")
+    print(f"    S_std:     [{', '.join(f'{v:.4f}' for v in stats['S_std'])}]")
+    print(f"    log_fric:  [{', '.join(f'{v:.3f}' for v in stats['log_fric_mean'])}] "
+          f"± [{', '.join(f'{v:.3f}' for v in stats['log_fric_std'])}]")
+    print(f"    fric_raw:  mean={stats['friction_mean']:.1f} max={stats['friction_max']:.0f}")
+    print(f"    settle:    [{', '.join(f'{v:.2f}' for v in stats['settle_mean'])}] "
+          f"(>{2}: {stats['settle_frac_gt2']:.1%})")
+    print(f"    char_c:    [{', '.join(f'{v:.4f}' for v in stats['char_coeffs_mean'])}]")
+    print(f"    refine:    mean={stats['refine_res_mean']:.2e} max={stats['refine_res_max']:.2e}")
+    spatial_cvs = [stats.get(f'friction_spatial_cv_mode{d}', 0) for d in range(D)]
+    print(f"    fric_cv:   [{', '.join(f'{v:.4f}' for v in spatial_cvs)}]")
+
+
+# ═══════════════════════════════════════════════════════════════
+# LOGGING
+# ═══════════════════════════════════════════════════════════════
+
+def load_log():
+    if os.path.exists(LOG_PATH):
+        with open(LOG_PATH) as f:
+            return json.load(f)
+    return {'version': VERSION, 'entries': []}
+
+
+def save_log(log):
+    os.makedirs(os.path.dirname(LOG_PATH), exist_ok=True)
+    with open(LOG_PATH, 'w') as f:
+        json.dump(log, f, indent=2)
+
+
+# ═══════════════════════════════════════════════════════════════
+# SAVE & PUSH
+# ═══════════════════════════════════════════════════════════════
+
+def save_checkpoint(model, opt, sched, epoch, val_mse, log,
+                    path, is_best=False):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    ckpt = {
+        'config': {
+            'V': CFG['V'], 'D': CFG['D'], 'patch_size': CFG['ps'],
+            'hidden': CFG['hidden'], 'depth': CFG['depth'],
+            'n_cross_layers': CFG['n_cross'],
+            'img_size': CFG['img_size'],
+            'model_type': 'v2',
+        },
+        'model_state_dict': model.state_dict(),
+        'optimizer_state_dict': opt.state_dict(),
+        'scheduler_state_dict': sched.state_dict(),
+        'epoch': epoch,
+        'val_mse': val_mse,
+    }
+    torch.save(ckpt, path)
+    size_mb = os.path.getsize(path) / (1024 * 1024)
+    print(f"  💾 {path} ({size_mb:.1f}MB, ep{epoch}, MSE={val_mse:.6f})")
+
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=path,
+            path_in_repo=f'{VERSION}/checkpoints/{os.path.basename(path)}',
+            repo_id=HF_REPO, repo_type='model',
+            commit_message=f'{VERSION} ep{epoch} mse={val_mse:.6f}')
+        if is_best:
+            api.upload_file(
+                path_or_fileobj=path,
+                path_in_repo=f'{VERSION}/checkpoints/best.pt',
+                repo_id=HF_REPO, repo_type='model',
+                commit_message=f'{VERSION} BEST ep{epoch} mse={val_mse:.6f}')
+        save_log(log)
+        api.upload_file(
+            path_or_fileobj=LOG_PATH,
+            path_in_repo=f'{VERSION}/training_log.json',
+            repo_id=HF_REPO, repo_type='model',
+            commit_message=f'{VERSION} log ep{epoch}')
+        print(f"  ☁️  Pushed ep{epoch}")
+    except Exception as e:
+        print(f"  ⚠️  Push failed: {e}")
+
+
+# ═══════════════════════════════════════════════════════════════
+# TRAINING
+# ═══════════════════════════════════════════════════════════════
+
+def train():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    os.makedirs(LOCAL_DIR, exist_ok=True)
+
+    print("\n" + "=" * 70)
+    print(f"SVAE v2 CONDUIT TRAINER — {VERSION}")
+    print("=" * 70)
+
+    # ── Fresh v2 model from random init ──
+    D = CFG['D']
+    model = PatchSVAEv2(
+        V=CFG['V'], D=D, ps=CFG['ps'],
+        hidden=CFG['hidden'], depth=CFG['depth'],
+        n_cross=CFG['n_cross'],
+    ).to(device)
+
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"\n  Fresh PatchSVAEv2 from random init")
+    print(f"  Total params: {n_params:,}")
+
+    # ── Data ──
+    print(f"\n  Dataset: 16 noise types, {CFG['ds_size']:,} samples/epoch")
+    print(f"  Image size: {CFG['img_size']}×{CFG['img_size']}")
+    print(f"  Batch size: {CFG['batch_size']}")
+
+    train_ds = OmegaNoiseDataset(size=CFG['ds_size'], img_size=CFG['img_size'])
+    val_ds = OmegaNoiseDataset(size=CFG['val_size'], img_size=CFG['img_size'])
+    train_loader = torch.utils.data.DataLoader(
+        train_ds, batch_size=CFG['batch_size'], shuffle=True,
+        num_workers=4, pin_memory=True, drop_last=True)
+    val_loader = torch.utils.data.DataLoader(
+        val_ds, batch_size=CFG['batch_size'], shuffle=False,
+        num_workers=4, pin_memory=True)
+
+    # ── Optimizer (all params — full model training) ──
+    opt = torch.optim.Adam(model.parameters(), lr=CFG['lr'])
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=CFG['epochs'])
+
+    # CV soft hand
+    target_cv = CFG['target_cv']
+    cv_weight = CFG['cv_weight']
+    boost = CFG['boost']
+    sigma = CFG['sigma']
+
+    # Log
+    log = load_log()
+    best_mse = float('inf')
+
+    # ── Initial conduit readout ──
+    print(f"\n  Initial conduit profile:")
+    sample_batch = next(iter(val_loader))[0][:64]
+    init_stats = conduit_readout(model, sample_batch, device)
+    print_conduit_readout(init_stats, D)
+
+    # ── Initial MSE (will be terrible — decoder is random) ──
+    model.eval()
+    init_mse = 0
+    init_n = 0
+    with torch.no_grad():
+        for imgs, _ in val_loader:
+            imgs = imgs.to(device)
+            out = model(imgs)
+            init_mse += F.mse_loss(out['recon'], imgs).item() * len(imgs)
+            init_n += len(imgs)
+            if init_n >= 2560:
+                break
+    init_mse /= init_n
+    print(f"\n  Initial MSE (random decoder): {init_mse:.4f}")
+    print("=" * 70)
+
+    # ═══════════════════════════════════════════════════════════
+    # TRAINING LOOP
+    # ═══════════════════════════════════════════════════════════
+
+    for epoch in range(1, CFG['epochs'] + 1):
+
+        model.train()
+        total_loss, total_recon, n = 0, 0, 0
+        last_cv = target_cv
+        t0 = time.time()
+
+        pbar = tqdm(train_loader,
+                    desc=f"Ep {epoch}/{CFG['epochs']}",
+                    bar_format='{l_bar}{bar:20}{r_bar}')
+
+        for batch_idx, (images, _) in enumerate(pbar):
+            images = images.to(device)
+            opt.zero_grad()
+            out = model(images)
+            recon_loss = F.mse_loss(out['recon'], images)
+
+            # CV soft hand
+            with torch.no_grad():
+                if batch_idx % 50 == 0:
+                    cur_cv = cv_of(out['svd']['M'][0, 0])
+                    if cur_cv > 0:
+                        last_cv = cur_cv
+                    delta = last_cv - target_cv
+                    prox = math.exp(-delta**2 / (2 * sigma**2))
+
+            recon_w = 1.0 + boost * prox
+            cv_pen = cv_weight * (1.0 - prox)
+            loss = recon_w * recon_loss + cv_pen * (last_cv - target_cv)**2
+            loss.backward()
+
+            torch.nn.utils.clip_grad_norm_(
+                model.cross_attn.parameters(), max_norm=0.5)
+
+            opt.step()
+
+            total_loss += loss.item() * len(images)
+            total_recon += recon_loss.item() * len(images)
+            n += len(images)
+            pbar.set_postfix_str(f"mse={recon_loss.item():.4f} cv={last_cv:.3f}")
+
+        sched.step()
+        epoch_time = time.time() - t0
+
+        # ── Validation ──
+        model.eval()
+        val_mse, val_n = 0, 0
+        with torch.no_grad():
+            for imgs, _ in val_loader:
+                imgs = imgs.to(device)
+                out = model(imgs)
+                val_mse += F.mse_loss(out['recon'], imgs).item() * len(imgs)
+                val_n += len(imgs)
+        val_mse /= val_n
+
+        # ── Geometry snapshot ──
+        with torch.no_grad():
+            sample = next(iter(val_loader))[0][:64].to(device)
+            out = model(sample)
+            S_mean = out['svd']['S_orig'].mean(dim=(0, 1))
+            S_coord = out['svd']['S'].mean(dim=(0, 1))
+            erank = model.effective_rank(
+                out['svd']['S'].reshape(-1, D)).mean().item()
+            s_delta = model.s_delta(out['svd']['S_orig'], out['svd']['S'])
+
+        # ── Conduit readout ──
+        cond_stats = conduit_readout(model, sample.cpu(), device)
+
+        # ── Print ──
+        is_best = val_mse < best_mse
+        if is_best:
+            best_mse = val_mse
+
+        print(f"\n  ep{epoch:3d} | recon={total_recon/n:.4f} val={val_mse:.4f} "
+              f"{'★ BEST' if is_best else ''} | "
+              f"er={erank:.2f} Sd={s_delta:.4f} cv={last_cv:.3f} | {epoch_time:.0f}s")
+        print_conduit_readout(cond_stats, D)
+
+        # ── Per-type eval ──
+        if epoch % CFG['val_per_type_every'] == 0 or epoch <= 3:
+            type_mse = eval_per_type(model, CFG['img_size'], device)
+            type_str = " ".join(
+                f"{NOISE_NAMES[t][:4]}={v:.3f}" for t, v in sorted(type_mse.items()))
+            print(f"    types: {type_str}")
+
+        # ── Log entry ──
+        log['entries'].append({
+            'epoch': epoch,
+            'train_mse': total_recon / n,
+            'val_mse': val_mse,
+            'cv': last_cv,
+            'erank': erank,
+            's_delta': s_delta,
+            'S_mean': S_mean.cpu().tolist(),
+            'conduit': cond_stats,
+            'epoch_time': epoch_time,
+            'lr': opt.param_groups[0]['lr'],
+        })
+
+        # ── Checkpoint ──
+        if is_best:
+            save_checkpoint(model, opt, sched, epoch, val_mse, log,
+                           os.path.join(LOCAL_DIR, 'best.pt'),
+                           is_best=True)
+
+        if epoch % CFG['save_every'] == 0:
+            save_checkpoint(model, opt, sched, epoch, val_mse, log,
+                           os.path.join(LOCAL_DIR, f'epoch_{epoch:04d}.pt'))
+
+    # ═══════════════════════════════════════════════════════════
+    # DONE
+    # ═══════════════════════════════════════════════════════════
+
+    print(f"\n{'=' * 70}")
+    print(f"v2 CONDUIT TRAINING COMPLETE — {VERSION}")
+    print(f"  Best MSE: {best_mse:.6f}")
+    print(f"  Epochs: {CFG['epochs']}")
+    print(f"  Params: {n_params:,}")
+    print(f"{'=' * 70}")
+
+    return model
+
+
+if __name__ == "__main__":
+    torch.set_float32_matmul_precision('high')
+    train()