""" VAE.py ====== WaveCompress Balanced Gated-ARD feature extractor for ICWDS Stage-2/Stage-3. Expected checkpoint: best_balanced_compact_4to7dims.pt, renamed to VAE.pt in Colab. Input: parts: (B, K, 47, 72), each slot is mask_k * E in [0, 1] presence: (B, K), optional active-slot gate Output dict: recon: (B, K, 47, 72), decoder logits; apply sigmoid() for [0, 1] mu: (B, K, latent_dim) logvar: (B, K, latent_dim) z: (B, K, latent_dim) z_used: gated latent used by decoder gate: (K, latent_dim) desc: (B, K, 9) physical descriptor prediction This file intentionally keeps aliases used by downstream notebooks: GatedCompressConfig, CompressConfig, VAEConfig, Config WaveCompressGatedARD, WaveCompress, VAE, Model """ import math from dataclasses import dataclass, asdict from typing import Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F @dataclass class BalancedCompressConfig: n_freqs: int = 47 n_dirs: int = 72 pad_freqs: int = 48 n_slots: int = 6 latent_dim: int = 16 enc_width: int = 48 gate_init_logit: float = 1.55 gate_hard_th: float = 0.58 desc_dim: int = 9 def to_dict(self): return asdict(self) class WaveSystemVAECore(nn.Module): def __init__(self, cfg: BalancedCompressConfig): super().__init__() self.cfg = cfg w = cfg.enc_width self.enc = nn.Sequential( nn.Conv2d(1, w, 3, stride=2, padding=1), nn.GroupNorm(min(8, w), w), nn.GELU(), nn.Conv2d(w, w * 2, 3, stride=2, padding=1), nn.GroupNorm(min(8, w * 2), w * 2), nn.GELU(), nn.Conv2d(w * 2, w * 2, 3, stride=2, padding=1), nn.GroupNorm(min(8, w * 2), w * 2), nn.GELU(), ) self.enc_pool = nn.AdaptiveAvgPool2d(1) self.to_stats = nn.Linear(w * 2, 2 * cfg.latent_dim) self.from_lat = nn.Linear(cfg.latent_dim, w * 2 * 6 * 9) self.dec = nn.Sequential( nn.Conv2d(w * 2, w * 8, 3, padding=1), nn.PixelShuffle(2), nn.GroupNorm(min(8, w * 2), w * 2), nn.GELU(), nn.Conv2d(w * 2, w * 4, 3, padding=1), nn.PixelShuffle(2), nn.GroupNorm(min(8, w), w), nn.GELU(), nn.Conv2d(w, w * 4, 3, padding=1), nn.PixelShuffle(2), nn.GroupNorm(min(8, w), w), nn.GELU(), nn.Conv2d(w, 1, 1), ) self._w = w def encode(self, x: torch.Tensor): h = self.enc_pool(self.enc(x)).flatten(1) mu, logvar = self.to_stats(h).chunk(2, dim=-1) return mu, logvar.clamp(-8.0, 5.0) def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor): if self.training: return mu + torch.randn_like(mu) * (0.5 * logvar).exp() return mu def decode(self, z: torch.Tensor): h = self.from_lat(z).view(-1, self._w * 2, 6, 9) y = self.dec(h) return y[:, :, :self.cfg.n_freqs, :] class WaveCompressBalancedGatedARD(nn.Module): """Balanced gated latent VAE for per-wave-system features.""" def __init__(self, cfg: Optional[BalancedCompressConfig] = None): super().__init__() self.cfg = cfg or BalancedCompressConfig() cfg = self.cfg self.vae = WaveSystemVAECore(cfg) self.gate_logits = nn.Parameter(torch.full((cfg.n_slots, cfg.latent_dim), float(cfg.gate_init_logit))) self.register_buffer('dim_cost', torch.linspace(0.75, 1.45, cfg.latent_dim)) self.desc_head = nn.Sequential( nn.Linear(cfg.latent_dim, 64), nn.GELU(), nn.Linear(64, 64), nn.GELU(), nn.Linear(64, cfg.desc_dim), ) def _pad(self, x: torch.Tensor): if self.cfg.pad_freqs > self.cfg.n_freqs: return F.pad(x, (0, 0, 0, self.cfg.pad_freqs - self.cfg.n_freqs)) return x def gates(self, temperature: float = 1.0, hard: bool = False): g = torch.sigmoid(self.gate_logits / max(float(temperature), 1e-4)) if hard: gh = (g > self.cfg.gate_hard_th).float() g = gh.detach() - g.detach() + g return g def ordered_mask(self, B: int, K: int, keep_min: int, keep_max: int, device): D = self.cfg.latent_dim lo = int(max(1, min(D, round(keep_min)))) hi = int(max(lo, min(D, round(keep_max)))) keep = torch.randint(lo, hi + 1, (B, K, 1), device=device) dim = torch.arange(D, device=device).view(1, 1, D) return (dim < keep).float() def forward( self, parts: torch.Tensor, presence: Optional[torch.Tensor] = None, gate_temperature: float = 1.0, ordered_keep: Optional[Tuple[int, int]] = None, hard_gate: bool = False, ): B, K, nf, nd = parts.shape D = self.cfg.latent_dim x = parts.reshape(B * K, 1, nf, nd) mu, logvar = self.vae.encode(self._pad(x)) z = self.vae.reparameterize(mu, logvar) mu = mu.view(B, K, D) logvar = logvar.view(B, K, D) z = z.view(B, K, D) gate = self.gates(gate_temperature, hard=hard_gate)[:K, :].unsqueeze(0) z_used = z * gate if self.training and ordered_keep is not None: z_used = z_used * self.ordered_mask(B, K, ordered_keep[0], ordered_keep[1], parts.device) if presence is not None: z_used = z_used * presence.float().clamp(0, 1).unsqueeze(-1) recon = self.vae.decode(z_used.reshape(B * K, D)).view(B, K, nf, nd) desc = self.desc_head(z_used.reshape(B * K, D)).view(B, K, self.cfg.desc_dim) return {'recon': recon, 'mu': mu, 'logvar': logvar, 'z': z, 'z_used': z_used, 'gate': gate.squeeze(0), 'desc': desc} @torch.no_grad() def gate_stats(self, temperature: float = 1.0): g = self.gates(temperature, hard=False) hard = (g > self.cfg.gate_hard_th).float() return { 'gate_mean': float(g.mean().detach().cpu()), 'gate_soft_dim_mean': float(g.sum(dim=1).mean().detach().cpu()), 'gate_hard_dim_mean': float(hard.sum(dim=1).mean().detach().cpu()), 'gate_soft_per_slot': g.sum(dim=1).detach().cpu().numpy(), 'gate_hard_per_slot': hard.sum(dim=1).detach().cpu().numpy(), 'gate_per_slot': g.detach().cpu().numpy(), } def num_params(self): return sum(p.numel() for p in self.parameters()) @torch.no_grad() def _safe_presence(presence: torch.Tensor): return presence.float().clamp(0, 1) def part_descriptors(parts: torch.Tensor): """Differentiable descriptors: log_mass, peak_f, sin_dir, cos_dir, spread_f, spread_dir, peak_val, area, compactness.""" B, K, nf, nd = parts.shape device = parts.device eps = 1e-8 E = parts.clamp_min(0.0) mass = E.sum(dim=(2, 3)).clamp_min(eps) fcoord = torch.linspace(0, 1, nf, device=device).view(1, 1, nf, 1) theta = torch.linspace(0, 2 * math.pi, nd + 1, device=device)[:nd].view(1, 1, 1, nd) log_mass = torch.log1p(mass) / math.log1p(float(nf * nd)) f_mean = (E * fcoord).sum(dim=(2, 3)) / mass f_var = (E * (fcoord - f_mean[:, :, None, None]).pow(2)).sum(dim=(2, 3)) / mass spread_f = torch.sqrt(f_var.clamp_min(0.0)) sx = (E * torch.sin(theta)).sum(dim=(2, 3)) / mass cx = (E * torch.cos(theta)).sum(dim=(2, 3)) / mass R = torch.sqrt(sx * sx + cx * cx).clamp(0, 1) spread_dir = torch.sqrt(torch.clamp(1.0 - R, min=0.0)) flat = E.reshape(B, K, nf * nd) peak_idx = flat.argmax(dim=-1) peak_f = (peak_idx // nd).float() / max(nf - 1, 1) peak_d = (peak_idx % nd).float() / nd * 2 * math.pi peak_sin = torch.sin(peak_d) peak_cos = torch.cos(peak_d) peak_val = flat.max(dim=-1).values.clamp(0, 1) area = (E > 1e-4).float().mean(dim=(2, 3)) compactness = peak_val / (mass / float(nf * nd) + peak_val + eps) return torch.stack([log_mass, peak_f, peak_sin, peak_cos, spread_f, spread_dir, peak_val, area, compactness], dim=-1) # Aliases expected by the Stage-3 notebook and older code. GatedCompressConfig = BalancedCompressConfig CompressConfig = BalancedCompressConfig VAEConfig = BalancedCompressConfig Config = BalancedCompressConfig WaveCompressGatedARD = WaveCompressBalancedGatedARD WaveCompress = WaveCompressBalancedGatedARD VAE = WaveCompressBalancedGatedARD Model = WaveCompressBalancedGatedARD