diff --git a/pytta/_plot.py b/pytta/_plot.py index c0a0970..f980b38 100644 --- a/pytta/_plot.py +++ b/pytta/_plot.py @@ -821,7 +821,7 @@ def spectrogram(sigObjs, winType, winSize, window size in samples * overlap (), (float): - window overlap in % + overlapping fraction of consecutive windows, 0 to 1 * xLabel ('Time [s]'), (str): x axis label. @@ -868,7 +868,7 @@ def spectrogram(sigObjs, winType, winSize, ax.set_snap(False) _spectrogram, _specTime, _specFreq = \ - _calc_spectrogram(data['timeSignal'], data['timeVector'], + _calc_spectrogram(data['timeSignal'], data['samplingRate'], overlap, winType, winSize, data['dBRef']) @@ -956,30 +956,18 @@ def _curve_data_extractor_spectrogram(sigObjs): return curveData -def _calc_spectrogram(timeSignal, timeVector, samplingRate, overlap, winType, - winSize, dBRef): - window = eval('ss.windows.' + winType)(winSize) - nextIdx = int(winSize*overlap) - rng = int(timeSignal.shape[0]/winSize/overlap - 1) - _spectrogram = np.zeros((winSize//2 + 1, rng)) - _specFreq = np.linspace(0, samplingRate//2, winSize//2 + 1) - _specTime = np.linspace(0, timeVector[-1], rng) - - for N in range(rng): - try: - strIdx = N*nextIdx - endIdx = winSize + N*nextIdx - sliceAudio = window*timeSignal[strIdx:endIdx] - sliceFFT = np.fft.rfft(sliceAudio, axis=0) - sliceMag = np.absolute(sliceFFT) * (2/sliceFFT.size) - _spectrogram[:, N] = 20*np.log10(sliceMag/dBRef) - - except IndexError: - sliceAudio = timeSignal[-winSize:] - sliceFFT = np.fft.rfft(sliceAudio, axis=0) - sliceMag = np.absolute(sliceFFT) * (2/sliceFFT.size) - _spectrogram[:, N] = 20*np.log10(sliceMag) - +def _calc_spectrogram(timeSignal, samplingRate, overlap, winType, winSize, + dBRef): + """Spectrogram in dB of peak amplitude re dBRef. `overlap` is the + overlapping fraction of consecutive windows; each column is at the + centre of its window.""" + window = getattr(ss.windows, winType)(winSize) + _specFreq, _specTime, magnitude = ss.spectrogram( + timeSignal, samplingRate, window=window, + noverlap=int(winSize * overlap), detrend=False, mode='magnitude', + scaling='spectrum') # |X| / sum(window): half the peak amplitude + with np.errstate(divide='ignore'): + _spectrogram = 20 * np.log10(2 * magnitude / dBRef) return _spectrogram, _specTime, _specFreq diff --git a/pytta/classes/signal.py b/pytta/classes/signal.py index 79151c3..91e63a3 100644 --- a/pytta/classes/signal.py +++ b/pytta/classes/signal.py @@ -753,7 +753,7 @@ def plot_spectrogram(self, winType:str='hann', winSize:int=1024, window size in samples * overlap (0.5), (float): - window overlap in % + overlapping fraction of consecutive windows, 0 to 1 * xLabel ('Time [s]'), (str): x axis label. diff --git a/pytta/functions.py b/pytta/functions.py index 9b39239..da0d0d4 100644 --- a/pytta/functions.py +++ b/pytta/functions.py @@ -550,7 +550,7 @@ def plot_spectrogram(*sigObjs, winType:str='hann', winSize:int=1024, window size in samples * overlap (0.5), (float): - window overlap in % + overlapping fraction of consecutive windows, 0 to 1 * xLabel ('Time [s]'), (str): x axis label. diff --git a/tests/test_spectrogram.py b/tests/test_spectrogram.py new file mode 100644 index 0000000..fe674e3 --- /dev/null +++ b/tests/test_spectrogram.py @@ -0,0 +1,63 @@ +""" +Tests for the plain spectrogram of pytta.plot_spectrogram +(pytta._plot._calc_spectrogram). + +Each test pins one rule: `overlap` is the overlapping fraction of +consecutive windows, each column sits at the centre of its window, and a +sine reads at its peak amplitude in dB whatever the window. +""" +import os +import unittest + +os.environ.setdefault('MPLBACKEND', 'Agg') + +import numpy as np + +from pytta import _plot + +FS = 8000 +W = 1024 + + +def spectrogram(x, overlap=0.5, winType='hann'): + S, t, f = _plot._calc_spectrogram(x, FS, overlap, winType, W, 1.0) + return S, t, f + + +class TestSpectrogram(unittest.TestCase): + + def test_overlap_is_the_overlapping_fraction(self): + # Hop = W (1 - overlap); every window lies inside the signal, so + # there are (n - W) // hop + 1 columns, for any overlap in [0, 1). + x = np.random.default_rng(0).standard_normal(FS) + for overlap in (0, 0.25, 0.5, 0.75, 0.9): + hop = W - int(W * overlap) + S, t, f = spectrogram(x, overlap) + self.assertEqual(S.shape, (W // 2 + 1, (FS - W) // hop + 1)) + self.assertTrue(np.allclose(np.diff(t), hop / FS)) + + def test_columns_sit_at_the_centre_of_their_window(self): + # The first column is centred at W / 2, and a click is loudest in + # the column whose centre is nearest to it. + x = np.zeros(FS) + x[5000] = 1 + S, t, f = spectrogram(x, 0.75) + self.assertAlmostEqual(t[0], W / 2 / FS) + loudest = t[np.argmax(S[W // 8])] + self.assertLessEqual(abs(loudest - 5000 / FS), 128 / FS) + + def test_sine_reads_at_its_peak_amplitude(self): + # A sine of amplitude 0.5 centred on a bin reads -6.02 dB re 1, with + # the coherent gain of each window taken out. + k = 64 + tt = np.arange(FS) / FS + x = 0.5 * np.sin(2 * np.pi * k * FS / W * tt) + for winType in ('hann', 'blackman', 'boxcar'): + S, t, f = spectrogram(x, 0.5, winType) + level = np.median(S[k]) + print(f'\n {winType:8s} {level:+.3f} dB (expected -6.021)') + self.assertAlmostEqual(level, 20 * np.log10(0.5), delta=0.05) + + +if __name__ == '__main__': + unittest.main()