Spaces:
Running
Running
| import math | |
| import torch | |
| from typing import List | |
| def butter(fc, fs: float = 2.0): | |
| """ | |
| Recall Butterworth polynomials | |
| N = 1 s + 1 | |
| N = 2 s^2 + sqrt(2s) + 1 | |
| N = 3 (s^2 + s + 1)(s + 1) | |
| N = 4 (s^2 + 0.76536s + 1)(s^2 + 1.84776s + 1) | |
| Scaling | |
| LP to LP: s -> s/w_c | |
| LP to HP: s -> w_c/s | |
| Bilinear transform: | |
| s = 2/T_d * (1 - z^-1)/(1 + z^-1) | |
| For 1-pole butterworth lowpass | |
| 1 / (s + 1) 1-pole prototype | |
| 1 / (s/w_c + 1) LP to LP | |
| 1 / (2/T_d * (1 - z^-1)/(1 + z^-1))/w_c + 1) Bilinear transform | |
| """ | |
| # apply pre-warping to the cutoff | |
| T_d = 1 / fs | |
| w_d = (2 * math.pi * fc) / fs | |
| # sys.exit() | |
| w_c = (2 / T_d) * torch.tan(w_d / 2) | |
| a0 = 2 + (T_d * w_c) | |
| a1 = (T_d * w_c) - 2 | |
| b0 = T_d * w_c | |
| b1 = T_d * w_c | |
| b = torch.stack([b0, b1], dim=0).view(-1) | |
| a = torch.stack([a0, a1], dim=0).view(-1) | |
| # normalize | |
| b = b.type_as(fc) / a0 | |
| a = a.type_as(fc) / a0 | |
| return b, a | |
| def biqaud( | |
| gain_dB: torch.Tensor, | |
| cutoff_freq: torch.Tensor, | |
| q_factor: torch.Tensor, | |
| sample_rate: float, | |
| filter_type: str = "peaking", | |
| ): | |
| # convert inputs to Tensors if needed | |
| # gain_dB = torch.tensor([gain_dB]) | |
| # cutoff_freq = torch.tensor([cutoff_freq]) | |
| # q_factor = torch.tensor([q_factor]) | |
| A = 10 ** (gain_dB / 40.0) | |
| w0 = 2 * math.pi * (cutoff_freq / sample_rate) | |
| alpha = torch.sin(w0) / (2 * q_factor) | |
| cos_w0 = torch.cos(w0) | |
| sqrt_A = torch.sqrt(A) | |
| if filter_type == "high_shelf": | |
| b0 = A * ((A + 1) + (A - 1) * cos_w0 + 2 * sqrt_A * alpha) | |
| b1 = -2 * A * ((A - 1) + (A + 1) * cos_w0) | |
| b2 = A * ((A + 1) + (A - 1) * cos_w0 - 2 * sqrt_A * alpha) | |
| a0 = (A + 1) - (A - 1) * cos_w0 + 2 * sqrt_A * alpha | |
| a1 = 2 * ((A - 1) - (A + 1) * cos_w0) | |
| a2 = (A + 1) - (A - 1) * cos_w0 - 2 * sqrt_A * alpha | |
| elif filter_type == "low_shelf": | |
| b0 = A * ((A + 1) - (A - 1) * cos_w0 + 2 * sqrt_A * alpha) | |
| b1 = 2 * A * ((A - 1) - (A + 1) * cos_w0) | |
| b2 = A * ((A + 1) - (A - 1) * cos_w0 - 2 * sqrt_A * alpha) | |
| a0 = (A + 1) + (A - 1) * cos_w0 + 2 * sqrt_A * alpha | |
| a1 = -2 * ((A - 1) + (A + 1) * cos_w0) | |
| a2 = (A + 1) + (A - 1) * cos_w0 - 2 * sqrt_A * alpha | |
| elif filter_type == "peaking": | |
| b0 = 1 + alpha * A | |
| b1 = -2 * cos_w0 | |
| b2 = 1 - alpha * A | |
| a0 = 1 + (alpha / A) | |
| a1 = -2 * cos_w0 | |
| a2 = 1 - (alpha / A) | |
| else: | |
| raise ValueError(f"Invalid filter_type: {filter_type}.") | |
| b = torch.stack([b0, b1, b2], dim=0).view(-1) | |
| a = torch.stack([a0, a1, a2], dim=0).view(-1) | |
| # normalize | |
| b = b.type_as(gain_dB) / a0 | |
| a = a.type_as(gain_dB) / a0 | |
| return b, a | |
| def freqz(b, a, n_fft: int = 512): | |
| B = torch.fft.rfft(b, n_fft) | |
| A = torch.fft.rfft(a, n_fft) | |
| H = B / A | |
| return H | |
| def freq_domain_filter(x, H, n_fft): | |
| X = torch.fft.rfft(x, n_fft) | |
| # move H to same device as input x | |
| H = H.type_as(X) | |
| Y = X * H | |
| y = torch.fft.irfft(Y, n_fft) | |
| return y | |
| def approx_iir_filter(b, a, x): | |
| """Approimxate the application of an IIR filter. | |
| Args: | |
| b (Tensor): The numerator coefficients. | |
| """ | |
| # round up to nearest power of 2 for FFT | |
| # n_fft = 2 ** math.ceil(math.log2(x.shape[-1] + x.shape[-1] - 1)) | |
| n_fft = 2 ** torch.ceil(torch.log2(torch.tensor(x.shape[-1] + x.shape[-1] - 1))) | |
| n_fft = n_fft.int() | |
| # move coefficients to same device as x | |
| b = b.type_as(x).view(-1) | |
| a = a.type_as(x).view(-1) | |
| # compute complex response | |
| H = freqz(b, a, n_fft=n_fft).view(-1) | |
| # apply filter | |
| y = freq_domain_filter(x, H, n_fft) | |
| # crop | |
| y = y[: x.shape[-1]] | |
| return y | |
| def approx_iir_filter_cascade( | |
| b_s: List[torch.Tensor], | |
| a_s: List[torch.Tensor], | |
| x: torch.Tensor, | |
| ): | |
| """Apply a cascade of IIR filters. | |
| Args: | |
| b (list[Tensor]): List of tensors of shape (3) | |
| a (list[Tensor]): List of tensors of (3) | |
| x (torch.Tensor): 1d Tensor. | |
| """ | |
| if len(b_s) != len(a_s): | |
| raise RuntimeError( | |
| f"Must have same number of coefficients. Got b: {len(b_s)} and a: {len(a_s)}." | |
| ) | |
| # round up to nearest power of 2 for FFT | |
| # n_fft = 2 ** math.ceil(math.log2(x.shape[-1] + x.shape[-1] - 1)) | |
| n_fft = 2 ** torch.ceil(torch.log2(torch.tensor(x.shape[-1] + x.shape[-1] - 1))) | |
| n_fft = n_fft.int() | |
| # this could be done in parallel | |
| b = torch.stack(b_s, dim=0).type_as(x) | |
| a = torch.stack(a_s, dim=0).type_as(x) | |
| H = freqz(b, a, n_fft=n_fft) | |
| H = torch.prod(H, dim=0).view(-1) | |
| # apply filter | |
| y = freq_domain_filter(x, H, n_fft) | |
| # crop | |
| y = y[: x.shape[-1]] | |
| return y | |