test_array_tools.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. import numpy as np
  2. from numpy.testing import assert_array_equal
  3. from pytest import raises as assert_raises
  4. from scipy.signal._arraytools import (axis_slice, axis_reverse,
  5. odd_ext, even_ext, const_ext, zero_ext)
  6. class TestArrayTools:
  7. def test_axis_slice(self):
  8. a = np.arange(12).reshape(3, 4)
  9. s = axis_slice(a, start=0, stop=1, axis=0)
  10. assert_array_equal(s, a[0:1, :])
  11. s = axis_slice(a, start=-1, axis=0)
  12. assert_array_equal(s, a[-1:, :])
  13. s = axis_slice(a, start=0, stop=1, axis=1)
  14. assert_array_equal(s, a[:, 0:1])
  15. s = axis_slice(a, start=-1, axis=1)
  16. assert_array_equal(s, a[:, -1:])
  17. s = axis_slice(a, start=0, step=2, axis=0)
  18. assert_array_equal(s, a[::2, :])
  19. s = axis_slice(a, start=0, step=2, axis=1)
  20. assert_array_equal(s, a[:, ::2])
  21. def test_axis_reverse(self):
  22. a = np.arange(12).reshape(3, 4)
  23. r = axis_reverse(a, axis=0)
  24. assert_array_equal(r, a[::-1, :])
  25. r = axis_reverse(a, axis=1)
  26. assert_array_equal(r, a[:, ::-1])
  27. def test_odd_ext(self):
  28. a = np.array([[1, 2, 3, 4, 5],
  29. [9, 8, 7, 6, 5]])
  30. odd = odd_ext(a, 2, axis=1)
  31. expected = np.array([[-1, 0, 1, 2, 3, 4, 5, 6, 7],
  32. [11, 10, 9, 8, 7, 6, 5, 4, 3]])
  33. assert_array_equal(odd, expected)
  34. odd = odd_ext(a, 1, axis=0)
  35. expected = np.array([[-7, -4, -1, 2, 5],
  36. [1, 2, 3, 4, 5],
  37. [9, 8, 7, 6, 5],
  38. [17, 14, 11, 8, 5]])
  39. assert_array_equal(odd, expected)
  40. assert_raises(ValueError, odd_ext, a, 2, axis=0)
  41. assert_raises(ValueError, odd_ext, a, 5, axis=1)
  42. def test_even_ext(self):
  43. a = np.array([[1, 2, 3, 4, 5],
  44. [9, 8, 7, 6, 5]])
  45. even = even_ext(a, 2, axis=1)
  46. expected = np.array([[3, 2, 1, 2, 3, 4, 5, 4, 3],
  47. [7, 8, 9, 8, 7, 6, 5, 6, 7]])
  48. assert_array_equal(even, expected)
  49. even = even_ext(a, 1, axis=0)
  50. expected = np.array([[9, 8, 7, 6, 5],
  51. [1, 2, 3, 4, 5],
  52. [9, 8, 7, 6, 5],
  53. [1, 2, 3, 4, 5]])
  54. assert_array_equal(even, expected)
  55. assert_raises(ValueError, even_ext, a, 2, axis=0)
  56. assert_raises(ValueError, even_ext, a, 5, axis=1)
  57. def test_const_ext(self):
  58. a = np.array([[1, 2, 3, 4, 5],
  59. [9, 8, 7, 6, 5]])
  60. const = const_ext(a, 2, axis=1)
  61. expected = np.array([[1, 1, 1, 2, 3, 4, 5, 5, 5],
  62. [9, 9, 9, 8, 7, 6, 5, 5, 5]])
  63. assert_array_equal(const, expected)
  64. const = const_ext(a, 1, axis=0)
  65. expected = np.array([[1, 2, 3, 4, 5],
  66. [1, 2, 3, 4, 5],
  67. [9, 8, 7, 6, 5],
  68. [9, 8, 7, 6, 5]])
  69. assert_array_equal(const, expected)
  70. def test_zero_ext(self):
  71. a = np.array([[1, 2, 3, 4, 5],
  72. [9, 8, 7, 6, 5]])
  73. zero = zero_ext(a, 2, axis=1)
  74. expected = np.array([[0, 0, 1, 2, 3, 4, 5, 0, 0],
  75. [0, 0, 9, 8, 7, 6, 5, 0, 0]])
  76. assert_array_equal(zero, expected)
  77. zero = zero_ext(a, 1, axis=0)
  78. expected = np.array([[0, 0, 0, 0, 0],
  79. [1, 2, 3, 4, 5],
  80. [9, 8, 7, 6, 5],
  81. [0, 0, 0, 0, 0]])
  82. assert_array_equal(zero, expected)