[FEATURE] add butterworth, notch and peak_eq filters

docs/pyttadoc/functions.rst

@@ -26,4 +26,6 @@ Functions
 .. autofunction:: pytta.plot_bars
 .. autofunction:: pytta.plot_spectrogram
 .. autofunction:: pytta.plot_waterfall
-
+.. autofunction:: pytta.butterworth
+.. autofunction:: pytta.notch
+.. autofunction:: pytta.peak_eq

pytta/__init__.py

@@ -113,7 +113,8 @@
     list_devices, print_devices, get_device_from_user,\
     fft_convolve, find_delay, corr_coef, resample,\
     peak_time, save, load, plot_time, plot_time_dB, plot_freq,\
-    plot_bars, plot_spectrogram, plot_waterfall, SPL
+    plot_bars, plot_spectrogram, plot_waterfall, SPL,\
+    butterworth, notch, peak_eq
 
 # Must go on v0.1.0 (DEPRECATED)
 from .functions import fft_degree
@@ -146,6 +147,9 @@
            'resample',
            'corr_coef',
            'peak_time',
+           'butterworth',
+           'notch',
+           'peak_eq',
            'save',
            'load',
            'weighting',

pytta/functions.py

@@ -23,6 +23,9 @@
     >>> pytta.plot_bars(signalObj1, signalObj2, ..., signalObjN)
     >>> pytta.save(fileName, obj1, ..., objN)
     >>> pytta.load(fileName)
+    >>> pytta.butterworth(signalObj, cutoff)
+    >>> pytta.notch(signalObj, fc)
+    >>> pytta.peak_eq(signalObj, fc, gain)
 
 For further information, check the function specific documentation.
 
@@ -274,6 +277,96 @@ def peak_time(signal):
     else:
         return peaks_time[0]
 
+def _apply_response(signal, response_of):
+    """signal filtered by the complex response response_of(f) at its rfft bins."""
+    x = np.asarray(signal.timeSignal, dtype=float)
+    n = x.shape[0]
+    f = np.fft.rfftfreq(n, 1 / signal.samplingRate)
+    y = np.fft.irfft(np.fft.rfft(x, axis=0) * response_of(f)[:, None], n, axis=0)
+    newSignal = SignalObj(y, 'time', samplingRate=signal.samplingRate,
+                          signalType=signal.signalType)
+    newSignal.channels = cp.deepcopy(signal.channels)
+    return newSignal
+
+
+def butterworth(signal, cutoff, btype='low', order=8, linkwitz_riley=False,
+                zero_phase=False):
+    """
+    signal filtered by a Butterworth or Linkwitz-Riley filter, applied to
+    its spectrum (circular), as ita_butterFR, ita_liriFR and ita_filter_LiRi.
+
+    Parameters
+    ----------
+    signal : SignalObj
+    cutoff : float or (float, float)
+        Cut-off frequency in Hz, or the two of a band-pass.
+    btype : str
+        'low', 'high' or 'band'.
+    order : int
+        Order of the filter; a Linkwitz-Riley filter of order 2k is a
+        Butterworth filter of order k squared. Default 8.
+    linkwitz_riley : bool
+        Linkwitz-Riley instead of Butterworth. Default False.
+    zero_phase : bool
+        Keep only the magnitude of the response. Default False.
+
+    The response is evaluated at the frequencies of the FFT bins of the
+    signal (ita_butterFR samples it with MATLAB freqz(Hd, nBins), from 0 to
+    pi without pi, so its bins sit slightly off the FFT bins).
+
+    >>> lowpassed = pytta.butterworth(signalObj, 1000)
+    """
+    if btype not in ('low', 'high', 'band'):
+        raise ValueError("btype must be 'low', 'high' or 'band'.")
+    k = order // 2 if linkwitz_riley else order
+    if linkwitz_riley and order % 2:
+        raise ValueError('a Linkwitz-Riley filter has an even order.')
+    sos = ss.butter(k, cutoff, btype={'band': 'bandpass'}.get(btype, btype),
+                    fs=signal.samplingRate, output='sos')
+
+    def response(f):
+        _, h = ss.sosfreqz(sos, f, fs=signal.samplingRate)
+        h = h * h if linkwitz_riley else h
+        return np.abs(h) if zero_phase else h
+    return _apply_response(signal, response)
+
+
+def _biquad_response(b, a, fs):
+    def response(f):
+        return ss.freqz(b, a, f, fs=fs)[1]
+    return response
+
+
+def notch(signal, fc, Q=2.):
+    """
+    signal through a second-order notch at fc (Hz) with quality factor Q,
+    applied to its spectrum, as ita_filter_notch (the biquad of the Audio
+    EQ Cookbook of R. Bristow-Johnson).
+
+    >>> pytta.notch(signalObj, 50, Q=10)
+    """
+    w0 = 2 * np.pi * fc / signal.samplingRate
+    alpha = np.sin(w0) / (2 * Q)
+    b = [1, -2 * np.cos(w0), 1]
+    a = [1 + alpha, -2 * np.cos(w0), 1 - alpha]
+    return _apply_response(signal, _biquad_response(b, a, signal.samplingRate))
+
+
+def peak_eq(signal, fc, gain, Q=2.):
+    """
+    signal through a second-order peak (bell) filter at fc (Hz) with gain
+    (dB) and quality factor Q, applied to its spectrum, as ita_filter_peak
+    (bilinear design with K = tan(pi fc / fs) and V = 10^(gain / 20)).
+
+    >>> pytta.peak_eq(signalObj, 1000, 6)
+    """
+    K = np.tan(np.pi * fc / signal.samplingRate)
+    V = 10 ** (gain / 20)
+    b = [1 + V * K / Q + K ** 2, 2 * (K ** 2 - 1), 1 - V * K / Q + K ** 2]
+    a = [1 + K / Q + K ** 2, 2 * (K ** 2 - 1), 1 - K / Q + K ** 2]
+    return _apply_response(signal, _biquad_response(b, a, signal.samplingRate))
+
+
 def fft_degree(*args,**kwargs):
     """
     DEPRECATED

tests/test_filters.py

@@ -0,0 +1,82 @@
+"""
+Tests for the filters of pytta.functions that follow the ITA-Toolbox
+(ita_butterFR, ita_liriFR, ita_filter_LiRi, ita_filter_notch,
+ita_filter_peak). Each test pins one rule of the filter design.
+
+The filters act on the spectrum, so an impulse gives the response itself
+at the FFT bins.
+"""
+import os
+import unittest
+
+import numpy as np
+
+os.environ.setdefault('MPLBACKEND', 'Agg')
+
+import pytta
+
+FS, N = 48000, 48000      # 1 Hz bins
+
+
+def response(sig):
+    return np.fft.rfft(sig.timeSignal[:, 0].astype(float))
+
+
+def impulse():
+    label = 'impulse'
+    x = np.zeros(N)
+    x[0] = 1
+    return pytta.SignalObj(x, 'time', FS)
+
+
+class TestFilters(unittest.TestCase):
+
+    def test_butterworth_is_3_dB_down_at_the_cutoff(self):
+        # |H(f)|^2 = 1 / (1 + (f / fc)^(2n)): 1/2 at fc for every order, and
+        # from 2 fc to 4 fc the level falls by 10 lg((1 + 16^n) / (1 + 4^n))
+        # (fc = 100 Hz, far below Nyquist, where the bilinear transform bends
+        # the response, and above the float32 floor of SignalObj).
+        for order in (2, 4, 8):
+            H = response(pytta.butterworth(impulse(), 1000, order=order))
+            at_fc = 20 * np.log10(np.abs(H[1000]))
+            H = response(pytta.butterworth(impulse(), 100, order=order))
+            slope = 20 * np.log10(np.abs(H[400]) / np.abs(H[200]))
+            print(f'\n  order {order}: {at_fc:.4f} dB at fc, {slope:.2f} dB per octave above', end='')
+            self.assertAlmostEqual(at_fc, -3.0103, delta=1e-3)
+            self.assertAlmostEqual(slope, -10 * np.log10((1 + 16 ** order) / (1 + 4 ** order)), delta=0.02)
+
+    def test_linkwitz_riley_pair_sums_to_an_allpass(self):
+        # A Linkwitz-Riley pair is 6 dB down at the crossover and the low
+        # and high branches sum to a flat magnitude.
+        lo = response(pytta.butterworth(impulse(), 2000, 'low', 4, linkwitz_riley=True))
+        hi = response(pytta.butterworth(impulse(), 2000, 'high', 4, linkwitz_riley=True))
+        flat = np.abs(np.abs(lo + hi)[20:20000] - 1).max()
+        at_fc = 20 * np.log10(np.abs(lo[2000]))
+        print(f'\n  {at_fc:.4f} dB at the crossover, |lo + hi| within {flat:.1e} of 1')
+        self.assertAlmostEqual(at_fc, -6.0206, delta=1e-3)
+        self.assertLess(flat, 1e-6)
+
+    def test_notch_and_peak(self):
+        # The notch is zero at fc and 1 at DC and Nyquist; the peak filter
+        # has its gain at fc and 1 at DC and Nyquist.
+        nH = np.abs(response(pytta.notch(impulse(), 1000, Q=5)))
+        pH = np.abs(response(pytta.peak_eq(impulse(), 3000, 6, Q=2)))
+        print(f'\n  notch: {nH[1000]:.1e} at fc, {nH[0]:.6f} at DC; '
+              f'peak: {20 * np.log10(pH[3000]):.4f} dB at fc, {pH[0]:.6f} at DC')
+        self.assertLess(nH[1000], 1e-6)
+        self.assertAlmostEqual(nH[0], 1, places=6)
+        self.assertAlmostEqual(20 * np.log10(pH[3000]), 6, delta=1e-4)
+        self.assertAlmostEqual(pH[0], 1, places=6)
+        self.assertAlmostEqual(pH[-1], 1, places=6)
+
+    def test_zero_phase_and_rejection(self):
+        H = response(pytta.butterworth(impulse(), 1000, zero_phase=True))
+        self.assertLess(np.abs(H.imag).max(), 1e-6)
+        with self.assertRaises(ValueError):
+            pytta.butterworth(impulse(), 1000, 'notch')
+        with self.assertRaises(ValueError):
+            pytta.butterworth(impulse(), 1000, order=3, linkwitz_riley=True)
+
+
+if __name__ == '__main__':
+    unittest.main()