Skip to content

Commit ff304e4

Browse files
tacaswelltacaswell
committed
Merge pull request matplotlib#2571 from toddrjen/specfix
This fixes thee probllem brought up in the mailing list with the recent spectrum improvements
2 parents c91589c + 0425fbf commit ff304e4

File tree

2 files changed

+330
-2
lines changed

2 files changed

+330
-2
lines changed

‎lib/matplotlib/mlab.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -730,10 +730,10 @@ def _spectral_helper(x, y=None, NFFT=None, Fs=None, detrend_func=None,
730730
numFreqs = pad_to//2 + 1
731731
scaling_factor = 2.
732732

733-
result = stride_windows(x, NFFT, noverlap)
733+
result = stride_windows(x, NFFT, noverlap, axis=0)
734+
result = detrend(result, detrend_func, axis=0)
734735
result, windowVals = apply_window(result, window, axis=0,
735736
return_window=True)
736-
result = detrend(result, detrend_func, axis=0)
737737
result = np.fft.fft(result, n=pad_to, axis=0)[:numFreqs, :]
738738
freqs = np.fft.fftfreq(pad_to, 1/Fs)[:numFreqs]
739739

‎lib/matplotlib/tests/test_mlab.py‎

Lines changed: 328 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,24 @@ def test_stride_windows_n13_noverlapn3_axis1(self):
259259
assert_equal((6, 13), y.shape)
260260
assert_true(self.get_base(y) is x)
261261

262+
def test_stride_windows_n32_noverlap0_axis0_unflatten(self):
263+
n = 32
264+
x = np.arange(n)[np.newaxis]
265+
x1 = np.tile(x, (21, 1))
266+
x2 = x1.flatten()
267+
y = mlab.stride_windows(x2, n)
268+
assert_equal(y.shape, x1.T.shape)
269+
assert_array_equal(y, x1.T)
270+
271+
def test_stride_windows_n32_noverlap0_axis1_unflatten(self):
272+
n = 32
273+
x = np.arange(n)[np.newaxis]
274+
x1 = np.tile(x, (21, 1))
275+
x2 = x1.flatten()
276+
y = mlab.stride_windows(x2, n, axis=1)
277+
assert_equal(y.shape, x1.shape)
278+
assert_array_equal(y, x1)
279+
262280

263281
class csv_testcase(CleanupTestCase):
264282
def setUp(self):
@@ -503,6 +521,51 @@ def test_apply_window_stride_windows_hanning_2D_n13_noverlapn3_axis0(self):
503521
assert_not_equal(x.shape, y.shape)
504522
assert_allclose(yt, y, atol=1e-06)
505523

524+
def test_apply_window_hanning_2D_stack_axis1(self):
525+
ydata = np.arange(32)
526+
ydata1 = ydata+5
527+
ydata2 = ydata+3.3
528+
ycontrol1 = mlab.apply_window(ydata1, mlab.window_hanning)
529+
ycontrol2 = mlab.window_hanning(ydata2)
530+
ydata = np.vstack([ydata1, ydata2])
531+
ycontrol = np.vstack([ycontrol1, ycontrol2])
532+
ydata = np.tile(ydata, (20, 1))
533+
ycontrol = np.tile(ycontrol, (20, 1))
534+
result = mlab.apply_window(ydata, mlab.window_hanning, axis=1,
535+
return_window=False)
536+
assert_allclose(ycontrol, result, atol=1e-08)
537+
538+
def test_apply_window_hanning_2D_stack_windows_axis1(self):
539+
ydata = np.arange(32)
540+
ydata1 = ydata+5
541+
ydata2 = ydata+3.3
542+
ycontrol1 = mlab.apply_window(ydata1, mlab.window_hanning)
543+
ycontrol2 = mlab.window_hanning(ydata2)
544+
ydata = np.vstack([ydata1, ydata2])
545+
ycontrol = np.vstack([ycontrol1, ycontrol2])
546+
ydata = np.tile(ydata, (20, 1))
547+
ycontrol = np.tile(ycontrol, (20, 1))
548+
result = mlab.apply_window(ydata, mlab.window_hanning, axis=1,
549+
return_window=False)
550+
assert_allclose(ycontrol, result, atol=1e-08)
551+
552+
def test_apply_window_hanning_2D_stack_windows_axis1_unflatten(self):
553+
n = 32
554+
ydata = np.arange(n)
555+
ydata1 = ydata+5
556+
ydata2 = ydata+3.3
557+
ycontrol1 = mlab.apply_window(ydata1, mlab.window_hanning)
558+
ycontrol2 = mlab.window_hanning(ydata2)
559+
ydata = np.vstack([ydata1, ydata2])
560+
ycontrol = np.vstack([ycontrol1, ycontrol2])
561+
ydata = np.tile(ydata, (20, 1))
562+
ycontrol = np.tile(ycontrol, (20, 1))
563+
ydata = ydata.flatten()
564+
ydata1 = mlab.stride_windows(ydata, 32, noverlap=0, axis=0)
565+
result = mlab.apply_window(ydata1, mlab.window_hanning, axis=0,
566+
return_window=False)
567+
assert_allclose(ycontrol.T, result, atol=1e-08)
568+
506569

507570
class detrend_testcase(CleanupTestCase):
508571
def setUp(self):
@@ -1460,6 +1523,243 @@ def test_psd(self):
14601523
assert_equal(spec.shape, freqs.shape)
14611524
self.check_freqs(spec, freqs, fsp, self.fstims)
14621525

1526+
def test_psd_detrend_mean_func_offset(self):
1527+
if self.NFFT_density is None:
1528+
return
1529+
freqs = self.freqs_density
1530+
ydata = np.zeros(self.NFFT_density)
1531+
ydata1 = ydata+5
1532+
ydata2 = ydata+3.3
1533+
ydata = np.vstack([ydata1, ydata2])
1534+
ydata = np.tile(ydata, (20, 1))
1535+
ydatab = ydata.T.flatten()
1536+
ydata = ydata.flatten()
1537+
ycontrol = np.zeros_like(ydata)
1538+
spec_g, fsp_g = mlab.psd(x=ydata,
1539+
NFFT=self.NFFT_density,
1540+
Fs=self.Fs,
1541+
noverlap=0,
1542+
sides=self.sides,
1543+
detrend=mlab.detrend_mean)
1544+
spec_b, fsp_b = mlab.psd(x=ydatab,
1545+
NFFT=self.NFFT_density,
1546+
Fs=self.Fs,
1547+
noverlap=0,
1548+
sides=self.sides,
1549+
detrend=mlab.detrend_mean)
1550+
spec_c, fsp_c = mlab.psd(x=ycontrol,
1551+
NFFT=self.NFFT_density,
1552+
Fs=self.Fs,
1553+
noverlap=0,
1554+
sides=self.sides)
1555+
assert_array_equal(fsp_g, fsp_c)
1556+
assert_array_equal(fsp_b, fsp_c)
1557+
assert_allclose(spec_g, spec_c, atol=1e-08)
1558+
# these should not be almost equal
1559+
assert_raises(AssertionError,
1560+
assert_allclose, spec_b, spec_c, atol=1e-08)
1561+
1562+
def test_psd_detrend_mean_str_offset(self):
1563+
if self.NFFT_density is None:
1564+
return
1565+
freqs = self.freqs_density
1566+
ydata = np.zeros(self.NFFT_density)
1567+
ydata1 = ydata+5
1568+
ydata2 = ydata+3.3
1569+
ydata = np.vstack([ydata1, ydata2])
1570+
ydata = np.tile(ydata, (20, 1))
1571+
ydatab = ydata.T.flatten()
1572+
ydata = ydata.flatten()
1573+
ycontrol = np.zeros_like(ydata)
1574+
spec_g, fsp_g = mlab.psd(x=ydata,
1575+
NFFT=self.NFFT_density,
1576+
Fs=self.Fs,
1577+
noverlap=0,
1578+
sides=self.sides,
1579+
detrend='mean')
1580+
spec_b, fsp_b = mlab.psd(x=ydatab,
1581+
NFFT=self.NFFT_density,
1582+
Fs=self.Fs,
1583+
noverlap=0,
1584+
sides=self.sides,
1585+
detrend='mean')
1586+
spec_c, fsp_c = mlab.psd(x=ycontrol,
1587+
NFFT=self.NFFT_density,
1588+
Fs=self.Fs,
1589+
noverlap=0,
1590+
sides=self.sides)
1591+
assert_array_equal(fsp_g, fsp_c)
1592+
assert_array_equal(fsp_b, fsp_c)
1593+
assert_allclose(spec_g, spec_c, atol=1e-08)
1594+
# these should not be almost equal
1595+
assert_raises(AssertionError,
1596+
assert_allclose, spec_b, spec_c, atol=1e-08)
1597+
1598+
def test_psd_detrend_linear_func_trend(self):
1599+
if self.NFFT_density is None:
1600+
return
1601+
freqs = self.freqs_density
1602+
ydata = np.arange(self.NFFT_density)
1603+
ydata1 = ydata+5
1604+
ydata2 = ydata+3.3
1605+
ydata = np.vstack([ydata1, ydata2])
1606+
ydata = np.tile(ydata, (20, 1))
1607+
ydatab = ydata.T.flatten()
1608+
ydata = ydata.flatten()
1609+
ycontrol = np.zeros_like(ydata)
1610+
spec_g, fsp_g = mlab.psd(x=ydata,
1611+
NFFT=self.NFFT_density,
1612+
Fs=self.Fs,
1613+
noverlap=0,
1614+
sides=self.sides,
1615+
detrend=mlab.detrend_linear)
1616+
spec_b, fsp_b = mlab.psd(x=ydatab,
1617+
NFFT=self.NFFT_density,
1618+
Fs=self.Fs,
1619+
noverlap=0,
1620+
sides=self.sides,
1621+
detrend=mlab.detrend_linear)
1622+
spec_c, fsp_c = mlab.psd(x=ycontrol,
1623+
NFFT=self.NFFT_density,
1624+
Fs=self.Fs,
1625+
noverlap=0,
1626+
sides=self.sides)
1627+
assert_array_equal(fsp_g, fsp_c)
1628+
assert_array_equal(fsp_b, fsp_c)
1629+
assert_allclose(spec_g, spec_c, atol=1e-08)
1630+
# these should not be almost equal
1631+
assert_raises(AssertionError,
1632+
assert_allclose, spec_b, spec_c, atol=1e-08)
1633+
1634+
def test_psd_detrend_linear_str_trend(self):
1635+
if self.NFFT_density is None:
1636+
return
1637+
freqs = self.freqs_density
1638+
ydata = np.arange(self.NFFT_density)
1639+
ydata1 = ydata+5
1640+
ydata2 = ydata+3.3
1641+
ydata = np.vstack([ydata1, ydata2])
1642+
ydata = np.tile(ydata, (20, 1))
1643+
ydatab = ydata.T.flatten()
1644+
ydata = ydata.flatten()
1645+
ycontrol = np.zeros_like(ydata)
1646+
spec_g, fsp_g = mlab.psd(x=ydata,
1647+
NFFT=self.NFFT_density,
1648+
Fs=self.Fs,
1649+
noverlap=0,
1650+
sides=self.sides,
1651+
detrend='linear')
1652+
spec_b, fsp_b = mlab.psd(x=ydatab,
1653+
NFFT=self.NFFT_density,
1654+
Fs=self.Fs,
1655+
noverlap=0,
1656+
sides=self.sides,
1657+
detrend='linear')
1658+
spec_c, fsp_c = mlab.psd(x=ycontrol,
1659+
NFFT=self.NFFT_density,
1660+
Fs=self.Fs,
1661+
noverlap=0,
1662+
sides=self.sides)
1663+
assert_array_equal(fsp_g, fsp_c)
1664+
assert_array_equal(fsp_b, fsp_c)
1665+
assert_allclose(spec_g, spec_c, atol=1e-08)
1666+
# these should not be almost equal
1667+
assert_raises(AssertionError,
1668+
assert_allclose, spec_b, spec_c, atol=1e-08)
1669+
1670+
def test_psd_window_hanning(self):
1671+
if self.NFFT_density is None:
1672+
return
1673+
freqs = self.freqs_density
1674+
ydata = np.arange(self.NFFT_density)
1675+
ydata1 = ydata+5
1676+
ydata2 = ydata+3.3
1677+
ycontrol1, windowVals = mlab.apply_window(ydata1,
1678+
mlab.window_hanning,
1679+
return_window=True)
1680+
ycontrol2 = mlab.window_hanning(ydata2)
1681+
ydata = np.vstack([ydata1, ydata2])
1682+
ycontrol = np.vstack([ycontrol1, ycontrol2])
1683+
ydata = np.tile(ydata, (20, 1))
1684+
ycontrol = np.tile(ycontrol, (20, 1))
1685+
ydatab = ydata.T.flatten()
1686+
ydataf = ydata.flatten()
1687+
ycontrol = ycontrol.flatten()
1688+
spec_g, fsp_g = mlab.psd(x=ydataf,
1689+
NFFT=self.NFFT_density,
1690+
Fs=self.Fs,
1691+
noverlap=0,
1692+
sides=self.sides,
1693+
window=mlab.window_hanning)
1694+
spec_b, fsp_b = mlab.psd(x=ydatab,
1695+
NFFT=self.NFFT_density,
1696+
Fs=self.Fs,
1697+
noverlap=0,
1698+
sides=self.sides,
1699+
window=mlab.window_hanning)
1700+
spec_c, fsp_c = mlab.psd(x=ycontrol,
1701+
NFFT=self.NFFT_density,
1702+
Fs=self.Fs,
1703+
noverlap=0,
1704+
sides=self.sides,
1705+
window=mlab.window_none)
1706+
spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum()
1707+
assert_array_equal(fsp_g, fsp_c)
1708+
assert_array_equal(fsp_b, fsp_c)
1709+
assert_allclose(spec_g, spec_c, atol=1e-08)
1710+
# these should not be almost equal
1711+
assert_raises(AssertionError,
1712+
assert_allclose, spec_b, spec_c, atol=1e-08)
1713+
1714+
def test_psd_window_hanning_detrend_linear(self):
1715+
if self.NFFT_density is None:
1716+
return
1717+
freqs = self.freqs_density
1718+
ydata = np.arange(self.NFFT_density)
1719+
ycontrol = np.zeros(self.NFFT_density)
1720+
ydata1 = ydata+5
1721+
ydata2 = ydata+3.3
1722+
ycontrol1 = ycontrol
1723+
ycontrol2 = ycontrol
1724+
ycontrol1, windowVals = mlab.apply_window(ycontrol1,
1725+
mlab.window_hanning,
1726+
return_window=True)
1727+
ycontrol2 = mlab.window_hanning(ycontrol2)
1728+
ydata = np.vstack([ydata1, ydata2])
1729+
ycontrol = np.vstack([ycontrol1, ycontrol2])
1730+
ydata = np.tile(ydata, (20, 1))
1731+
ycontrol = np.tile(ycontrol, (20, 1))
1732+
ydatab = ydata.T.flatten()
1733+
ydataf = ydata.flatten()
1734+
ycontrol = ycontrol.flatten()
1735+
spec_g, fsp_g = mlab.psd(x=ydataf,
1736+
NFFT=self.NFFT_density,
1737+
Fs=self.Fs,
1738+
noverlap=0,
1739+
sides=self.sides,
1740+
detrend=mlab.detrend_linear,
1741+
window=mlab.window_hanning)
1742+
spec_b, fsp_b = mlab.psd(x=ydatab,
1743+
NFFT=self.NFFT_density,
1744+
Fs=self.Fs,
1745+
noverlap=0,
1746+
sides=self.sides,
1747+
detrend=mlab.detrend_linear,
1748+
window=mlab.window_hanning)
1749+
spec_c, fsp_c = mlab.psd(x=ycontrol,
1750+
NFFT=self.NFFT_density,
1751+
Fs=self.Fs,
1752+
noverlap=0,
1753+
sides=self.sides,
1754+
window=mlab.window_none)
1755+
spec_c *= len(ycontrol1)/(np.abs(windowVals)**2).sum()
1756+
assert_array_equal(fsp_g, fsp_c)
1757+
assert_array_equal(fsp_b, fsp_c)
1758+
assert_allclose(spec_g, spec_c, atol=1e-08)
1759+
# these should not be almost equal
1760+
assert_raises(AssertionError,
1761+
assert_allclose, spec_b, spec_c, atol=1e-08)
1762+
14631763
def test_psd_windowarray(self):
14641764
freqs = self.freqs_density
14651765
spec, fsp = mlab.psd(x=self.y,
@@ -1472,6 +1772,34 @@ def test_psd_windowarray(self):
14721772
assert_allclose(fsp, freqs, atol=1e-06)
14731773
assert_equal(spec.shape, freqs.shape)
14741774

1775+
def test_psd_windowarray_scale_by_freq(self):
1776+
freqs = self.freqs_density
1777+
spec, fsp = mlab.psd(x=self.y,
1778+
NFFT=self.NFFT_density,
1779+
Fs=self.Fs,
1780+
noverlap=self.nover_density,
1781+
pad_to=self.pad_to_density,
1782+
sides=self.sides)
1783+
spec_s, fsp_s = mlab.psd(x=self.y,
1784+
NFFT=self.NFFT_density,
1785+
Fs=self.Fs,
1786+
noverlap=self.nover_density,
1787+
pad_to=self.pad_to_density,
1788+
sides=self.sides,
1789+
scale_by_freq=True)
1790+
spec_n, fsp_n = mlab.psd(x=self.y,
1791+
NFFT=self.NFFT_density,
1792+
Fs=self.Fs,
1793+
noverlap=self.nover_density,
1794+
pad_to=self.pad_to_density,
1795+
sides=self.sides,
1796+
scale_by_freq=False)
1797+
1798+
assert_array_equal(fsp, fsp_s)
1799+
assert_array_equal(fsp, fsp_n)
1800+
assert_array_equal(spec, spec_s)
1801+
assert_allclose(spec_s, spec_n/self.Fs, atol=1e-08)
1802+
14751803
def test_complex_spectrum(self):
14761804
freqs = self.freqs_spectrum
14771805
spec, fsp = mlab.complex_spectrum(x=self.y,

0 commit comments

Comments
 (0)