Repository: pyfar/pyfar
Pull request #666 (2024-08-17 08:42:31)

fix squeeze in dsp.spectrogram to keep cshape
separate pull to fix squeeze in dsp.spectrogram to keep cshape of input signal

merges into #659 


Patch:
diff --git a/pyfar/dsp/dsp.py b/pyfar/dsp/dsp.py
index b903fc98..2b3d51b0 100644
--- a/pyfar/dsp/dsp.py
+++ b/pyfar/dsp/dsp.py
@@ -67,7 +67,9 @@ def group_delay(signal, frequencies=None, method='fft'):
     Returns
     -------
     group_delay : numpy array
-        Frequency dependent group delay in samples.
+        Frequency dependent group delay of shape
+        (:py:func:`~pyfar.classes.audio.Signal.cshape`,
+        frequencies).
 
     References
     ----------
@@ -115,11 +117,6 @@ def group_delay(signal, frequencies=None, method='fft'):
         raise ValueError(
             "Invalid method, needs to be either 'scipy' or 'fft'.")
 
-    # squeeze last dimension. If only one frequency is given, shape (1, ) is
-    # returned instead of (1, 1)
-    if frequencies is not None and frequencies.size == 1:
-        group_delay = np.squeeze(group_delay, axis=-1)
-
     return group_delay
 
 
@@ -299,7 +296,7 @@ def spectrogram(signal, window='hann', window_length=1024,
     window = sgn.get_window(window, window_length)
 
     frequencies, times, spectrogram = sgn.spectrogram(
-        x=signal.time.squeeze(), fs=signal.sampling_rate, window=window,
+        x=signal.time, fs=signal.sampling_rate, window=window,
         noverlap=window_overlap, mode='magnitude', scaling='spectrum')
 
     # remove normalization from scipy.signal.spectrogram
diff --git a/pyfar/plot/_two_d.py b/pyfar/plot/_two_d.py
index 8a46f0d9..7e15ff8d 100644
--- a/pyfar/plot/_two_d.py
+++ b/pyfar/plot/_two_d.py
@@ -337,6 +337,11 @@ def _spectrogram(signal, dB=True, log_prefix=None, log_reference=1,
     frequencies, times, spectrogram = dsp.spectrogram(
         signal[first_channel], window, window_length, window_overlap_fct)
 
+    # squeeze dsp.spectrogram returns for correct plotting
+    times = times.squeeze()
+    frequencies = frequencies.squeeze()
+    spectrogram = spectrogram.squeeze()
+
     # get magnitude data in dB
     if dB:
         if log_prefix is None:


Test patch:
diff --git a/tests/test_dsp.py b/tests/test_dsp.py
index 0d24cba3..78a2f483 100644
--- a/tests/test_dsp.py
+++ b/tests/test_dsp.py
@@ -97,7 +97,7 @@ def test_group_delay_custom_frequencies(impulse_group_delay):
     frequency = 1000
     frequency_idx = np.abs(signal.frequencies-frequency).argmin()
     grp = dsp.group_delay(signal, frequency, method='scipy')
-    assert grp.shape == (1, )
+    assert grp.shape == (1, 1)
     npt.assert_allclose(
         grp, impulse_group_delay[1][0, frequency_idx], atol=1e-10)
 
@@ -119,7 +119,7 @@ def test_group_delay_cshape(shape):
 
     grp = dsp.group_delay(impulse)
 
-    assert grp.shape == (*impulse.cshape, grp.shape[-1])
+    assert grp.shape == (*impulse.cshape, impulse.n_bins)
 
 
 def test_linear_phase():
diff --git a/tests/test_dsp_spectrogram.py b/tests/test_dsp_spectrogram.py
index 24479f49..db2266af 100644
--- a/tests/test_dsp_spectrogram.py
+++ b/tests/test_dsp_spectrogram.py
@@ -31,9 +31,9 @@ def test_return_values():
     npt.assert_allclose(times, [0, 512/1024, 1])
 
     # check middle slice
-    npt.assert_allclose(spectro[:256, 1], 0, atol=1e-13)
-    npt.assert_allclose(spectro[256, 1], 1, atol=1e-13)
-    npt.assert_allclose(spectro[257:, 1], 0, atol=1e-13)
+    npt.assert_allclose(spectro[0, :256, 1], 0, atol=1e-13)
+    npt.assert_allclose(spectro[0, 256, 1], 1, atol=1e-13)
+    npt.assert_allclose(spectro[0, 257:, 1], 0, atol=1e-13)
 
 
 @pytest.mark.parametrize('window,value', [
@@ -47,7 +47,7 @@ def test_window(window, value):
     _, _, spectro = pf.dsp.spectrogram(signal, window=window)
 
     # check middle slice
-    npt.assert_allclose(spectro[255:258, 1], value, atol=1e-13)
+    npt.assert_allclose(spectro[0, 255:258, 1], value, atol=1e-13)
 
 
 def test_normalize(sine):
@@ -55,3 +55,15 @@ def test_normalize(sine):
     sine.fft_norm = 'amplitude'
     assert pf.dsp.spectrogram(sine)[-1].max() < 1
     assert pf.dsp.spectrogram(sine, normalize=False)[-1].max() > 1
+
+
+@pytest.mark.parametrize('shape', [(2, 1), (1, 2), (1,), (1, 1)])
+def test_spectrogram_shape(shape):
+    """Test cshape of spectrogram returns"""
+    impulse = pf.signals.impulse(2048, 0, np.ones((shape)))
+    freq, time, spectro = pf.dsp.spectrogram(impulse)
+
+    assert spectro.shape == (*shape, *freq.shape, *time.shape)
+    assert time.ndim == 1
+    assert freq.ndim == 1
+