test_wavelets.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  1. from __future__ import division, print_function, absolute_import
  2. import numpy as np
  3. from numpy.testing import assert_equal, \
  4. assert_array_equal, assert_array_almost_equal, assert_array_less, assert_
  5. from scipy._lib.six import xrange
  6. from scipy.signal import wavelets
  7. class TestWavelets(object):
  8. def test_qmf(self):
  9. assert_array_equal(wavelets.qmf([1, 1]), [1, -1])
  10. def test_daub(self):
  11. for i in xrange(1, 15):
  12. assert_equal(len(wavelets.daub(i)), i * 2)
  13. def test_cascade(self):
  14. for J in xrange(1, 7):
  15. for i in xrange(1, 5):
  16. lpcoef = wavelets.daub(i)
  17. k = len(lpcoef)
  18. x, phi, psi = wavelets.cascade(lpcoef, J)
  19. assert_(len(x) == len(phi) == len(psi))
  20. assert_equal(len(x), (k - 1) * 2 ** J)
  21. def test_morlet(self):
  22. x = wavelets.morlet(50, 4.1, complete=True)
  23. y = wavelets.morlet(50, 4.1, complete=False)
  24. # Test if complete and incomplete wavelet have same lengths:
  25. assert_equal(len(x), len(y))
  26. # Test if complete wavelet is less than incomplete wavelet:
  27. assert_array_less(x, y)
  28. x = wavelets.morlet(10, 50, complete=False)
  29. y = wavelets.morlet(10, 50, complete=True)
  30. # For large widths complete and incomplete wavelets should be
  31. # identical within numerical precision:
  32. assert_equal(x, y)
  33. # miscellaneous tests:
  34. x = np.array([1.73752399e-09 + 9.84327394e-25j,
  35. 6.49471756e-01 + 0.00000000e+00j,
  36. 1.73752399e-09 - 9.84327394e-25j])
  37. y = wavelets.morlet(3, w=2, complete=True)
  38. assert_array_almost_equal(x, y)
  39. x = np.array([2.00947715e-09 + 9.84327394e-25j,
  40. 7.51125544e-01 + 0.00000000e+00j,
  41. 2.00947715e-09 - 9.84327394e-25j])
  42. y = wavelets.morlet(3, w=2, complete=False)
  43. assert_array_almost_equal(x, y, decimal=2)
  44. x = wavelets.morlet(10000, s=4, complete=True)
  45. y = wavelets.morlet(20000, s=8, complete=True)[5000:15000]
  46. assert_array_almost_equal(x, y, decimal=2)
  47. x = wavelets.morlet(10000, s=4, complete=False)
  48. assert_array_almost_equal(y, x, decimal=2)
  49. y = wavelets.morlet(20000, s=8, complete=False)[5000:15000]
  50. assert_array_almost_equal(x, y, decimal=2)
  51. x = wavelets.morlet(10000, w=3, s=5, complete=True)
  52. y = wavelets.morlet(20000, w=3, s=10, complete=True)[5000:15000]
  53. assert_array_almost_equal(x, y, decimal=2)
  54. x = wavelets.morlet(10000, w=3, s=5, complete=False)
  55. assert_array_almost_equal(y, x, decimal=2)
  56. y = wavelets.morlet(20000, w=3, s=10, complete=False)[5000:15000]
  57. assert_array_almost_equal(x, y, decimal=2)
  58. x = wavelets.morlet(10000, w=7, s=10, complete=True)
  59. y = wavelets.morlet(20000, w=7, s=20, complete=True)[5000:15000]
  60. assert_array_almost_equal(x, y, decimal=2)
  61. x = wavelets.morlet(10000, w=7, s=10, complete=False)
  62. assert_array_almost_equal(x, y, decimal=2)
  63. y = wavelets.morlet(20000, w=7, s=20, complete=False)[5000:15000]
  64. assert_array_almost_equal(x, y, decimal=2)
  65. def test_ricker(self):
  66. w = wavelets.ricker(1.0, 1)
  67. expected = 2 / (np.sqrt(3 * 1.0) * (np.pi ** 0.25))
  68. assert_array_equal(w, expected)
  69. lengths = [5, 11, 15, 51, 101]
  70. for length in lengths:
  71. w = wavelets.ricker(length, 1.0)
  72. assert_(len(w) == length)
  73. max_loc = np.argmax(w)
  74. assert_(max_loc == (length // 2))
  75. points = 100
  76. w = wavelets.ricker(points, 2.0)
  77. half_vec = np.arange(0, points // 2)
  78. #Wavelet should be symmetric
  79. assert_array_almost_equal(w[half_vec], w[-(half_vec + 1)])
  80. #Check zeros
  81. aas = [5, 10, 15, 20, 30]
  82. points = 99
  83. for a in aas:
  84. w = wavelets.ricker(points, a)
  85. vec = np.arange(0, points) - (points - 1.0) / 2
  86. exp_zero1 = np.argmin(np.abs(vec - a))
  87. exp_zero2 = np.argmin(np.abs(vec + a))
  88. assert_array_almost_equal(w[exp_zero1], 0)
  89. assert_array_almost_equal(w[exp_zero2], 0)
  90. def test_cwt(self):
  91. widths = [1.0]
  92. delta_wavelet = lambda s, t: np.array([1])
  93. len_data = 100
  94. test_data = np.sin(np.pi * np.arange(0, len_data) / 10.0)
  95. #Test delta function input gives same data as output
  96. cwt_dat = wavelets.cwt(test_data, delta_wavelet, widths)
  97. assert_(cwt_dat.shape == (len(widths), len_data))
  98. assert_array_almost_equal(test_data, cwt_dat.flatten())
  99. #Check proper shape on output
  100. widths = [1, 3, 4, 5, 10]
  101. cwt_dat = wavelets.cwt(test_data, wavelets.ricker, widths)
  102. assert_(cwt_dat.shape == (len(widths), len_data))
  103. widths = [len_data * 10]
  104. #Note: this wavelet isn't defined quite right, but is fine for this test
  105. flat_wavelet = lambda l, w: np.ones(w) / w
  106. cwt_dat = wavelets.cwt(test_data, flat_wavelet, widths)
  107. assert_array_almost_equal(cwt_dat, np.mean(test_data))