nicktup/reverb-extractor

1"""White-box system identification: fit a DifferentiableFDN to a target RIR.2 3Single-RIR optimisation (no batching), following the paper: Adam with a fairly4large learning rate on the perceptual EDC + soft-EDP loss. Every FDN parameter,5including the delay lengths, is updated by backpropagation.6"""7 8from __future__ import annotations9 10from dataclasses import dataclass, field11import torch12 13from .fdn import DifferentiableFDN, FDNConfig14from .losses import total_loss, EDPConfig15 16 17@dataclass18class TrainConfig:19    iters: int = 50020    lr: float = 0.121    lam: float = 0.1              # weight on the soft-EDP term22    betas: tuple = (0.9, 0.999)23    log_every: int = 2524    fdn: FDNConfig = field(default_factory=FDNConfig)25    edp: EDPConfig = field(default_factory=EDPConfig)26    seed: int = 027 28 29def fit_fdn(30    h_target: torch.Tensor,31    cfg: TrainConfig | None = None,32    verbose: bool = True,33):34    """Fit an FDN to ``h_target``. Returns ``(model, history, h_init, h_final)``."""35    cfg = cfg or TrainConfig()36    torch.manual_seed(cfg.seed)37 38    L = h_target.shape[-1]39    # Ensure the frequency-sampling grid is comfortably longer than the RIR to40    # curb time-domain aliasing of the IIR response.41    n_fft = 142    while n_fft < 2 * L:43        n_fft *= 244    fdn_cfg = FDNConfig(45        n_delays=cfg.fdn.n_delays,46        sample_rate=cfg.fdn.sample_rate,47        n_fft=n_fft,48        max_delay_ms=cfg.fdn.max_delay_ms,49        beta_a=cfg.fdn.beta_a,50        beta_b=cfg.fdn.beta_b,51    )52    edp_cfg = EDPConfig(53        window_ms=cfg.edp.window_ms,54        sample_rate=cfg.fdn.sample_rate,55        kappa_start=cfg.edp.kappa_start,56        kappa_end=cfg.edp.kappa_end,57    )58 59    model = DifferentiableFDN(fdn_cfg).to(h_target.device)60    with torch.no_grad():61        h_init = model.render(L).detach().clone()62 63    opt = torch.optim.Adam(model.parameters(), lr=cfg.lr, betas=cfg.betas)64 65    history = {"total": [], "edc": [], "edp": []}66    for it in range(cfg.iters):67        opt.zero_grad()68        h_pred = model.render(L)69        loss, l_edc, l_edp = total_loss(h_pred, h_target, lam=cfg.lam, edp_cfg=edp_cfg)70        loss.backward()71        opt.step()72 73        history["total"].append(float(loss))74        history["edc"].append(float(l_edc))75        history["edp"].append(float(l_edp))76        if verbose and (it % cfg.log_every == 0 or it == cfg.iters - 1):77            dmin = float(model.delays.min())78            dmax = float(model.delays.max())79            print(80                f"iter {it:4d} | L={float(loss):.5e} "81                f"(EDC={float(l_edc):.5e}, EDP={float(l_edp):.5e}) "82                f"| delays[{dmin:.1f},{dmax:.1f}] samp"83            )84 85    with torch.no_grad():86        h_final = model.render(L).detach().clone()87 88    return model, history, h_init, h_final89