From dbb4e62c8fb4f9b8ca18f6bc5f21e89c10fbd95a Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Fri, 15 Nov 2019 23:22:04 +0100 Subject: [PATCH] Parametrize some spectral tests. --- lib/matplotlib/tests/test_mlab.py | 290 ++++-------------------------- 1 file changed, 31 insertions(+), 259 deletions(-) diff --git a/lib/matplotlib/tests/test_mlab.py b/lib/matplotlib/tests/test_mlab.py index 9447cfb3707c..bbe9ac8bafaf 100644 --- a/lib/matplotlib/tests/test_mlab.py +++ b/lib/matplotlib/tests/test_mlab.py @@ -1259,25 +1259,11 @@ def check_maxfreq(self, spec, fsp, fstims): del fstimst[-1] spect[maxind-5:maxind+5] = 0 - def test_spectral_helper_raises_complex_same_data(self): - # test that mode 'complex' cannot be used if x is not y + @pytest.mark.parametrize( # Modes that require `x is y`. + 'mode', ['complex', 'magnitude', 'angle', 'phase']) + def test_spectral_helper_mode_requires_x_is_y(self, mode): with pytest.raises(ValueError): - mlab._spectral_helper(x=self.y, y=self.y+1, mode='complex') - - def test_spectral_helper_raises_magnitude_same_data(self): - # test that mode 'magnitude' cannot be used if x is not y - with pytest.raises(ValueError): - mlab._spectral_helper(x=self.y, y=self.y+1, mode='magnitude') - - def test_spectral_helper_raises_angle_same_data(self): - # test that mode 'angle' cannot be used if x is not y - with pytest.raises(ValueError): - mlab._spectral_helper(x=self.y, y=self.y+1, mode='angle') - - def test_spectral_helper_raises_phase_same_data(self): - # test that mode 'phase' cannot be used if x is not y - with pytest.raises(ValueError): - mlab._spectral_helper(x=self.y, y=self.y+1, mode='phase') + mlab._spectral_helper(x=self.y, y=self.y+1, mode=mode) def test_spectral_helper_raises_unknown_mode(self): # test that unknown value for mode cannot be used @@ -1305,15 +1291,10 @@ def test_spectral_helper_raises_winlen_ne_NFFT(self): mlab._spectral_helper(x=self.y, y=self.y, NFFT=10, window=np.ones(9)) - def test_single_spectrum_helper_raises_mode_default(self): - # test that mode 'default' cannot be used with _single_spectrum_helper + @pytest.mark.parametrize('mode', ['default', 'psd']) + def test_single_spectrum_helper_unsupported_modes(self, mode): with pytest.raises(ValueError): - mlab._single_spectrum_helper(x=self.y, mode='default') - - def test_single_spectrum_helper_raises_mode_psd(self): - # test that mode 'psd' cannot be used with _single_spectrum_helper - with pytest.raises(ValueError): - mlab._single_spectrum_helper(x=self.y, mode='psd') + mlab._single_spectrum_helper(x=self.y, mode=mode) def test_spectral_helper_psd(self): freqs = self.freqs_density @@ -1397,10 +1378,14 @@ def test_psd(self): assert spec.shape == freqs.shape self.check_freqs(spec, freqs, fsp, self.fstims) - def test_psd_detrend_mean_func_offset(self): + @pytest.mark.parametrize( + 'make_data, detrend', + [(np.zeros, mlab.detrend_mean), (np.zeros, 'mean'), + (np.arange, mlab.detrend_linear), (np.arange, 'linear')]) + def test_psd_detrend(self, make_data, detrend): if self.NFFT_density is None: return - ydata = np.zeros(self.NFFT_density) + ydata = make_data(self.NFFT_density) ydata1 = ydata+5 ydata2 = ydata+3.3 ydata = np.vstack([ydata1, ydata2]) @@ -1413,118 +1398,13 @@ def test_psd_detrend_mean_func_offset(self): Fs=self.Fs, noverlap=0, sides=self.sides, - detrend=mlab.detrend_mean) + detrend=detrend) spec_b, fsp_b = mlab.psd(x=ydatab, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, - detrend=mlab.detrend_mean) - spec_c, fsp_c = mlab.psd(x=ycontrol, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides) - assert_array_equal(fsp_g, fsp_c) - assert_array_equal(fsp_b, fsp_c) - assert_allclose(spec_g, spec_c, atol=1e-08) - # these should not be almost equal - with pytest.raises(AssertionError): - assert_allclose(spec_b, spec_c, atol=1e-08) - - def test_psd_detrend_mean_str_offset(self): - if self.NFFT_density is None: - return - ydata = np.zeros(self.NFFT_density) - ydata1 = ydata+5 - ydata2 = ydata+3.3 - ydata = np.vstack([ydata1, ydata2]) - ydata = np.tile(ydata, (20, 1)) - ydatab = ydata.T.flatten() - ydata = ydata.flatten() - ycontrol = np.zeros_like(ydata) - spec_g, fsp_g = mlab.psd(x=ydata, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides, - detrend='mean') - spec_b, fsp_b = mlab.psd(x=ydatab, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides, - detrend='mean') - spec_c, fsp_c = mlab.psd(x=ycontrol, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides) - assert_array_equal(fsp_g, fsp_c) - assert_array_equal(fsp_b, fsp_c) - assert_allclose(spec_g, spec_c, atol=1e-08) - # these should not be almost equal - with pytest.raises(AssertionError): - assert_allclose(spec_b, spec_c, atol=1e-08) - - def test_psd_detrend_linear_func_trend(self): - if self.NFFT_density is None: - return - ydata = np.arange(self.NFFT_density) - ydata1 = ydata+5 - ydata2 = ydata+3.3 - ydata = np.vstack([ydata1, ydata2]) - ydata = np.tile(ydata, (20, 1)) - ydatab = ydata.T.flatten() - ydata = ydata.flatten() - ycontrol = np.zeros_like(ydata) - spec_g, fsp_g = mlab.psd(x=ydata, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides, - detrend=mlab.detrend_linear) - spec_b, fsp_b = mlab.psd(x=ydatab, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides, - detrend=mlab.detrend_linear) - spec_c, fsp_c = mlab.psd(x=ycontrol, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides) - assert_array_equal(fsp_g, fsp_c) - assert_array_equal(fsp_b, fsp_c) - assert_allclose(spec_g, spec_c, atol=1e-08) - # these should not be almost equal - with pytest.raises(AssertionError): - assert_allclose(spec_b, spec_c, atol=1e-08) - - def test_psd_detrend_linear_str_trend(self): - if self.NFFT_density is None: - return - ydata = np.arange(self.NFFT_density) - ydata1 = ydata+5 - ydata2 = ydata+3.3 - ydata = np.vstack([ydata1, ydata2]) - ydata = np.tile(ydata, (20, 1)) - ydatab = ydata.T.flatten() - ydata = ydata.flatten() - ycontrol = np.zeros_like(ydata) - spec_g, fsp_g = mlab.psd(x=ydata, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides, - detrend='linear') - spec_b, fsp_b = mlab.psd(x=ydatab, - NFFT=self.NFFT_density, - Fs=self.Fs, - noverlap=0, - sides=self.sides, - detrend='linear') + detrend=detrend) spec_c, fsp_c = mlab.psd(x=ycontrol, NFFT=self.NFFT_density, Fs=self.Fs, @@ -1710,95 +1590,11 @@ def test_phase_spectrum(self): assert_allclose(fsp, freqs, atol=1e-06) assert spec.shape == freqs.shape - def test_specgram_auto(self): - freqs = self.freqs_specgram - spec, fsp, t = mlab.specgram(x=self.y, - NFFT=self.NFFT_specgram, - Fs=self.Fs, - noverlap=self.nover_specgram, - pad_to=self.pad_to_specgram, - sides=self.sides) - specm = np.mean(spec, axis=1) - - assert_allclose(fsp, freqs, atol=1e-06) - assert_allclose(t, self.t_specgram, atol=1e-06) - - assert spec.shape[0] == freqs.shape[0] - assert spec.shape[1] == self.t_specgram.shape[0] - - # since we are using a single freq, all time slices - # should be about the same - if np.abs(spec.max()) != 0: - assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0, - atol=1e-02) - self.check_freqs(specm, freqs, fsp, self.fstims) - - def test_specgram_default(self): - freqs = self.freqs_specgram - spec, fsp, t = mlab.specgram(x=self.y, - NFFT=self.NFFT_specgram, - Fs=self.Fs, - noverlap=self.nover_specgram, - pad_to=self.pad_to_specgram, - sides=self.sides, - mode='default') - specm = np.mean(spec, axis=1) - - assert_allclose(fsp, freqs, atol=1e-06) - assert_allclose(t, self.t_specgram, atol=1e-06) - - assert spec.shape[0] == freqs.shape[0] - assert spec.shape[1] == self.t_specgram.shape[0] - - # since we are using a single freq, all time slices - # should be about the same - if np.abs(spec.max()) != 0: - assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0, - atol=1e-02) - self.check_freqs(specm, freqs, fsp, self.fstims) - - def test_specgram_psd(self): - freqs = self.freqs_specgram - spec, fsp, t = mlab.specgram(x=self.y, - NFFT=self.NFFT_specgram, - Fs=self.Fs, - noverlap=self.nover_specgram, - pad_to=self.pad_to_specgram, - sides=self.sides, - mode='psd') - specm = np.mean(spec, axis=1) - - assert_allclose(fsp, freqs, atol=1e-06) - assert_allclose(t, self.t_specgram, atol=1e-06) - - assert spec.shape[0] == freqs.shape[0] - assert spec.shape[1] == self.t_specgram.shape[0] - # since we are using a single freq, all time slices - # should be about the same - if np.abs(spec.max()) != 0: - assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0, - atol=1e-02) - self.check_freqs(specm, freqs, fsp, self.fstims) - - def test_specgram_complex(self): - freqs = self.freqs_specgram - spec, fsp, t = mlab.specgram(x=self.y, - NFFT=self.NFFT_specgram, - Fs=self.Fs, - noverlap=self.nover_specgram, - pad_to=self.pad_to_specgram, - sides=self.sides, - mode='complex') - specm = np.mean(np.abs(spec), axis=1) - assert_allclose(fsp, freqs, atol=1e-06) - assert_allclose(t, self.t_specgram, atol=1e-06) - - assert spec.shape[0] == freqs.shape[0] - assert spec.shape[1] == self.t_specgram.shape[0] - - self.check_freqs(specm, freqs, fsp, self.fstims) - - def test_specgram_magnitude(self): + @pytest.mark.parametrize( + 'kwargs', + [{}, {'mode': 'default'}, {'mode': 'psd'}, {'mode': 'magnitude'}, + {'mode': 'complex'}, {'mode': 'angle'}, {'mode': 'phase'}]) + def test_specgram(self, kwargs): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, NFFT=self.NFFT_specgram, @@ -1806,49 +1602,25 @@ def test_specgram_magnitude(self): noverlap=self.nover_specgram, pad_to=self.pad_to_specgram, sides=self.sides, - mode='magnitude') + **kwargs) + if kwargs.get('mode') == 'complex': + spec = np.abs(spec) specm = np.mean(spec, axis=1) - assert_allclose(fsp, freqs, atol=1e-06) - assert_allclose(t, self.t_specgram, atol=1e-06) - assert spec.shape[0] == freqs.shape[0] - assert spec.shape[1] == self.t_specgram.shape[0] - # since we are using a single freq, all time slices - # should be about the same - if np.abs(spec.max()) != 0: - assert_allclose(np.diff(spec, axis=1).max()/np.abs(spec.max()), 0, - atol=1e-02) - self.check_freqs(specm, freqs, fsp, self.fstims) - - def test_specgram_angle(self): - freqs = self.freqs_specgram - spec, fsp, t = mlab.specgram(x=self.y, - NFFT=self.NFFT_specgram, - Fs=self.Fs, - noverlap=self.nover_specgram, - pad_to=self.pad_to_specgram, - sides=self.sides, - mode='angle') assert_allclose(fsp, freqs, atol=1e-06) assert_allclose(t, self.t_specgram, atol=1e-06) assert spec.shape[0] == freqs.shape[0] assert spec.shape[1] == self.t_specgram.shape[0] - def test_specgram_phase(self): - freqs = self.freqs_specgram - spec, fsp, t = mlab.specgram(x=self.y, - NFFT=self.NFFT_specgram, - Fs=self.Fs, - noverlap=self.nover_specgram, - pad_to=self.pad_to_specgram, - sides=self.sides, - mode='phase') - assert_allclose(fsp, freqs, atol=1e-06) - assert_allclose(t, self.t_specgram, atol=1e-06) - - assert spec.shape[0] == freqs.shape[0] - assert spec.shape[1] == self.t_specgram.shape[0] + if kwargs.get('mode') not in ['complex', 'angle', 'phase']: + # using a single freq, so all time slices should be about the same + if np.abs(spec.max()) != 0: + assert_allclose( + np.diff(spec, axis=1).max() / np.abs(spec.max()), 0, + atol=1e-02) + if kwargs.get('mode') not in ['angle', 'phase']: + self.check_freqs(specm, freqs, fsp, self.fstims) def test_specgram_warn_only1seg(self): """Warning should be raised if len(x) <= NFFT."""