diff options
Diffstat (limited to 'numpy/lib/tests/test_index_tricks.py')
| -rw-r--r-- | numpy/lib/tests/test_index_tricks.py | 67 |
1 files changed, 62 insertions, 5 deletions
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py index c17ee5d6a..40f75936e 100644 --- a/numpy/lib/tests/test_index_tricks.py +++ b/numpy/lib/tests/test_index_tricks.py @@ -4,12 +4,69 @@ from numpy import ( array, ones, r_, mgrid, unravel_index, zeros, where, ndenumerate, fill_diagonal, diag_indices, diag_indices_from, s_, index_exp ) -class TestUnravelIndex(TestCase): +class TestRavelUnravelIndex(TestCase): def test_basic(self): - assert unravel_index(2,(2,2)) == (1,0) - assert unravel_index(254,(17,94)) == (2, 66) - assert_raises(ValueError, unravel_index, 4,(2,2)) - + assert_equal(np.unravel_index(2,(2,2)), (1,0)) + assert_equal(np.ravel_coords((1,0),(2,2)), 2) + assert_equal(np.unravel_index(254,(17,94)), (2,66)) + assert_equal(np.ravel_coords((2,66),(17,94)), 254) + assert_raises(ValueError, np.unravel_index, -1, (2,2)) + assert_raises(TypeError, np.unravel_index, 0.5, (2,2)) + assert_raises(ValueError, np.unravel_index, 4, (2,2)) + assert_raises(ValueError, np.ravel_coords, (-3,1), (2,2)) + assert_raises(ValueError, np.ravel_coords, (2,1), (2,2)) + assert_raises(ValueError, np.ravel_coords, (0,-3), (2,2)) + assert_raises(ValueError, np.ravel_coords, (0,2), (2,2)) + assert_raises(TypeError, np.ravel_coords, (0.1,0.), (2,2)) + + assert_equal(np.unravel_index((2*3 + 1)*6 + 4, (4,3,6)), [2,1,4]) + assert_equal(np.ravel_coords([2,1,4], (4,3,6)), (2*3 + 1)*6 + 4) + + arr = np.array([[3,6,6],[4,5,1]]) + assert_equal(np.ravel_coords(arr, (7,6)), [22,41,37]) + assert_equal(np.ravel_coords(arr, (7,6), order='F'), [31,41,13]) + assert_equal(np.ravel_coords(arr, (4,6), mode='clip'), [22,23,19]) + assert_equal(np.ravel_coords(arr, (4,4), mode=('clip','wrap')), + [12,13,13]) + assert_equal(np.ravel_coords((3,1,4,1), (6,7,8,9)), 1621) + + assert_equal(np.unravel_index(np.array([22, 41, 37]), (7,6)), + [[3, 6, 6],[4, 5, 1]]) + assert_equal(np.unravel_index(np.array([31, 41, 13]), (7,6), order='F'), + [[3, 6, 6], [4, 5, 1]]) + assert_equal(np.unravel_index(1621, (6,7,8,9)), [3,1,4,1]) + + def test_dtypes(self): + # Test with different data types + for dtype in [np.int16, np.uint16, np.int32, + np.uint32, np.int64, np.uint64]: + coords = np.array([[1,0,1,2,3,4],[1,6,1,3,2,0]], dtype=dtype) + shape = (5,8) + uncoords = 8*coords[0]+coords[1] + assert_equal(np.ravel_coords(coords, shape), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape)) + uncoords = coords[0]+5*coords[1] + assert_equal(np.ravel_coords(coords, shape, order='F'), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) + + coords = np.array([[1,0,1,2,3,4],[1,6,1,3,2,0],[1,3,1,0,9,5]], + dtype=dtype) + shape = (5,8,10) + uncoords = 10*(8*coords[0]+coords[1])+coords[2] + assert_equal(np.ravel_coords(coords, shape), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape)) + uncoords = coords[0]+5*(coords[1]+8*coords[2]) + assert_equal(np.ravel_coords(coords, shape, order='F'), uncoords) + assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) + + def test_clipmodes(self): + # Test clipmodes + assert_equal(np.ravel_coords([5,1,-1,2], (4,3,7,12), mode='wrap'), + np.ravel_coords([1,1,6,2], (4,3,7,12))) + assert_equal(np.ravel_coords([5,1,-1,2], (4,3,7,12), + mode=('wrap','raise','clip','raise')), + np.ravel_coords([1,1,0,2], (4,3,7,12))) + assert_raises(ValueError, np.ravel_coords, [5,1,-1,2], (4,3,7,12)) class TestGrid(TestCase): def test_basic(self): |