diff options
Diffstat (limited to 'numpy/lib/tests/test_index_tricks.py')
| -rw-r--r-- | numpy/lib/tests/test_index_tricks.py | 176 |
1 files changed, 73 insertions, 103 deletions
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py index 0520ce580..089a7589a 100644 --- a/numpy/lib/tests/test_index_tricks.py +++ b/numpy/lib/tests/test_index_tricks.py @@ -2,9 +2,8 @@ from __future__ import division, absolute_import, print_function import numpy as np from numpy.testing import ( - run_module_suite, assert_, assert_equal, assert_array_equal, - assert_almost_equal, assert_array_almost_equal, assert_raises, - assert_raises_regex + assert_, assert_equal, assert_array_equal, assert_almost_equal, + assert_array_almost_equal, assert_raises, assert_raises_regex ) from numpy.lib.index_tricks import ( mgrid, ndenumerate, fill_diagonal, diag_indices, diag_indices_from, @@ -185,37 +184,6 @@ class TestConcatenator(object): assert_array_equal(d[:5, :], b) assert_array_equal(d[5:, :], c) - def test_matrix(self): - a = [1, 2] - b = [3, 4] - - ab_r = np.r_['r', a, b] - ab_c = np.r_['c', a, b] - - assert_equal(type(ab_r), np.matrix) - assert_equal(type(ab_c), np.matrix) - - assert_equal(np.array(ab_r), [[1,2,3,4]]) - assert_equal(np.array(ab_c), [[1],[2],[3],[4]]) - - assert_raises(ValueError, lambda: np.r_['rc', a, b]) - - def test_matrix_scalar(self): - r = np.r_['r', [1, 2], 3] - assert_equal(type(r), np.matrix) - assert_equal(np.array(r), [[1,2,3]]) - - def test_matrix_builder(self): - a = np.array([1]) - b = np.array([2]) - c = np.array([3]) - d = np.array([4]) - actual = np.r_['a, b; c, d'] - expected = np.bmat([[a, b], [c, d]]) - - assert_equal(actual, expected) - assert_equal(type(actual), type(expected)) - def test_0d(self): assert_equal(r_[0, np.array(1), 2], [0, 1, 2]) assert_equal(r_[[0, 1, 2], np.array(3)], [0, 1, 2, 3]) @@ -285,71 +253,77 @@ def test_c_(): assert_equal(a, [[1, 2, 3, 0, 0, 4, 5, 6]]) -def test_fill_diagonal(): - a = np.zeros((3, 3), int) - fill_diagonal(a, 5) - yield (assert_array_equal, a, - np.array([[5, 0, 0], - [0, 5, 0], - [0, 0, 5]])) - - #Test tall matrix - a = np.zeros((10, 3), int) - fill_diagonal(a, 5) - yield (assert_array_equal, a, - np.array([[5, 0, 0], - [0, 5, 0], - [0, 0, 5], - [0, 0, 0], - [0, 0, 0], - [0, 0, 0], - [0, 0, 0], - [0, 0, 0], - [0, 0, 0], - [0, 0, 0]])) - - #Test tall matrix wrap - a = np.zeros((10, 3), int) - fill_diagonal(a, 5, True) - yield (assert_array_equal, a, - np.array([[5, 0, 0], - [0, 5, 0], - [0, 0, 5], - [0, 0, 0], - [5, 0, 0], - [0, 5, 0], - [0, 0, 5], - [0, 0, 0], - [5, 0, 0], - [0, 5, 0]])) - - #Test wide matrix - a = np.zeros((3, 10), int) - fill_diagonal(a, 5) - yield (assert_array_equal, a, - np.array([[5, 0, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 5, 0, 0, 0, 0, 0, 0, 0, 0], - [0, 0, 5, 0, 0, 0, 0, 0, 0, 0]])) - - # The same function can operate on a 4-d array: - a = np.zeros((3, 3, 3, 3), int) - fill_diagonal(a, 4) - i = np.array([0, 1, 2]) - yield (assert_equal, np.where(a != 0), (i, i, i, i)) +class TestFillDiagonal(object): + def test_basic(self): + a = np.zeros((3, 3), int) + fill_diagonal(a, 5) + assert_array_equal( + a, np.array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5]]) + ) + + def test_tall_matrix(self): + a = np.zeros((10, 3), int) + fill_diagonal(a, 5) + assert_array_equal( + a, np.array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0]]) + ) + + def test_tall_matrix_wrap(self): + a = np.zeros((10, 3), int) + fill_diagonal(a, 5, True) + assert_array_equal( + a, np.array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5], + [0, 0, 0], + [5, 0, 0], + [0, 5, 0], + [0, 0, 5], + [0, 0, 0], + [5, 0, 0], + [0, 5, 0]]) + ) + + def test_wide_matrix(self): + a = np.zeros((3, 10), int) + fill_diagonal(a, 5) + assert_array_equal( + a, np.array([[5, 0, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 5, 0, 0, 0, 0, 0, 0, 0, 0], + [0, 0, 5, 0, 0, 0, 0, 0, 0, 0]]) + ) + + def test_operate_4d_array(self): + a = np.zeros((3, 3, 3, 3), int) + fill_diagonal(a, 4) + i = np.array([0, 1, 2]) + assert_equal(np.where(a != 0), (i, i, i, i)) def test_diag_indices(): di = diag_indices(4) a = np.array([[1, 2, 3, 4], - [5, 6, 7, 8], - [9, 10, 11, 12], - [13, 14, 15, 16]]) + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16]]) a[di] = 100 - yield (assert_array_equal, a, - np.array([[100, 2, 3, 4], - [5, 100, 7, 8], - [9, 10, 100, 12], - [13, 14, 15, 100]])) + assert_array_equal( + a, np.array([[100, 2, 3, 4], + [5, 100, 7, 8], + [9, 10, 100, 12], + [13, 14, 15, 100]]) + ) # Now, we create indices to manipulate a 3-d array: d3 = diag_indices(2, 3) @@ -357,12 +331,12 @@ def test_diag_indices(): # And use it to set the diagonal of a zeros array to 1: a = np.zeros((2, 2, 2), int) a[d3] = 1 - yield (assert_array_equal, a, - np.array([[[1, 0], - [0, 0]], - - [[0, 0], - [0, 1]]])) + assert_array_equal( + a, np.array([[[1, 0], + [0, 0]], + [[0, 0], + [0, 1]]]) + ) def test_diag_indices_from(): @@ -394,7 +368,3 @@ def test_ndindex(): # Make sure 0-sized ndindex works correctly x = list(ndindex(*[0])) assert_equal(x, []) - - -if __name__ == "__main__": - run_module_suite() |