[FEATURE] add harmonic_distortion for sweep measurements

pytta/__init__.py

@@ -113,7 +113,7 @@
     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, harmonic_distortion
 
 # Must go on v0.1.0 (DEPRECATED)
 from .functions import fft_degree
@@ -146,6 +146,7 @@
            'resample',
            'corr_coef',
            'peak_time',
+           'harmonic_distortion',
            'save',
            'load',
            'weighting',

pytta/functions.py

@@ -17,6 +17,7 @@
     >>> pytta.corr_coef(signalObj1, signalObj2)
     >>> pytta.resample(signalObj, newSamplingRate)
     >>> pytta.peak_time(signalObj1, signalObj2, ..., signalObjN)
+    >>> pytta.harmonic_distortion(ir, sweep)
     >>> pytta.plot_time(signalObj1, signalObj2, ..., signalObjN)
     >>> pytta.plot_time_dB(signalObj1, signalObj2, ..., signalObjN)
     >>> pytta.plot_freq(signalObj1, signalObj2, ..., signalObjN)
@@ -274,6 +275,82 @@ def peak_time(signal):
     else:
         return peaks_time[0]
 
+def harmonic_distortion(ir, sweep, numHarmonics=5):
+    """
+    Harmonic impulse responses and total harmonic distortion (THD) from the
+    impulse response of a system measured with an exponential sweep.
+
+    Deconvolution by an exponential sweep puts the harmonic of order k of a
+    nonlinear system L ln(k) seconds before the linear response, where L is
+    the time the sweep takes to multiply its frequency by e (Farina, AES
+    108th Convention, 2000). Each harmonic is cut with the same window: its
+    length is the gap between harmonics numHarmonics and numHarmonics + 1,
+    the shortest one, 10 % of it before the peak of the harmonic, with a
+    Tukey taper of 10 % at each end.
+
+    Parameters (default), (type):
+    -----------------------------
+        * ir (), (SignalObj):
+            impulse response with one channel, from a circular
+            deconvolution, so that the harmonics wrap to its end;
+        * sweep (), (SignalObj or float):
+            the exponential sweep used as excitation, whose rate is read
+            from its group delay between 2 freqMin and freqMax / 2, or its
+            rate in octaves per second;
+        * numHarmonics (5), (int):
+            highest harmonic order, 2 or more.
+
+    Return:
+    -------
+        * harmonics: SignalObj with one channel per order, from 1 (the
+          linear response) to numHarmonics;
+        * f: excitation frequencies, in Hz, up to fs / (2 numHarmonics);
+        * thd: total harmonic distortion at each frequency of f, as a
+          ratio: sqrt(sum |H_k(k f)| ** 2 for k >= 2) / |H_1(f)|.
+
+    >>> harmonics, f, thd = pytta.harmonic_distortion(ir, sweep, 3)
+    """
+    if not isinstance(ir, SignalObj):
+        raise TypeError('ir must be a SignalObj.')
+    if ir.numChannels != 1:
+        raise ValueError('ir must have one channel; see SignalObj.split.')
+    if numHarmonics < 2:
+        raise ValueError('numHarmonics must be 2 or more.')
+    rate = _sweep_rate(sweep) if isinstance(sweep, SignalObj) \
+        else float(sweep)
+    fs = ir.samplingRate
+    L = 1 / (rate * np.log(2))  # s for the frequency to grow by e
+    M = int(L * np.log((numHarmonics + 1) / numHarmonics) * fs)
+    h = np.asarray(ir.timeSignal[:, 0], dtype=float)
+    orders = np.arange(1, numHarmonics + 1)
+    # ponytail: positions rounded to the sample; enough for magnitudes
+    starts = np.argmax(np.abs(h)) - M // 10 \
+        - np.round(L * np.log(orders) * fs).astype(int)
+    harmonics = np.take(h, starts[None, :] + np.arange(M)[:, None],
+                        mode='wrap') * ss.windows.tukey(M, 0.2)[:, None]
+    H = np.abs(np.fft.rfft(harmonics, axis=0))
+    i = np.arange((M // 2) // numHarmonics + 1)
+    distortion = np.sqrt(sum(H[k * i, k - 1] ** 2
+                             for k in orders[1:]))
+    with np.errstate(divide='ignore', invalid='ignore'):
+        thd = distortion / H[i, 0]
+    harmonics = SignalObj(harmonics, 'time', samplingRate=fs,
+                          signalType='energy')
+    return harmonics, i * fs / M, thd
+
+
+def _sweep_rate(sweep):
+    """Rate of an exponential sweep in octaves per second, from the slope of
+    its group delay over the octaves between 2 freqMin and freqMax / 2."""
+    x = np.asarray(sweep.timeSignal[:, 0], dtype=float)
+    f = np.fft.rfftfreq(x.size, 1 / sweep.samplingRate)
+    # group delay with no phase unwrapping: Re{FFT(n x) / FFT(x)} / fs
+    delay = np.real(np.fft.rfft(np.arange(x.size) * x)
+                    / np.fft.rfft(x)) / sweep.samplingRate
+    band = (f >= 2 * sweep.freqMin) & (f <= sweep.freqMax / 2)
+    return 1 / np.polyfit(np.log2(f[band]), delay[band], 1)[0]
+
+
 def fft_degree(*args,**kwargs):
     """
     DEPRECATED

tests/test_harmonic_distortion.py

@@ -0,0 +1,98 @@
+"""
+Tests for pytta.harmonic_distortion.
+
+An exponential sweep separates the harmonics of a nonlinear system in the
+impulse response obtained by deconvolution: the harmonic of order k lands
+L ln(k) seconds before the linear response, where L is the time the sweep
+takes to multiply its frequency by e (Farina, AES 108th Convention, 2000).
+
+The tests use a memoryless system y = x + a x**2 + b x**3 driven by a sweep
+of amplitude A, whose harmonics are known: the linear response has gain
+1 + 3 b A**2 / 4, the second harmonic a A / 2 and the third b A**2 / 4,
+relative to the sweep. Each test pins one rule.
+"""
+import os
+import unittest
+
+os.environ.setdefault('MPLBACKEND', 'Agg')
+
+import numpy as np
+import scipy.signal as ss
+
+import pytta
+
+FS = 48000
+F1, F2, T = 20.0, 7000.0, 2.0     # 3 x F2 stays below Nyquist
+A = 0.5
+RATE = np.log2(F2 / F1) / T       # octaves per second
+
+
+def sweep_and_ir(system):
+    """Sweep of amplitude A, 0.5 s of silence after it, and the impulse
+    response of `system` by regularized circular deconvolution."""
+    t = np.arange(int(T * FS)) / FS
+    x = np.concatenate((A * ss.chirp(t, F1, T, F2, 'logarithmic', phi=-90),
+                        np.zeros(FS // 2)))
+    X, Y = np.fft.rfft(x), np.fft.rfft(system(x))
+    h = np.fft.irfft(Y * np.conj(X) / (np.abs(X) ** 2 + 1e-8 * np.max(
+        np.abs(X) ** 2)), n=x.size)
+    return x / A, h
+
+
+class TestHarmonicDistortion(unittest.TestCase):
+
+    def test_harmonic_levels_of_a_known_nonlinearity(self):
+        a, b = 0.1, 0.2
+        x, h = sweep_and_ir(lambda x: x + a * x ** 2 + b * x ** 3)
+        ir = pytta.SignalObj(h, 'time', FS)
+        harmonics, f, thd = pytta.harmonic_distortion(ir, RATE,
+                                                      numHarmonics=3)
+        band = (f >= 100) & (f <= 2000)
+        H = np.abs(np.fft.rfft(harmonics.timeSignal, axis=0))
+        i = np.flatnonzero(band)
+        gains = [np.median(H[i, 0]), np.median(H[2 * i, 1] / H[i, 0]),
+                 np.median(H[3 * i, 2] / H[i, 0])]
+        H1 = 1 + 3 * b * A ** 2 / 4
+        expected = [H1, a * A / 2 / H1, b * A ** 2 / 4 / H1]
+        for k, (got, exp) in enumerate(zip(gains, expected), start=1):
+            print(f'\n  harmonic {k}: {got:.5f} (expected {exp:.5f})')
+            self.assertAlmostEqual(got / exp, 1, delta=0.02)
+        THD = np.hypot(expected[1], expected[2])
+        worst = np.max(np.abs(thd[band] / THD - 1))
+        print(f'\n  THD {THD:.5f}, worst relative error 100 Hz to 2 kHz '
+              f'{worst:.1e}')
+        self.assertLess(worst, 0.05)
+
+    def test_linear_system_shows_no_distortion(self):
+        # A resonance whose impulse response decays well before the
+        # windows of the harmonics: no energy may leak into them.
+        b, a = ss.iirpeak(1000, 5, fs=FS)
+        x, h = sweep_and_ir(lambda x: ss.lfilter(b, a, x))
+        ir = pytta.SignalObj(h, 'time', FS)
+        _, f, thd = pytta.harmonic_distortion(ir, RATE, numHarmonics=3)
+        band = (f >= 100) & (f <= 2000)
+        print(f'\n  largest THD 100 Hz to 2 kHz: {thd[band].max():.1e}')
+        self.assertLess(thd[band].max(), 1e-3)
+
+    def test_rate_is_read_from_the_sweep(self):
+        # Passing the sweep instead of its rate gives the same result.
+        x, h = sweep_and_ir(lambda x: x + 0.1 * x ** 2)
+        ir = pytta.SignalObj(h, 'time', FS)
+        sweep = pytta.SignalObj(x, 'time', FS, freqMin=F1, freqMax=F2)
+        _, _, fromRate = pytta.harmonic_distortion(ir, RATE, 2)
+        _, _, fromSweep = pytta.harmonic_distortion(ir, sweep, 2)
+        self.assertTrue(np.allclose(fromSweep, fromRate, rtol=1e-3))
+
+    def test_rejects_invalid_arguments(self):
+        ir = pytta.SignalObj(np.zeros((FS, 2)), 'time', FS)
+        with self.assertRaises(ValueError):
+            pytta.harmonic_distortion(ir, RATE)            # two channels
+        ir = pytta.SignalObj(np.zeros(FS), 'time', FS)
+        with self.assertRaises(ValueError):
+            pytta.harmonic_distortion(ir, RATE, numHarmonics=1)
+        with self.assertRaises(TypeError):
+            pytta.harmonic_distortion(np.zeros(FS), RATE)
+
+
+if __name__ == '__main__':
+    unittest.main()