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)