[FEATURE] add sine_distortion and snr for measurements with a sine

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

pytta/functions.py

@@ -17,6 +17,8 @@
     >>> pytta.corr_coef(signalObj1, signalObj2)
     >>> pytta.resample(signalObj, newSamplingRate)
     >>> pytta.peak_time(signalObj1, signalObj2, ..., signalObjN)
+    >>> pytta.sine_distortion(signalObj, freq)
+    >>> pytta.snr(signalObj, noiseObj)
     >>> pytta.plot_time(signalObj1, signalObj2, ..., signalObjN)
     >>> pytta.plot_time_dB(signalObj1, signalObj2, ..., signalObjN)
     >>> pytta.plot_freq(signalObj1, signalObj2, ..., signalObjN)
@@ -274,6 +276,54 @@ def peak_time(signal):
     else:
         return peaks_time[0]
 
+def sine_distortion(signalObj, freq, numHarmonics=5):
+    """
+    Total harmonic distortion (THD) and THD plus noise (THD+N) of a
+    recorded sine, by channel, as ratios.
+
+    The power of each band comes from a Blackman-Harris periodogram,
+    integrated over 6 bins on each side of the fundamental and of its
+    harmonics up to numHarmonics (and below Nyquist):
+
+        * THD: rms of the harmonics re the rms of the fundamental;
+        * THD+N: rms of everything except the fundamental and DC, re the
+          total rms (as in AES17).
+
+    The recording should hold a whole number of cycles or be long enough
+    for the window to make that irrelevant (a second or more).
+
+    >>> thd, thdn = pytta.sine_distortion(recording, 1000)
+    """
+    fs = signalObj.samplingRate
+    if not 0 < freq < fs / 2:
+        raise ValueError('freq must lie between 0 and fs / 2.')
+    f, psd = ss.periodogram(signalObj.timeSignal, fs, window='blackmanharris',
+                            scaling='density', axis=0)
+    df = f[1]
+    half = 6 * df
+
+    def band(center):
+        return np.sum(psd[np.abs(f - center) <= half], axis=0) * df
+
+    fundamental = band(freq)
+    harmonics = sum(band(k * freq) for k in range(2, numHarmonics + 1)
+                    if k * freq + half < fs / 2)
+    total = np.sum(psd, axis=0) * df
+    rest = total - band(0) - fundamental
+    return np.sqrt(harmonics / fundamental), np.sqrt(rest / total)
+
+
+def snr(signal, noise):
+    """
+    Signal to noise ratio in dB, by channel: the rms of a recording of the
+    signal re the rms of a recording of the noise alone.
+
+    >>> pytta.snr(recording, backgroundNoise)
+    """
+    return 20 * np.log10(np.std(signal.timeSignal, axis=0)
+                         / np.std(noise.timeSignal, axis=0))
+
+
 def fft_degree(*args,**kwargs):
     """
     DEPRECATED

tests/test_sine_distortion.py

@@ -0,0 +1,64 @@
+"""
+Tests for pytta.sine_distortion and pytta.snr, for measurements with a
+sine, such as those of line inputs and outputs. Each test pins one rule on
+signals whose distortion and noise are known.
+"""
+import os
+import unittest
+
+os.environ.setdefault('MPLBACKEND', 'Agg')
+
+import numpy as np
+
+import pytta
+
+FS = 48000
+F0 = 997.0                 # off the bin grid, as test signals usually are
+tt = np.arange(FS) / FS
+
+
+class TestSineDistortion(unittest.TestCase):
+
+    def test_harmonics_give_the_thd(self):
+        # Second and third harmonics at 1 % and 0.5 % of the fundamental:
+        # THD = sqrt(0.01**2 + 0.005**2) re the fundamental, and THD+N the
+        # same rms re the total rms.
+        x = np.sin(2 * np.pi * F0 * tt) + 0.01 * np.sin(4 * np.pi * F0 * tt) \
+            + 0.005 * np.sin(6 * np.pi * F0 * tt)
+        thd, thdn = pytta.sine_distortion(pytta.SignalObj(x, 'time', FS), F0)
+        expected = np.hypot(0.01, 0.005)
+        print(f'\n  THD {thd[0]:.6f}, THD+N {thdn[0]:.6f} '
+              f'(expected {expected:.6f})')
+        self.assertAlmostEqual(thd[0] / expected, 1, delta=1e-3)
+        self.assertAlmostEqual(thdn[0] / (expected / np.sqrt(1 + expected ** 2)),
+                               1, delta=1e-3)
+
+    def test_noise_enters_thd_plus_n_and_not_thd(self):
+        # White noise of rms 0.01 on a sine of rms 1/sqrt(2): THD+N is the
+        # noise rms re the total rms, while THD only sees the noise in the
+        # narrow bands of the harmonics.
+        rng = np.random.default_rng(0)
+        noise = 0.01 * rng.standard_normal(FS)
+        x = np.sin(2 * np.pi * F0 * tt) + noise
+        thd, thdn = pytta.sine_distortion(pytta.SignalObj(x, 'time', FS), F0)
+        expected = np.std(noise) / np.sqrt(0.5 + np.var(noise))
+        print(f'\n  THD+N {thdn[0]:.5f} (expected {expected:.5f}), '
+              f'THD {thd[0]:.5f}')
+        self.assertAlmostEqual(thdn[0] / expected, 1, delta=0.01)
+        self.assertLess(thd[0], expected / 10)
+
+    def test_snr_of_signal_and_noise_recordings(self):
+        rng = np.random.default_rng(1)
+        sig = pytta.SignalObj(np.sin(2 * np.pi * F0 * tt), 'time', FS)
+        noise = pytta.SignalObj(1e-3 * rng.standard_normal(FS), 'time', FS)
+        expected = 20 * np.log10(np.sqrt(0.5) / np.std(noise.timeSignal))
+        self.assertAlmostEqual(pytta.snr(sig, noise)[0], expected, delta=0.01)
+
+    def test_rejects_a_fundamental_beyond_nyquist(self):
+        sig = pytta.SignalObj(np.zeros(FS), 'time', FS)
+        with self.assertRaises(ValueError):
+            pytta.sine_distortion(sig, FS / 2)
+
+
+if __name__ == '__main__':
+    unittest.main()