diff options
Diffstat (limited to 'numpy/lib/tests')
| -rw-r--r-- | numpy/lib/tests/test_index_tricks.py | 46 | ||||
| -rw-r--r-- | numpy/lib/tests/test_twodim_base.py | 110 |
2 files changed, 137 insertions, 19 deletions
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py index 641737d43..50dc5ac15 100644 --- a/numpy/lib/tests/test_index_tricks.py +++ b/numpy/lib/tests/test_index_tricks.py @@ -1,5 +1,6 @@ from numpy.testing import * -from numpy import array, ones, r_, mgrid, unravel_index, ndenumerate +from numpy import ( array, ones, r_, mgrid, unravel_index, zeros, where, + fill_diagonal, diag_indices, diag_indices_from ) class TestUnravelIndex(TestCase): def test_basic(self): @@ -68,5 +69,48 @@ class TestNdenumerate(TestCase): assert_equal(list(ndenumerate(a)), [((0,0), 1), ((0,1), 2), ((1,0), 3), ((1,1), 4)]) + +def test_fill_diagonal(): + a = zeros((3, 3),int) + fill_diagonal(a, 5) + yield (assert_array_equal, a, + array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5]])) + + # The same function can operate on a 4-d array: + a = zeros((3, 3, 3, 3), int) + fill_diagonal(a, 4) + i = array([0, 1, 2]) + yield (assert_equal, where(a != 0), (i, i, i, i)) + + +def test_diag_indices(): + di = diag_indices(4) + a = array([[1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16]]) + a[di] = 100 + yield (assert_array_equal, a, + 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) + + # And use it to set the diagonal of a zeros array to 1: + a = zeros((2, 2, 2),int) + a[d3] = 1 + yield (assert_array_equal, a, + array([[[1, 0], + [0, 0]], + + [[0, 0], + [0, 1]]]) ) + + if __name__ == "__main__": run_module_suite() diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py index 32c4ca58e..a1e58d53b 100644 --- a/numpy/lib/tests/test_twodim_base.py +++ b/numpy/lib/tests/test_twodim_base.py @@ -3,8 +3,11 @@ """ from numpy.testing import * -from numpy import arange, rot90, add, fliplr, flipud, zeros, ones, eye, \ - array, diag, histogram2d, tri + +from numpy import ( arange, rot90, add, fliplr, flipud, zeros, ones, eye, + array, diag, histogram2d, tri, mask_indices, triu_indices, + triu_indices_from, tril_indices, tril_indices_from ) + import numpy as np def get_mat(n): @@ -52,32 +55,32 @@ class TestEye(TestCase): class TestDiag(TestCase): def test_vector(self): - vals = (100*arange(5)).astype('l') - b = zeros((5,5)) + vals = (100 * arange(5)).astype('l') + b = zeros((5, 5)) for k in range(5): - b[k,k] = vals[k] - assert_equal(diag(vals),b) - b = zeros((7,7)) + b[k, k] = vals[k] + assert_equal(diag(vals), b) + b = zeros((7, 7)) c = b.copy() for k in range(5): - b[k,k+2] = vals[k] - c[k+2,k] = vals[k] - assert_equal(diag(vals,k=2), b) - assert_equal(diag(vals,k=-2), c) + b[k, k + 2] = vals[k] + c[k + 2, k] = vals[k] + assert_equal(diag(vals, k=2), b) + assert_equal(diag(vals, k=-2), c) def test_matrix(self): - vals = (100*get_mat(5)+1).astype('l') + vals = (100 * get_mat(5) + 1).astype('l') b = zeros((5,)) for k in range(5): b[k] = vals[k,k] - assert_equal(diag(vals),b) - b = b*0 + assert_equal(diag(vals), b) + b = b * 0 for k in range(3): - b[k] = vals[k,k+2] - assert_equal(diag(vals,2),b[:3]) + b[k] = vals[k, k + 2] + assert_equal(diag(vals, 2), b[:3]) for k in range(3): - b[k] = vals[k+2,k] - assert_equal(diag(vals,-2),b[:3]) + b[k] = vals[k + 2, k] + assert_equal(diag(vals, -2), b[:3]) class TestFliplr(TestCase): def test_basic(self): @@ -193,5 +196,76 @@ class TestTri(TestCase): assert_array_equal(tri(3,dtype=bool),out.astype(bool)) +def test_mask_indices(): + # simple test without offset + iu = mask_indices(3, np.triu) + a = np.arange(9).reshape(3, 3) + yield (assert_array_equal, a[iu], array([0, 1, 2, 4, 5, 8])) + # Now with an offset + iu1 = mask_indices(3, np.triu, 1) + yield (assert_array_equal, a[iu1], array([1, 2, 5])) + + +def test_tril_indices(): + # indices without and with offset + il1 = tril_indices(4) + il2 = tril_indices(4, 2) + + a = np.array([[1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16]]) + + # indexing: + yield (assert_array_equal, a[il1], + array([ 1, 5, 6, 9, 10, 11, 13, 14, 15, 16]) ) + + # And for assigning values: + a[il1] = -1 + yield (assert_array_equal, a, + array([[-1, 2, 3, 4], + [-1, -1, 7, 8], + [-1, -1, -1, 12], + [-1, -1, -1, -1]]) ) + + # These cover almost the whole array (two diagonals right of the main one): + a[il2] = -10 + yield (assert_array_equal, a, + array([[-10, -10, -10, 4], + [-10, -10, -10, -10], + [-10, -10, -10, -10], + [-10, -10, -10, -10]]) ) + + +def test_triu_indices(): + iu1 = triu_indices(4) + iu2 = triu_indices(4, 2) + + a = np.array([[1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + [13, 14, 15, 16]]) + + # Both for indexing: + yield (assert_array_equal, a[iu1], + array([ 1, 5, 6, 9, 10, 11, 13, 14, 15, 16]) ) + + # And for assigning values: + a[iu1] = -1 + yield (assert_array_equal, a, + array([[-1, -1, -1, -1], + [ 5, -1, -1, -1], + [ 9, 10, -1, -1], + [13, 14, 15, -1]]) ) + + # These cover almost the whole array (two diagonals right of the main one): + a[iu2] = -10 + yield ( assert_array_equal, a, + array([[ -1, -1, -10, -10], + [ 5, -1, -1, -10], + [ 9, 10, -1, -1], + [ 13, 14, 15, -1]]) ) + + if __name__ == "__main__": run_module_suite() |