| """ |
| 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) |
|
|
|
|
| |
| GatedCompressConfig = BalancedCompressConfig |
| CompressConfig = BalancedCompressConfig |
| VAEConfig = BalancedCompressConfig |
| Config = BalancedCompressConfig |
|
|
| WaveCompressGatedARD = WaveCompressBalancedGatedARD |
| WaveCompress = WaveCompressBalancedGatedARD |
| VAE = WaveCompressBalancedGatedARD |
| Model = WaveCompressBalancedGatedARD |
|
|
