from numpy.testing import (assert_allclose, assert_almost_equal, assert_array_equal, assert_array_almost_equal_nulp) import numpy as np import pytest import matplotlib.mlab as mlab class TestStride: def get_base(self, x): y = x while y.base is not None: y = y.base return y def calc_window_target(self, x, NFFT, noverlap=0, axis=0): """ This is an adaptation of the original window extraction algorithm. This is here to test to make sure the new implementation has the same result. """ step = NFFT - noverlap ind = np.arange(0, len(x) - NFFT + 1, step) n = len(ind) result = np.zeros((NFFT, n)) # do the ffts of the slices for i in range(n): result[:, i] = x[ind[i]:ind[i]+NFFT] if axis == 1: result = result.T return result @pytest.mark.parametrize('shape', [(), (10, 1)], ids=['0D', '2D']) def test_stride_windows_invalid_input_shape(self, shape): x = np.arange(np.prod(shape)).reshape(shape) with pytest.raises(ValueError): mlab.stride_windows(x, 5) @pytest.mark.parametrize('n, noverlap', [(0, None), (11, None), (2, 2), (2, 3)], ids=['n less than 1', 'n greater than input', 'noverlap greater than n', 'noverlap equal to n']) def test_stride_windows_invalid_params(self, n, noverlap): x = np.arange(10) with pytest.raises(ValueError): mlab.stride_windows(x, n, noverlap) @pytest.mark.parametrize('axis', [0, 1], ids=['axis0', 'axis1']) @pytest.mark.parametrize('n, noverlap', [(1, 0), (5, 0), (15, 2), (13, -3)], ids=['n1-noverlap0', 'n5-noverlap0', 'n15-noverlap2', 'n13-noverlapn3']) def test_stride_windows(self, n, noverlap, axis): x = np.arange(100) y = mlab.stride_windows(x, n, noverlap=noverlap, axis=axis) expected_shape = [0, 0] expected_shape[axis] = n expected_shape[1 - axis] = 100 // (n - noverlap) yt = self.calc_window_target(x, n, noverlap=noverlap, axis=axis) assert yt.shape == y.shape assert_array_equal(yt, y) assert tuple(expected_shape) == y.shape assert self.get_base(y) is x @pytest.mark.parametrize('axis', [0, 1], ids=['axis0', 'axis1']) def test_stride_windows_n32_noverlap0_unflatten(self, axis): n = 32 x = np.arange(n)[np.newaxis] x1 = np.tile(x, (21, 1)) x2 = x1.flatten() y = mlab.stride_windows(x2, n, axis=axis) if axis == 0: x1 = x1.T assert y.shape == x1.shape assert_array_equal(y, x1) def test_window(): np.random.seed(0) n = 1000 rand = np.random.standard_normal(n) + 100 ones = np.ones(n) assert_array_equal(mlab.window_none(ones), ones) assert_array_equal(mlab.window_none(rand), rand) assert_array_equal(np.hanning(len(rand)) * rand, mlab.window_hanning(rand)) assert_array_equal(np.hanning(len(ones)), mlab.window_hanning(ones)) class TestDetrend: def setup(self): np.random.seed(0) n = 1000 x = np.linspace(0., 100, n) self.sig_zeros = np.zeros(n) self.sig_off = self.sig_zeros + 100. self.sig_slope = np.linspace(-10., 90., n) self.sig_slope_mean = x - x.mean() self.sig_base = ( np.random.standard_normal(n) + np.sin(x*2*np.pi/(n/100))) self.sig_base -= self.sig_base.mean() def allclose(self, *args): assert_allclose(*args, atol=1e-8) def test_detrend_none(self): assert mlab.detrend_none(0.) == 0. assert mlab.detrend_none(0., axis=1) == 0. assert mlab.detrend(0., key="none") == 0. assert mlab.detrend(0., key=mlab.detrend_none) == 0. for sig in [ 5.5, self.sig_off, self.sig_slope, self.sig_base, (self.sig_base + self.sig_slope + self.sig_off).tolist(), np.vstack([self.sig_base, # 2D case. self.sig_base + self.sig_off, self.sig_base + self.sig_slope, self.sig_base + self.sig_off + self.sig_slope]), np.vstack([self.sig_base, # 2D transposed case. self.sig_base + self.sig_off, self.sig_base + self.sig_slope, self.sig_base + self.sig_off + self.sig_slope]).T, ]: if isinstance(sig, np.ndarray): assert_array_equal(mlab.detrend_none(sig), sig) else: assert mlab.detrend_none(sig) == sig def test_detrend_mean(self): for sig in [0., 5.5]: # 0D. assert mlab.detrend_mean(sig) == 0. assert mlab.detrend(sig, key="mean") == 0. assert mlab.detrend(sig, key=mlab.detrend_mean) == 0. # 1D. self.allclose(mlab.detrend_mean(self.sig_zeros), self.sig_zeros) self.allclose(mlab.detrend_mean(self.sig_base), self.sig_base) self.allclose(mlab.detrend_mean(self.sig_base + self.sig_off), self.sig_base) self.allclose(mlab.detrend_mean(self.sig_base + self.sig_slope), self.sig_base + self.sig_slope_mean) self.allclose( mlab.detrend_mean(self.sig_base + self.sig_slope + self.sig_off), self.sig_base + self.sig_slope_mean) def test_detrend_mean_1d_base_slope_off_list_andor_axis0(self): input = self.sig_base + self.sig_slope + self.sig_off target = self.sig_base + self.sig_slope_mean self.allclose(mlab.detrend_mean(input, axis=0), target) self.allclose(mlab.detrend_mean(input.tolist()), target) self.allclose(mlab.detrend_mean(input.tolist(), axis=0), target) def test_detrend_mean_2d(self): input = np.vstack([self.sig_off, self.sig_base + self.sig_off]) target = np.vstack([self.sig_zeros, self.sig_base]) self.allclose(mlab.detrend_mean(input), target) self.allclose(mlab.detrend_mean(input, axis=None), target) self.allclose(mlab.detrend_mean(input.T, axis=None).T, target) self.allclose(mlab.detrend(input), target) self.allclose(mlab.detrend(input, axis=None), target) self.allclose( mlab.detrend(input.T, key="constant", axis=None), target.T) input = np.vstack([self.sig_base, self.sig_base + self.sig_off, self.sig_base + self.sig_slope, self.sig_base + self.sig_off + self.sig_slope]) target = np.vstack([self.sig_base, self.sig_base, self.sig_base + self.sig_slope_mean, self.sig_base + self.sig_slope_mean]) self.allclose(mlab.detrend_mean(input.T, axis=0), target.T) self.allclose(mlab.detrend_mean(input, axis=1), target) self.allclose(mlab.detrend_mean(input, axis=-1), target) self.allclose(mlab.detrend(input, key="default", axis=1), target) self.allclose(mlab.detrend(input.T, key="mean", axis=0), target.T) self.allclose( mlab.detrend(input.T, key=mlab.detrend_mean, axis=0), target.T) def test_detrend_ValueError(self): for signal, kwargs in [ (self.sig_slope[np.newaxis], {"key": "spam"}), (self.sig_slope[np.newaxis], {"key": 5}), (5.5, {"axis": 0}), (self.sig_slope, {"axis": 1}), (self.sig_slope[np.newaxis], {"axis": 2}), ]: with pytest.raises(ValueError): mlab.detrend(signal, **kwargs) def test_detrend_mean_ValueError(self): for signal, kwargs in [ (5.5, {"axis": 0}), (self.sig_slope, {"axis": 1}), (self.sig_slope[np.newaxis], {"axis": 2}), ]: with pytest.raises(ValueError): mlab.detrend_mean(signal, **kwargs) def test_detrend_linear(self): # 0D. assert mlab.detrend_linear(0.) == 0. assert mlab.detrend_linear(5.5) == 0. assert mlab.detrend(5.5, key="linear") == 0. assert mlab.detrend(5.5, key=mlab.detrend_linear) == 0. for sig in [ # 1D. self.sig_off, self.sig_slope, self.sig_slope + self.sig_off, ]: self.allclose(mlab.detrend_linear(sig), self.sig_zeros) def test_detrend_str_linear_1d(self): input = self.sig_slope + self.sig_off target = self.sig_zeros self.allclose(mlab.detrend(input, key="linear"), target) self.allclose(mlab.detrend(input, key=mlab.detrend_linear), target) self.allclose(mlab.detrend_linear(input.tolist()), target) def test_detrend_linear_2d(self): input = np.vstack([self.sig_off, self.sig_slope, self.sig_slope + self.sig_off]) target = np.vstack([self.sig_zeros, self.sig_zeros, self.sig_zeros]) self.allclose( mlab.detrend(input.T, key="linear", axis=0), target.T) self.allclose( mlab.detrend(input.T, key=mlab.detrend_linear, axis=0), target.T) self.allclose( mlab.detrend(input, key="linear", axis=1), target) self.allclose( mlab.detrend(input, key=mlab.detrend_linear, axis=1), target) with pytest.raises(ValueError): mlab.detrend_linear(self.sig_slope[np.newaxis]) @pytest.mark.parametrize('iscomplex', [False, True], ids=['real', 'complex'], scope='class') @pytest.mark.parametrize('sides', ['onesided', 'twosided', 'default'], scope='class') @pytest.mark.parametrize( 'fstims,len_x,NFFT_density,nover_density,pad_to_density,pad_to_spectrum', [ ([], None, -1, -1, -1, -1), ([4], None, -1, -1, -1, -1), ([4, 5, 10], None, -1, -1, -1, -1), ([], None, None, -1, -1, None), ([], None, -1, -1, None, None), ([], None, None, -1, None, None), ([], 1024, 512, -1, -1, 128), ([], 256, -1, -1, 33, 257), ([], 255, 33, -1, -1, None), ([], 256, 128, -1, 256, 256), ([], None, -1, 32, -1, -1), ], ids=[ 'nosig', 'Fs4', 'FsAll', 'nosig_noNFFT', 'nosig_nopad_to', 'nosig_noNFFT_no_pad_to', 'nosig_trim', 'nosig_odd', 'nosig_oddlen', 'nosig_stretch', 'nosig_overlap', ], scope='class') class TestSpectral: @pytest.fixture(scope='class', autouse=True) def stim(self, request, fstims, iscomplex, sides, len_x, NFFT_density, nover_density, pad_to_density, pad_to_spectrum): Fs = 100. x = np.arange(0, 10, 1 / Fs) if len_x is not None: x = x[:len_x] # get the stimulus frequencies, defaulting to None fstims = [Fs / fstim for fstim in fstims] # get the constants, default to calculated values if NFFT_density is None: NFFT_density_real = 256 elif NFFT_density < 0: NFFT_density_real = NFFT_density = 100 else: NFFT_density_real = NFFT_density if nover_density is None: nover_density_real = 0 elif nover_density < 0: nover_density_real = nover_density = NFFT_density_real // 2 else: nover_density_real = nover_density if pad_to_density is None: pad_to_density_real = NFFT_density_real elif pad_to_density < 0: pad_to_density = int(2**np.ceil(np.log2(NFFT_density_real))) pad_to_density_real = pad_to_density else: pad_to_density_real = pad_to_density if pad_to_spectrum is None: pad_to_spectrum_real = len(x) elif pad_to_spectrum < 0: pad_to_spectrum_real = pad_to_spectrum = len(x) else: pad_to_spectrum_real = pad_to_spectrum if pad_to_spectrum is None: NFFT_spectrum_real = NFFT_spectrum = pad_to_spectrum_real else: NFFT_spectrum_real = NFFT_spectrum = len(x) nover_spectrum = 0 NFFT_specgram = NFFT_density nover_specgram = nover_density pad_to_specgram = pad_to_density NFFT_specgram_real = NFFT_density_real nover_specgram_real = nover_density_real if sides == 'onesided' or (sides == 'default' and not iscomplex): # frequencies for specgram, psd, and csd # need to handle even and odd differently if pad_to_density_real % 2: freqs_density = np.linspace(0, Fs / 2, num=pad_to_density_real, endpoint=False)[::2] else: freqs_density = np.linspace(0, Fs / 2, num=pad_to_density_real // 2 + 1) # frequencies for complex, magnitude, angle, and phase spectrums # need to handle even and odd differently if pad_to_spectrum_real % 2: freqs_spectrum = np.linspace(0, Fs / 2, num=pad_to_spectrum_real, endpoint=False)[::2] else: freqs_spectrum = np.linspace(0, Fs / 2, num=pad_to_spectrum_real // 2 + 1) else: # frequencies for specgram, psd, and csd # need to handle even and odd differently if pad_to_density_real % 2: freqs_density = np.linspace(-Fs / 2, Fs / 2, num=2 * pad_to_density_real, endpoint=False)[1::2] else: freqs_density = np.linspace(-Fs / 2, Fs / 2, num=pad_to_density_real, endpoint=False) # frequencies for complex, magnitude, angle, and phase spectrums # need to handle even and odd differently if pad_to_spectrum_real % 2: freqs_spectrum = np.linspace(-Fs / 2, Fs / 2, num=2 * pad_to_spectrum_real, endpoint=False)[1::2] else: freqs_spectrum = np.linspace(-Fs / 2, Fs / 2, num=pad_to_spectrum_real, endpoint=False) freqs_specgram = freqs_density # time points for specgram t_start = NFFT_specgram_real // 2 t_stop = len(x) - NFFT_specgram_real // 2 + 1 t_step = NFFT_specgram_real - nover_specgram_real t_specgram = x[t_start:t_stop:t_step] if NFFT_specgram_real % 2: t_specgram += 1 / Fs / 2 if len(t_specgram) == 0: t_specgram = np.array([NFFT_specgram_real / (2 * Fs)]) t_spectrum = np.array([NFFT_spectrum_real / (2 * Fs)]) t_density = t_specgram y = np.zeros_like(x) for i, fstim in enumerate(fstims): y += np.sin(fstim * x * np.pi * 2) * 10**i if iscomplex: y = y.astype('complex') # Interestingly, the instance on which this fixture is called is not # the same as the one on which a test is run. So we need to modify the # class itself when using a class-scoped fixture. cls = request.cls cls.Fs = Fs cls.sides = sides cls.fstims = fstims cls.NFFT_density = NFFT_density cls.nover_density = nover_density cls.pad_to_density = pad_to_density cls.NFFT_spectrum = NFFT_spectrum cls.nover_spectrum = nover_spectrum cls.pad_to_spectrum = pad_to_spectrum cls.NFFT_specgram = NFFT_specgram cls.nover_specgram = nover_specgram cls.pad_to_specgram = pad_to_specgram cls.t_specgram = t_specgram cls.t_density = t_density cls.t_spectrum = t_spectrum cls.y = y cls.freqs_density = freqs_density cls.freqs_spectrum = freqs_spectrum cls.freqs_specgram = freqs_specgram cls.NFFT_density_real = NFFT_density_real def check_freqs(self, vals, targfreqs, resfreqs, fstims): assert resfreqs.argmin() == 0 assert resfreqs.argmax() == len(resfreqs)-1 assert_allclose(resfreqs, targfreqs, atol=1e-06) for fstim in fstims: i = np.abs(resfreqs - fstim).argmin() assert vals[i] > vals[i+2] assert vals[i] > vals[i-2] def check_maxfreq(self, spec, fsp, fstims): # skip the test if there are no frequencies if len(fstims) == 0: return # if twosided, do the test for each side if fsp.min() < 0: fspa = np.abs(fsp) zeroind = fspa.argmin() self.check_maxfreq(spec[:zeroind], fspa[:zeroind], fstims) self.check_maxfreq(spec[zeroind:], fspa[zeroind:], fstims) return fstimst = fstims[:] spect = spec.copy() # go through each peak and make sure it is correctly the maximum peak while fstimst: maxind = spect.argmax() maxfreq = fsp[maxind] assert_almost_equal(maxfreq, fstimst[-1]) del fstimst[-1] spect[maxind-5:maxind+5] = 0 def test_spectral_helper_raises(self): # We don't use parametrize here to handle ``y = self.y``. for kwargs in [ # Various error conditions: {"y": self.y+1, "mode": "complex"}, # Modes requiring ``x is y``. {"y": self.y+1, "mode": "magnitude"}, {"y": self.y+1, "mode": "angle"}, {"y": self.y+1, "mode": "phase"}, {"mode": "spam"}, # Bad mode. {"y": self.y, "sides": "eggs"}, # Bad sides. {"y": self.y, "NFFT": 10, "noverlap": 20}, # noverlap > NFFT. {"NFFT": 10, "noverlap": 10}, # noverlap == NFFT. {"y": self.y, "NFFT": 10, "window": np.ones(9)}, # len(win) != NFFT. ]: with pytest.raises(ValueError): mlab._spectral_helper(x=self.y, **kwargs) @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=mode) @pytest.mark.parametrize("mode, case", [ ("psd", "density"), ("magnitude", "specgram"), ("magnitude", "spectrum"), ]) def test_spectral_helper_psd(self, mode, case): freqs = getattr(self, f"freqs_{case}") spec, fsp, t = mlab._spectral_helper( x=self.y, y=self.y, NFFT=getattr(self, f"NFFT_{case}"), Fs=self.Fs, noverlap=getattr(self, f"nover_{case}"), pad_to=getattr(self, f"pad_to_{case}"), sides=self.sides, mode=mode) assert_allclose(fsp, freqs, atol=1e-06) assert_allclose(t, getattr(self, f"t_{case}"), atol=1e-06) assert spec.shape[0] == freqs.shape[0] assert spec.shape[1] == getattr(self, f"t_{case}").shape[0] def test_csd(self): freqs = self.freqs_density spec, fsp = mlab.csd(x=self.y, y=self.y+1, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides) assert_allclose(fsp, freqs, atol=1e-06) assert spec.shape == freqs.shape def test_csd_padding(self): """Test zero padding of csd().""" if self.NFFT_density is None: # for derived classes return sargs = dict(x=self.y, y=self.y+1, Fs=self.Fs, window=mlab.window_none, sides=self.sides) spec0, _ = mlab.csd(NFFT=self.NFFT_density, **sargs) spec1, _ = mlab.csd(NFFT=self.NFFT_density*2, **sargs) assert_almost_equal(np.sum(np.conjugate(spec0)*spec0).real, np.sum(np.conjugate(spec1/2)*spec1/2).real) def test_psd(self): freqs = self.freqs_density spec, fsp = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides) assert spec.shape == freqs.shape self.check_freqs(spec, freqs, fsp, self.fstims) @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 = make_data(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=detrend) spec_b, fsp_b = mlab.psd(x=ydatab, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, detrend=detrend) 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_window_hanning(self): if self.NFFT_density is None: return ydata = np.arange(self.NFFT_density) ydata1 = ydata+5 ydata2 = ydata+3.3 windowVals = mlab.window_hanning(np.ones_like(ydata1)) ycontrol1 = ydata1 * windowVals ycontrol2 = mlab.window_hanning(ydata2) ydata = np.vstack([ydata1, ydata2]) ycontrol = np.vstack([ycontrol1, ycontrol2]) ydata = np.tile(ydata, (20, 1)) ycontrol = np.tile(ycontrol, (20, 1)) ydatab = ydata.T.flatten() ydataf = ydata.flatten() ycontrol = ycontrol.flatten() spec_g, fsp_g = mlab.psd(x=ydataf, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, window=mlab.window_hanning) spec_b, fsp_b = mlab.psd(x=ydatab, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, window=mlab.window_hanning) spec_c, fsp_c = mlab.psd(x=ycontrol, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, window=mlab.window_none) spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum() 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_window_hanning_detrend_linear(self): if self.NFFT_density is None: return ydata = np.arange(self.NFFT_density) ycontrol = np.zeros(self.NFFT_density) ydata1 = ydata+5 ydata2 = ydata+3.3 ycontrol1 = ycontrol ycontrol2 = ycontrol windowVals = mlab.window_hanning(np.ones_like(ycontrol1)) ycontrol1 = ycontrol1 * windowVals ycontrol2 = mlab.window_hanning(ycontrol2) ydata = np.vstack([ydata1, ydata2]) ycontrol = np.vstack([ycontrol1, ycontrol2]) ydata = np.tile(ydata, (20, 1)) ycontrol = np.tile(ycontrol, (20, 1)) ydatab = ydata.T.flatten() ydataf = ydata.flatten() ycontrol = ycontrol.flatten() spec_g, fsp_g = mlab.psd(x=ydataf, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, detrend=mlab.detrend_linear, window=mlab.window_hanning) spec_b, fsp_b = mlab.psd(x=ydatab, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, detrend=mlab.detrend_linear, window=mlab.window_hanning) spec_c, fsp_c = mlab.psd(x=ycontrol, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=0, sides=self.sides, window=mlab.window_none) spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum() 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_windowarray(self): freqs = self.freqs_density spec, fsp = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides, window=np.ones(self.NFFT_density_real)) assert_allclose(fsp, freqs, atol=1e-06) assert spec.shape == freqs.shape def test_psd_windowarray_scale_by_freq(self): win = mlab.window_hanning(np.ones(self.NFFT_density_real)) spec, fsp = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides, window=mlab.window_hanning) spec_s, fsp_s = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides, window=mlab.window_hanning, scale_by_freq=True) spec_n, fsp_n = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides, window=mlab.window_hanning, scale_by_freq=False) assert_array_equal(fsp, fsp_s) assert_array_equal(fsp, fsp_n) assert_array_equal(spec, spec_s) assert_allclose(spec_s*(win**2).sum(), spec_n/self.Fs*win.sum()**2, atol=1e-08) @pytest.mark.parametrize( "kind", ["complex", "magnitude", "angle", "phase"]) def test_spectrum(self, kind): freqs = self.freqs_spectrum spec, fsp = getattr(mlab, f"{kind}_spectrum")( x=self.y, Fs=self.Fs, sides=self.sides, pad_to=self.pad_to_spectrum) assert_allclose(fsp, freqs, atol=1e-06) assert spec.shape == freqs.shape if kind == "magnitude": self.check_maxfreq(spec, fsp, self.fstims) self.check_freqs(spec, freqs, fsp, self.fstims) @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, Fs=self.Fs, noverlap=self.nover_specgram, pad_to=self.pad_to_specgram, sides=self.sides, **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] 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.""" with pytest.warns(UserWarning, match="Only one segment is calculated"): mlab.specgram(x=self.y, NFFT=len(self.y), Fs=self.Fs) def test_psd_csd_equal(self): Pxx, freqsxx = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides) Pxy, freqsxy = mlab.csd(x=self.y, y=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides) assert_array_almost_equal_nulp(Pxx, Pxy) assert_array_equal(freqsxx, freqsxy) @pytest.mark.parametrize("mode", ["default", "psd"]) def test_specgram_auto_default_psd_equal(self, mode): """ Test that mlab.specgram without mode and with mode 'default' and 'psd' are all the same. """ speca, freqspeca, ta = 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) specb, freqspecb, tb = 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=mode) assert_array_equal(speca, specb) assert_array_equal(freqspeca, freqspecb) assert_array_equal(ta, tb) @pytest.mark.parametrize( "mode, conv", [ ("magnitude", np.abs), ("angle", np.angle), ("phase", lambda x: np.unwrap(np.angle(x), axis=0)) ]) def test_specgram_complex_equivalent(self, mode, conv): specc, freqspecc, tc = 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, freqspecm, tm = 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=mode) assert_array_equal(freqspecc, freqspecm) assert_array_equal(tc, tm) assert_allclose(conv(specc), specm, atol=1e-06) def test_psd_windowarray_equal(self): win = mlab.window_hanning(np.ones(self.NFFT_density_real)) speca, fspa = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides, window=win) specb, fspb = mlab.psd(x=self.y, NFFT=self.NFFT_density, Fs=self.Fs, noverlap=self.nover_density, pad_to=self.pad_to_density, sides=self.sides) assert_array_equal(fspa, fspb) assert_allclose(speca, specb, atol=1e-08) # extra test for cohere... def test_cohere(): N = 1024 np.random.seed(19680801) x = np.random.randn(N) # phase offset y = np.roll(x, 20) # high-freq roll-off y = np.convolve(y, np.ones(20) / 20., mode='same') cohsq, f = mlab.cohere(x, y, NFFT=256, Fs=2, noverlap=128) assert_allclose(np.mean(cohsq), 0.837, atol=1.e-3) assert np.isreal(np.mean(cohsq)) # ***************************************************************** # These Tests were taken from SCIPY with some minor modifications # this can be retrieved from: # https://github.com/scipy/scipy/blob/master/scipy/stats/tests/test_kdeoth.py # ***************************************************************** class TestGaussianKDE: def test_kde_integer_input(self): """Regression test for #1181.""" x1 = np.arange(5) kde = mlab.GaussianKDE(x1) y_expected = [0.13480721, 0.18222869, 0.19514935, 0.18222869, 0.13480721] np.testing.assert_array_almost_equal(kde(x1), y_expected, decimal=6) def test_gaussian_kde_covariance_caching(self): x1 = np.array([-7, -5, 1, 4, 5], dtype=float) xs = np.linspace(-10, 10, num=5) # These expected values are from scipy 0.10, before some changes to # gaussian_kde. They were not compared with any external reference. y_expected = [0.02463386, 0.04689208, 0.05395444, 0.05337754, 0.01664475] # set it to the default bandwidth. kde2 = mlab.GaussianKDE(x1, 'scott') y2 = kde2(xs) np.testing.assert_array_almost_equal(y_expected, y2, decimal=7) def test_kde_bandwidth_method(self): np.random.seed(8765678) n_basesample = 50 xn = np.random.randn(n_basesample) # Default gkde = mlab.GaussianKDE(xn) # Supply a callable gkde2 = mlab.GaussianKDE(xn, 'scott') # Supply a scalar gkde3 = mlab.GaussianKDE(xn, bw_method=gkde.factor) xs = np.linspace(-7, 7, 51) kdepdf = gkde.evaluate(xs) kdepdf2 = gkde2.evaluate(xs) assert kdepdf.all() == kdepdf2.all() kdepdf3 = gkde3.evaluate(xs) assert kdepdf.all() == kdepdf3.all() class TestGaussianKDECustom: def test_no_data(self): """Pass no data into the GaussianKDE class.""" with pytest.raises(ValueError): mlab.GaussianKDE([]) def test_single_dataset_element(self): """Pass a single dataset element into the GaussianKDE class.""" with pytest.raises(ValueError): mlab.GaussianKDE([42]) def test_silverman_multidim_dataset(self): """Test silverman's for a multi-dimensional array.""" x1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) with pytest.raises(np.linalg.LinAlgError): mlab.GaussianKDE(x1, "silverman") def test_silverman_singledim_dataset(self): """Test silverman's output for a single dimension list.""" x1 = np.array([-7, -5, 1, 4, 5]) mygauss = mlab.GaussianKDE(x1, "silverman") y_expected = 0.76770389927475502 assert_almost_equal(mygauss.covariance_factor(), y_expected, 7) def test_scott_multidim_dataset(self): """Test scott's output for a multi-dimensional array.""" x1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) with pytest.raises(np.linalg.LinAlgError): mlab.GaussianKDE(x1, "scott") def test_scott_singledim_dataset(self): """Test scott's output a single-dimensional array.""" x1 = np.array([-7, -5, 1, 4, 5]) mygauss = mlab.GaussianKDE(x1, "scott") y_expected = 0.72477966367769553 assert_almost_equal(mygauss.covariance_factor(), y_expected, 7) def test_scalar_empty_dataset(self): """Test the scalar's cov factor for an empty array.""" with pytest.raises(ValueError): mlab.GaussianKDE([], bw_method=5) def test_scalar_covariance_dataset(self): """Test a scalar's cov factor.""" np.random.seed(8765678) n_basesample = 50 multidim_data = [np.random.randn(n_basesample) for i in range(5)] kde = mlab.GaussianKDE(multidim_data, bw_method=0.5) assert kde.covariance_factor() == 0.5 def test_callable_covariance_dataset(self): """Test the callable's cov factor for a multi-dimensional array.""" np.random.seed(8765678) n_basesample = 50 multidim_data = [np.random.randn(n_basesample) for i in range(5)] def callable_fun(x): return 0.55 kde = mlab.GaussianKDE(multidim_data, bw_method=callable_fun) assert kde.covariance_factor() == 0.55 def test_callable_singledim_dataset(self): """Test the callable's cov factor for a single-dimensional array.""" np.random.seed(8765678) n_basesample = 50 multidim_data = np.random.randn(n_basesample) kde = mlab.GaussianKDE(multidim_data, bw_method='silverman') y_expected = 0.48438841363348911 assert_almost_equal(kde.covariance_factor(), y_expected, 7) def test_wrong_bw_method(self): """Test the error message that should be called when bw is invalid.""" np.random.seed(8765678) n_basesample = 50 data = np.random.randn(n_basesample) with pytest.raises(ValueError): mlab.GaussianKDE(data, bw_method="invalid") class TestGaussianKDEEvaluate: def test_evaluate_diff_dim(self): """ Test the evaluate method when the dim's of dataset and points have different dimensions. """ x1 = np.arange(3, 10, 2) kde = mlab.GaussianKDE(x1) x2 = np.arange(3, 12, 2) y_expected = [ 0.08797252, 0.11774109, 0.11774109, 0.08797252, 0.0370153 ] y = kde.evaluate(x2) np.testing.assert_array_almost_equal(y, y_expected, 7) def test_evaluate_inv_dim(self): """ Invert the dimensions; i.e., for a dataset of dimension 1 [3, 2, 4], the points should have a dimension of 3 [[3], [2], [4]]. """ np.random.seed(8765678) n_basesample = 50 multidim_data = np.random.randn(n_basesample) kde = mlab.GaussianKDE(multidim_data) x2 = [[1], [2], [3]] with pytest.raises(ValueError): kde.evaluate(x2) def test_evaluate_dim_and_num(self): """Tests if evaluated against a one by one array""" x1 = np.arange(3, 10, 2) x2 = np.array([3]) kde = mlab.GaussianKDE(x1) y_expected = [0.08797252] y = kde.evaluate(x2) np.testing.assert_array_almost_equal(y, y_expected, 7) def test_evaluate_point_dim_not_one(self): x1 = np.arange(3, 10, 2) x2 = [np.arange(3, 10, 2), np.arange(3, 10, 2)] kde = mlab.GaussianKDE(x1) with pytest.raises(ValueError): kde.evaluate(x2) def test_evaluate_equal_dim_and_num_lt(self): x1 = np.arange(3, 10, 2) x2 = np.arange(3, 8, 2) kde = mlab.GaussianKDE(x1) y_expected = [0.08797252, 0.11774109, 0.11774109] y = kde.evaluate(x2) np.testing.assert_array_almost_equal(y, y_expected, 7) def test_psd_onesided_norm(): u = np.array([0, 1, 2, 3, 1, 2, 1]) dt = 1.0 Su = np.abs(np.fft.fft(u) * dt)**2 / (dt * u.size) P, f = mlab.psd(u, NFFT=u.size, Fs=1/dt, window=mlab.window_none, detrend=mlab.detrend_none, noverlap=0, pad_to=None, scale_by_freq=None, sides='onesided') Su_1side = np.append([Su[0]], Su[1:4] + Su[4:][::-1]) assert_allclose(P, Su_1side, atol=1e-06) def test_psd_oversampling(): """Test the case len(x) < NFFT for psd().""" u = np.array([0, 1, 2, 3, 1, 2, 1]) dt = 1.0 Su = np.abs(np.fft.fft(u) * dt)**2 / (dt * u.size) P, f = mlab.psd(u, NFFT=u.size*2, Fs=1/dt, window=mlab.window_none, detrend=mlab.detrend_none, noverlap=0, pad_to=None, scale_by_freq=None, sides='onesided') Su_1side = np.append([Su[0]], Su[1:4] + Su[4:][::-1]) assert_almost_equal(np.sum(P), np.sum(Su_1side)) # same energy