Skip to content

Commit b2c8f88

Browse files
authored
Test function for scale factor of flattop window
Add new test function for scale factor of flattop window. Also remove the unnecessary `np.abs()` on window in other functions.
1 parent cebd3d7 commit b2c8f88

File tree

1 file changed

+40
-2
lines changed

1 file changed

+40
-2
lines changed

lib/matplotlib/tests/test_mlab.py

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -615,7 +615,7 @@ def test_psd_window_hanning(self):
615615
noverlap=0,
616616
sides=self.sides,
617617
window=mlab.window_none)
618-
spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum()
618+
spec_c *= len(ycontrol1)/(windowVals**2).sum()
619619
assert_array_equal(fsp_g, fsp_c)
620620
assert_array_equal(fsp_b, fsp_c)
621621
assert_allclose(spec_g, spec_c, atol=1e-08)
@@ -662,14 +662,52 @@ def test_psd_window_hanning_detrend_linear(self):
662662
noverlap=0,
663663
sides=self.sides,
664664
window=mlab.window_none)
665-
spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum()
665+
spec_c *= len(ycontrol1)/(windowVals**2).sum()
666666
assert_array_equal(fsp_g, fsp_c)
667667
assert_array_equal(fsp_b, fsp_c)
668668
assert_allclose(spec_g, spec_c, atol=1e-08)
669669
# these should not be almost equal
670670
with pytest.raises(AssertionError):
671671
assert_allclose(spec_b, spec_c, atol=1e-08)
672672

673+
def test_psd_window_flattop(self):
674+
# flattop window
675+
# adaption from https://github.com/scipy/scipy/blob/v1.10.0/scipy/signal/windows/_windows.py#L562-L622
676+
if self.NFFT_density_real <=1:
677+
win = np.ones(self.NFFT_density_real)
678+
else:
679+
a = [0.21557895, 0.41663158, 0.277263158, 0.083578947, 0.006947368]
680+
fac = np.linspace(-np.pi, np.pi, self.NFFT_density_real)
681+
win = np.zeros(self.NFFT_density_real)
682+
for k in range(len(a)):
683+
win += a[k] * np.cos(k * fac)
684+
685+
spec, fsp = mlab.psd(x=self.y,
686+
NFFT=self.NFFT_density,
687+
Fs=self.Fs,
688+
noverlap=0,
689+
sides=self.sides,
690+
window=win,
691+
scale_by_freq=False)
692+
spec_a, fsp_a = mlab.psd(x=self.y,
693+
NFFT=self.NFFT_density,
694+
Fs=self.Fs,
695+
noverlap=0,
696+
sides=self.sides,
697+
window=win)
698+
spec_b, fsp_b = mlab.psd(x=self.y * win,
699+
NFFT=self.NFFT_density,
700+
Fs=self.Fs,
701+
noverlap=0,
702+
sides=self.sides,
703+
window=mlab.window_none)
704+
assert_allclose(spec*win.sum()**2,
705+
spec_a*self.Fs*(win**2).sum(),
706+
atol=1e-08)
707+
assert_allclose(spec*win.sum()**2,
708+
spec_b*self.Fs*self.NFFT_density,
709+
atol=1e-08)
710+
673711
def test_psd_windowarray(self):
674712
freqs = self.freqs_density
675713
spec, fsp = mlab.psd(x=self.y,

0 commit comments

Comments
 (0)