Repository: pyfar/pyfar
Pull request #659 (2024-08-01 16:06:37)

Fix squeeze in dsp functions to keep cshape of input signal
Adjusted the dsp-functions so signals won't be squeezed and the original cshape is kept.

separate PR for dsp.spectrogram #666 


Patch:
diff --git a/pyfar/dsp/dsp.py b/pyfar/dsp/dsp.py
index ba0d2357..2b3d51b0 100644
--- a/pyfar/dsp/dsp.py
+++ b/pyfar/dsp/dsp.py
@@ -67,8 +67,9 @@ def group_delay(signal, frequencies=None, method='fft'):
     Returns
     -------
     group_delay : numpy array
-        Frequency dependent group delay in samples. The array is flattened if
-        a single channel signal was passed to the function.
+        Frequency dependent group delay of shape
+        (:py:func:`~pyfar.classes.audio.Signal.cshape`,
+        frequencies).
 
     References
     ----------
@@ -116,10 +117,6 @@ def group_delay(signal, frequencies=None, method='fft'):
         raise ValueError(
             "Invalid method, needs to be either 'scipy' or 'fft'.")
 
-    # flatten in numpy fashion if a single channel is returned
-    if signal.cshape == (1, ):
-        group_delay = np.squeeze(group_delay)
-
     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
@@ -1430,7 +1427,7 @@ def find_impulse_response_start(
                     f'No values below threshold found found for channel {ch}',
                     'defaulting to 0')
 
-    return np.squeeze(start_sample)
+    return start_sample
 
 
 @rename_arg({"freq_range": "frequency_range"},
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 0788f5f0..78a2f483 100644
--- a/tests/test_dsp.py
+++ b/tests/test_dsp.py
@@ -49,20 +49,20 @@ def test_group_delay_single_channel(impulse_group_delay):
         dsp.group_delay(signal, method='fft', frequencies=[1, 2, 3])
 
     grp = dsp.group_delay(signal, method='scipy')
-    assert grp.shape == (signal.n_bins, )
+    assert grp.shape == (1, signal.n_bins)
     npt.assert_allclose(
-        grp, impulse_group_delay[1].flatten(), rtol=1e-10, atol=1e-10)
+        grp[0, :], impulse_group_delay[1].flatten(), rtol=1e-10, atol=1e-10)
 
     grp = dsp.group_delay(signal, method='fft')
-    assert grp.shape == (signal.n_bins, )
+    assert grp.shape == (1, signal.n_bins)
     npt.assert_allclose(
-        grp, impulse_group_delay[1].flatten(), rtol=1e-10, atol=1e-10)
+        grp[0, :], impulse_group_delay[1].flatten(), rtol=1e-10, atol=1e-10)
 
     grp = dsp.group_delay(
         signal, method='fft')
-    assert grp.shape == (signal.n_bins, )
+    assert grp.shape == (1, signal.n_bins)
     npt.assert_allclose(
-        grp, impulse_group_delay[1].flatten(), rtol=1e-10, atol=1e-10)
+        grp[0, :], impulse_group_delay[1].flatten(), rtol=1e-10, atol=1e-10)
 
 
 def test_group_delay_two_channel(impulse_group_delay_two_channel):
@@ -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 == ()
+    assert grp.shape == (1, 1)
     npt.assert_allclose(
         grp, impulse_group_delay[1][0, frequency_idx], atol=1e-10)
 
@@ -106,9 +106,20 @@ def test_group_delay_custom_frequencies(impulse_group_delay):
     frequency_idx = np.abs(
         signal.frequencies-frequency[..., np.newaxis]).argmin(axis=-1)
     grp = dsp.group_delay(signal, frequency, method='scipy')
-    assert grp.shape == (2,)
+    assert grp.shape == (1, 2)
     npt.assert_allclose(
-        grp, impulse_group_delay[1][0, frequency_idx], atol=1e-10)
+        grp[0, :], impulse_group_delay[1][0, frequency_idx], atol=1e-10)
+
+
+@pytest.mark.parametrize("shape", [(4, 1), (1, 4), (1, ), (1, 1)])
+def test_group_delay_cshape(shape):
+    """Test if group delay function keeps cshape of signals of shape `m x n`
+     with `m = 1` or `n = 1`."""
+    impulse = pf.signals.impulse(256, 0, np.ones(shape))
+
+    grp = dsp.group_delay(impulse)
+
+    assert grp.shape == (*impulse.cshape, impulse.n_bins)
 
 
 def test_linear_phase():
@@ -121,13 +132,13 @@ def test_linear_phase():
     y = dsp.linear_phase(x, N/2)
     # test output
     assert isinstance(y, pf.Signal)
-    npt.assert_allclose(dsp.group_delay(y), N / 2 * np.ones(y.n_bins))
+    npt.assert_allclose(dsp.group_delay(y)[0, :], N / 2 * np.ones(y.n_bins))
     # test if input did not change
     npt.assert_allclose(x.time, pf.signals.impulse(N).time)
 
     # test group delay in seconds
     y = dsp.linear_phase(x, N / 2 / fs, unit="s")
-    npt.assert_allclose(dsp.group_delay(y), N / 2 * np.ones(y.n_bins))
+    npt.assert_allclose(dsp.group_delay(y)[0, :], N / 2 * np.ones(y.n_bins))
 
     # test assertion
     with pytest.raises(TypeError, match="signal must be a pyfar Signal"):
@@ -144,13 +155,13 @@ def test_linear_phase_multichannel():
 
     # test with scalar group delay
     y = dsp.linear_phase(x, N/2)
-    npt.assert_allclose(dsp.group_delay(y[0]), N / 2 * np.ones(y.n_bins))
-    npt.assert_allclose(dsp.group_delay(y[1]), N / 2 * np.ones(y.n_bins))
+    npt.assert_allclose(dsp.group_delay(y[0])[0, :], N / 2 * np.ones(y.n_bins))
+    npt.assert_allclose(dsp.group_delay(y[1])[0, :], N / 2 * np.ones(y.n_bins))
 
     # test with array like group delay
     y = dsp.linear_phase(x, [N/2, N/4])
-    npt.assert_allclose(dsp.group_delay(y[0]), N / 2 * np.ones(y.n_bins))
-    npt.assert_allclose(dsp.group_delay(y[1]), N / 4 * np.ones(y.n_bins))
+    npt.assert_allclose(dsp.group_delay(y[0])[0, :], N / 2 * np.ones(y.n_bins))
+    npt.assert_allclose(dsp.group_delay(y[1])[0, :], N / 4 * np.ones(y.n_bins))
 
 
 def test_zero_phase():
@@ -654,6 +665,16 @@ def test_impulse_response_delay_multidim():
     npt.assert_allclose(start_sample_est, start_sample, atol=1e-2)
 
 
+@pytest.mark.parametrize("shape", [(4, 1), (1, 4), (1, ), (1, 1)])
+def test_impulse_response_delay_cshape(shape):
+    """Test if find_impulse_response_delay function keeps cshape of signals of
+    shape `m x n` with `m = 1` or `n = 1`."""
+    ir = pf.signals.impulse(256, 24, np.ones(shape))
+    start_sample = dsp.find_impulse_response_delay(ir)
+
+    assert start_sample.shape == ir.cshape
+
+
 def test_impulse_response_start_insufficient_snr():
     n_samples = 2**9
     snr = 15
@@ -756,6 +777,16 @@ def test_impulse_response_start_multidim():
     npt.assert_allclose(start_sample_est, start_samples - 1)
 
 
+@pytest.mark.parametrize("shape", [(4, 1), (1, 4), (1, ), (1, 1)])
+def test_impulse_response_start_cshape(shape):
+    """Test if find_impulse_response_start function keeps cshape of signals of
+    shape `m x n` with `m = 1` or `n = 1`."""
+    ir = pf.signals.impulse(256, 24, np.ones(shape))
+    start_sample = dsp.find_impulse_response_start(ir)
+
+    assert start_sample.shape == ir.cshape
+
+
 def test_convolve_default():
     x = pf.Signal([1, 0.5, 0.25, 0], 44100)
     y = pf.Signal([1, -1, 0], 44100)
diff --git a/tests/test_dsp_interpolation.py b/tests/test_dsp_interpolation.py
index d607c78b..81a2a1c4 100644
--- a/tests/test_dsp_interpolation.py
+++ b/tests/test_dsp_interpolation.py
@@ -203,7 +203,7 @@ def test_fractional_delay_mode_cyclic(delay):
     delayed = fractional_time_shift(signal, delay, mode="cyclic")
 
     # if the delay is too large, it is cyclicly shifted
-    group_delay = pf.dsp.group_delay(delayed)[0]
+    group_delay = pf.dsp.group_delay(delayed)[0, 0]
     npt.assert_allclose(group_delay, (16+delay) % 32, atol=.05)
 
 
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
+