nicktup/reverb-extractor

1"""The differentiable Feedback Delay Network (FDN).2 3State-space description of the single-input single-output FDN (paper Sec. 2)::4 5    y[n]        = c^T s[n] + d u[n]6    s_i[n]      = ( b_i u[n] + sum_j A_ij s_j[n] )  delayed by m_i samples7 8with feedback matrix ``A = U @ diag(gamma)``. The transfer function is9 10    H(z) = c^T ( I - D(z) A )^{-1} D(z) b + d ,   D(z) = diag(z^{-m_i}) .11 12We render the RIR by *frequency sampling*: evaluate ``H`` on the DFT grid using13the differentiable fractional-delay response ``D_i[k]`` from ``delay.py`` and14take the inverse real FFT. Every parameter -- input/output gains ``b, c``,15direct gain ``d``, feedback matrix ``A`` and, crucially, the delay lengths16``m`` -- is a differentiable function of an unconstrained proxy, so the whole17network is trainable by backpropagation.18 19Reparameterisations (paper Sec. 3-4):20    * delays   m   = clamp(|m_tilde|, 0, Q-1)         (nonnegative, causal)21    * orthog.  U   = expm(triu(W) - triu(W)^T)        (unilossless feedback)22    * absorp.  gamma = sigmoid(gamma_tilde) in (0,1)  (decoupled from m)23    * gains    b,c,d = |.|                             (nonnegative)24"""25 26from __future__ import annotations27 28from dataclasses import dataclass29import math30import torch31import torch.nn as nn32 33from .delay import fractional_delay_frequency_response34 35 36@dataclass37class FDNConfig:38    n_delays: int = 6          # N, number of delay lines39    sample_rate: int = 16000   # fs (Hz)40    n_fft: int = 16384         # K, frequency-sampling grid (even; also sets Q=K/2)41    max_delay_ms: float = 64.0  # init scale psi -> ~64 ms max (paper)42    beta_a: float = 1.1        # Beta(alpha, beta) delay init43    beta_b: float = 6.044 45 46class DifferentiableFDN(nn.Module):47    """A fully differentiable FDN whose delay lines are learnable."""48 49    def __init__(self, config: FDNConfig | None = None):50        super().__init__()51        self.cfg = config or FDNConfig()52        N = self.cfg.n_delays53        self.Q = self.cfg.n_fft // 2  # buffer length; delays bounded by Q-154 55        # --- delay lines: proxy m_tilde, m = clamp(|m_tilde|, 0, Q-1) ---56        psi = self.cfg.max_delay_ms * 1e-3 * self.cfg.sample_rate  # samples at max57        beta = torch.distributions.Beta(self.cfg.beta_a, self.cfg.beta_b)58        m0 = psi * beta.sample((N,))  # mean ~ psi * a/(a+b) ~ 10 ms59        self.m_tilde = nn.Parameter(m0)60 61        # --- orthogonal feedback matrix via skew-symmetric exponential map ---62        self.W = nn.Parameter(torch.randn(N, N) / math.sqrt(N))63 64        # --- absorption coefficients gamma = sigmoid(gamma_tilde) in (0,1) ---65        # init ~0.9 (light per-hop loss -> audible decay). logit(0.9) ~ 2.266        self.gamma_tilde = nn.Parameter(torch.full((N,), 2.2))67 68        # --- input / output / direct gains (nonnegative via abs) ---69        self.b_tilde = nn.Parameter(torch.randn(N) / math.sqrt(N))70        self.c_tilde = nn.Parameter(torch.full((N,), 1.0 / N))71        self.d_tilde = nn.Parameter(torch.tensor(1.0))72 73    # ------------------------------------------------------------------ #74    # Constrained parameters (differentiable functions of the proxies)   #75    # ------------------------------------------------------------------ #76    @property77    def delays(self) -> torch.Tensor:78        """Delay lengths m in samples, in [0, Q-1]."""79        return torch.clamp(torch.abs(self.m_tilde), 0.0, self.Q - 1)80 81    @property82    def U(self) -> torch.Tensor:83        """Orthogonal (unilossless) matrix via matrix exponential of a skew form."""84        tri = torch.triu(self.W, diagonal=1)85        skew = tri - tri.transpose(-1, -2)86        return torch.matrix_exp(skew)87 88    @property89    def gamma(self) -> torch.Tensor:90        return torch.sigmoid(self.gamma_tilde)91 92    @property93    def feedback_matrix(self) -> torch.Tensor:94        """A = U @ diag(gamma). Spectral radius < 1 => stable."""95        return self.U * self.gamma.unsqueeze(0)  # column scaling == U @ diag(gamma)96 97    @property98    def b(self) -> torch.Tensor:99        return torch.abs(self.b_tilde)100 101    @property102    def c(self) -> torch.Tensor:103        return torch.abs(self.c_tilde)104 105    @property106    def d(self) -> torch.Tensor:107        return torch.abs(self.d_tilde)108 109    # ------------------------------------------------------------------ #110    # Rendering                                                          #111    # ------------------------------------------------------------------ #112    def transfer_function(self, n_fft: int | None = None) -> torch.Tensor:113        """Evaluate H on the rfft grid. Returns complex tensor (n_fft//2+1,)."""114        K = n_fft or self.cfg.n_fft115        m = self.delays116        D = fractional_delay_frequency_response(m, K)  # (N, n_bins) complex117        D = D.transpose(0, 1)                          # (n_bins, N)118        n_bins = D.shape[0]119        N = self.cfg.n_delays120 121        A = self.feedback_matrix.to(D.dtype)           # (N, N) complex122        I = torch.eye(N, dtype=D.dtype, device=D.device)123 124        # M[k] = I - diag(D[k]) @ A  ==  I - D[k, :, None] * A125        M = I.unsqueeze(0) - D.unsqueeze(-1) * A.unsqueeze(0)  # (n_bins, N, N)126        rhs = (D * self.b.to(D.dtype)).unsqueeze(-1)           # (n_bins, N, 1)127        x = torch.linalg.solve(M, rhs).squeeze(-1)            # (n_bins, N)128        H = (x * self.c.to(D.dtype)).sum(-1) + self.d.to(D.dtype)  # (n_bins,)129        return H130 131    def render(self, length: int, n_fft: int | None = None) -> torch.Tensor:132        """Render the RIR (real, length ``length``)."""133        K = n_fft or self.cfg.n_fft134        if K < length:135            raise ValueError(f"n_fft ({K}) must be >= length ({length})")136        H = self.transfer_function(K)137        h = torch.fft.irfft(H, n=K)138        return h[:length]139 140    def forward(self, length: int) -> torch.Tensor:141        return self.render(length)142