import sys import pytest import textwrap import subprocess import numpy as np import numpy.core._multiarray_tests as _multiarray_tests from numpy import array, arange, nditer, all from numpy.testing import ( assert_, assert_equal, assert_array_equal, assert_raises, IS_WASM, HAS_REFCOUNT, suppress_warnings, break_cycles ) def iter_multi_index(i): ret = [] while not i.finished: ret.append(i.multi_index) i.iternext() return ret def iter_indices(i): ret = [] while not i.finished: ret.append(i.index) i.iternext() return ret def iter_iterindices(i): ret = [] while not i.finished: ret.append(i.iterindex) i.iternext() return ret @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") def test_iter_refcount(): # Make sure the iterator doesn't leak # Basic a = arange(6) dt = np.dtype('f4').newbyteorder() rc_a = sys.getrefcount(a) rc_dt = sys.getrefcount(dt) with nditer(a, [], [['readwrite', 'updateifcopy']], casting='unsafe', op_dtypes=[dt]) as it: assert_(not it.iterationneedsapi) assert_(sys.getrefcount(a) > rc_a) assert_(sys.getrefcount(dt) > rc_dt) # del 'it' it = None assert_equal(sys.getrefcount(a), rc_a) assert_equal(sys.getrefcount(dt), rc_dt) # With a copy a = arange(6, dtype='f4') dt = np.dtype('f4') rc_a = sys.getrefcount(a) rc_dt = sys.getrefcount(dt) it = nditer(a, [], [['readwrite']], op_dtypes=[dt]) rc2_a = sys.getrefcount(a) rc2_dt = sys.getrefcount(dt) it2 = it.copy() assert_(sys.getrefcount(a) > rc2_a) assert_(sys.getrefcount(dt) > rc2_dt) it = None assert_equal(sys.getrefcount(a), rc2_a) assert_equal(sys.getrefcount(dt), rc2_dt) it2 = None assert_equal(sys.getrefcount(a), rc_a) assert_equal(sys.getrefcount(dt), rc_dt) del it2 # avoid pyflakes unused variable warning def test_iter_best_order(): # The iterator should always find the iteration order # with increasing memory addresses # Test the ordering for 1-D to 5-D shapes for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: a = arange(np.prod(shape)) # Test each combination of positive and negative strides for dirs in range(2**len(shape)): dirs_index = [slice(None)]*len(shape) for bit in range(len(shape)): if ((2**bit) & dirs): dirs_index[bit] = slice(None, None, -1) dirs_index = tuple(dirs_index) aview = a.reshape(shape)[dirs_index] # C-order i = nditer(aview, [], [['readonly']]) assert_equal([x for x in i], a) # Fortran-order i = nditer(aview.T, [], [['readonly']]) assert_equal([x for x in i], a) # Other order if len(shape) > 2: i = nditer(aview.swapaxes(0, 1), [], [['readonly']]) assert_equal([x for x in i], a) def test_iter_c_order(): # Test forcing C order # Test the ordering for 1-D to 5-D shapes for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: a = arange(np.prod(shape)) # Test each combination of positive and negative strides for dirs in range(2**len(shape)): dirs_index = [slice(None)]*len(shape) for bit in range(len(shape)): if ((2**bit) & dirs): dirs_index[bit] = slice(None, None, -1) dirs_index = tuple(dirs_index) aview = a.reshape(shape)[dirs_index] # C-order i = nditer(aview, order='C') assert_equal([x for x in i], aview.ravel(order='C')) # Fortran-order i = nditer(aview.T, order='C') assert_equal([x for x in i], aview.T.ravel(order='C')) # Other order if len(shape) > 2: i = nditer(aview.swapaxes(0, 1), order='C') assert_equal([x for x in i], aview.swapaxes(0, 1).ravel(order='C')) def test_iter_f_order(): # Test forcing F order # Test the ordering for 1-D to 5-D shapes for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: a = arange(np.prod(shape)) # Test each combination of positive and negative strides for dirs in range(2**len(shape)): dirs_index = [slice(None)]*len(shape) for bit in range(len(shape)): if ((2**bit) & dirs): dirs_index[bit] = slice(None, None, -1) dirs_index = tuple(dirs_index) aview = a.reshape(shape)[dirs_index] # C-order i = nditer(aview, order='F') assert_equal([x for x in i], aview.ravel(order='F')) # Fortran-order i = nditer(aview.T, order='F') assert_equal([x for x in i], aview.T.ravel(order='F')) # Other order if len(shape) > 2: i = nditer(aview.swapaxes(0, 1), order='F') assert_equal([x for x in i], aview.swapaxes(0, 1).ravel(order='F')) def test_iter_c_or_f_order(): # Test forcing any contiguous (C or F) order # Test the ordering for 1-D to 5-D shapes for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: a = arange(np.prod(shape)) # Test each combination of positive and negative strides for dirs in range(2**len(shape)): dirs_index = [slice(None)]*len(shape) for bit in range(len(shape)): if ((2**bit) & dirs): dirs_index[bit] = slice(None, None, -1) dirs_index = tuple(dirs_index) aview = a.reshape(shape)[dirs_index] # C-order i = nditer(aview, order='A') assert_equal([x for x in i], aview.ravel(order='A')) # Fortran-order i = nditer(aview.T, order='A') assert_equal([x for x in i], aview.T.ravel(order='A')) # Other order if len(shape) > 2: i = nditer(aview.swapaxes(0, 1), order='A') assert_equal([x for x in i], aview.swapaxes(0, 1).ravel(order='A')) def test_nditer_multi_index_set(): # Test the multi_index set a = np.arange(6).reshape(2, 3) it = np.nditer(a, flags=['multi_index']) # Removes the iteration on two first elements of a[0] it.multi_index = (0, 2,) assert_equal([i for i in it], [2, 3, 4, 5]) @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts") def test_nditer_multi_index_set_refcount(): # Test if the reference count on index variable is decreased index = 0 i = np.nditer(np.array([111, 222, 333, 444]), flags=['multi_index']) start_count = sys.getrefcount(index) i.multi_index = (index,) end_count = sys.getrefcount(index) assert_equal(start_count, end_count) def test_iter_best_order_multi_index_1d(): # The multi-indices should be correct with any reordering a = arange(4) # 1D order i = nditer(a, ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0,), (1,), (2,), (3,)]) # 1D reversed order i = nditer(a[::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(3,), (2,), (1,), (0,)]) def test_iter_best_order_multi_index_2d(): # The multi-indices should be correct with any reordering a = arange(6) # 2D C-order i = nditer(a.reshape(2, 3), ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]) # 2D Fortran-order i = nditer(a.reshape(2, 3).copy(order='F'), ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2)]) # 2D reversed C-order i = nditer(a.reshape(2, 3)[::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(1, 0), (1, 1), (1, 2), (0, 0), (0, 1), (0, 2)]) i = nditer(a.reshape(2, 3)[:, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 2), (0, 1), (0, 0), (1, 2), (1, 1), (1, 0)]) i = nditer(a.reshape(2, 3)[::-1, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(1, 2), (1, 1), (1, 0), (0, 2), (0, 1), (0, 0)]) # 2D reversed Fortran-order i = nditer(a.reshape(2, 3).copy(order='F')[::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(1, 0), (0, 0), (1, 1), (0, 1), (1, 2), (0, 2)]) i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 2), (1, 2), (0, 1), (1, 1), (0, 0), (1, 0)]) i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(1, 2), (0, 2), (1, 1), (0, 1), (1, 0), (0, 0)]) def test_iter_best_order_multi_index_3d(): # The multi-indices should be correct with any reordering a = arange(12) # 3D C-order i = nditer(a.reshape(2, 3, 2), ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1)]) # 3D Fortran-order i = nditer(a.reshape(2, 3, 2).copy(order='F'), ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0), (0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1)]) # 3D reversed C-order i = nditer(a.reshape(2, 3, 2)[::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1), (0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1)]) i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 2, 0), (0, 2, 1), (0, 1, 0), (0, 1, 1), (0, 0, 0), (0, 0, 1), (1, 2, 0), (1, 2, 1), (1, 1, 0), (1, 1, 1), (1, 0, 0), (1, 0, 1)]) i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 0, 1), (0, 0, 0), (0, 1, 1), (0, 1, 0), (0, 2, 1), (0, 2, 0), (1, 0, 1), (1, 0, 0), (1, 1, 1), (1, 1, 0), (1, 2, 1), (1, 2, 0)]) # 3D reversed Fortran-order i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(1, 0, 0), (0, 0, 0), (1, 1, 0), (0, 1, 0), (1, 2, 0), (0, 2, 0), (1, 0, 1), (0, 0, 1), (1, 1, 1), (0, 1, 1), (1, 2, 1), (0, 2, 1)]) i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 2, 0), (1, 2, 0), (0, 1, 0), (1, 1, 0), (0, 0, 0), (1, 0, 0), (0, 2, 1), (1, 2, 1), (0, 1, 1), (1, 1, 1), (0, 0, 1), (1, 0, 1)]) i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1], ['multi_index'], [['readonly']]) assert_equal(iter_multi_index(i), [(0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1), (0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0)]) def test_iter_best_order_c_index_1d(): # The C index should be correct with any reordering a = arange(4) # 1D order i = nditer(a, ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 1, 2, 3]) # 1D reversed order i = nditer(a[::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [3, 2, 1, 0]) def test_iter_best_order_c_index_2d(): # The C index should be correct with any reordering a = arange(6) # 2D C-order i = nditer(a.reshape(2, 3), ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5]) # 2D Fortran-order i = nditer(a.reshape(2, 3).copy(order='F'), ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 3, 1, 4, 2, 5]) # 2D reversed C-order i = nditer(a.reshape(2, 3)[::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [3, 4, 5, 0, 1, 2]) i = nditer(a.reshape(2, 3)[:, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [2, 1, 0, 5, 4, 3]) i = nditer(a.reshape(2, 3)[::-1, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0]) # 2D reversed Fortran-order i = nditer(a.reshape(2, 3).copy(order='F')[::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [3, 0, 4, 1, 5, 2]) i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [2, 5, 1, 4, 0, 3]) i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [5, 2, 4, 1, 3, 0]) def test_iter_best_order_c_index_3d(): # The C index should be correct with any reordering a = arange(12) # 3D C-order i = nditer(a.reshape(2, 3, 2), ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) # 3D Fortran-order i = nditer(a.reshape(2, 3, 2).copy(order='F'), ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11]) # 3D reversed C-order i = nditer(a.reshape(2, 3, 2)[::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5]) i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7]) i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10]) # 3D reversed Fortran-order i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5]) i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7]) i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1], ['c_index'], [['readonly']]) assert_equal(iter_indices(i), [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10]) def test_iter_best_order_f_index_1d(): # The Fortran index should be correct with any reordering a = arange(4) # 1D order i = nditer(a, ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 1, 2, 3]) # 1D reversed order i = nditer(a[::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [3, 2, 1, 0]) def test_iter_best_order_f_index_2d(): # The Fortran index should be correct with any reordering a = arange(6) # 2D C-order i = nditer(a.reshape(2, 3), ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 2, 4, 1, 3, 5]) # 2D Fortran-order i = nditer(a.reshape(2, 3).copy(order='F'), ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5]) # 2D reversed C-order i = nditer(a.reshape(2, 3)[::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [1, 3, 5, 0, 2, 4]) i = nditer(a.reshape(2, 3)[:, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [4, 2, 0, 5, 3, 1]) i = nditer(a.reshape(2, 3)[::-1, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [5, 3, 1, 4, 2, 0]) # 2D reversed Fortran-order i = nditer(a.reshape(2, 3).copy(order='F')[::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [1, 0, 3, 2, 5, 4]) i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [4, 5, 2, 3, 0, 1]) i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0]) def test_iter_best_order_f_index_3d(): # The Fortran index should be correct with any reordering a = arange(12) # 3D C-order i = nditer(a.reshape(2, 3, 2), ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11]) # 3D Fortran-order i = nditer(a.reshape(2, 3, 2).copy(order='F'), ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) # 3D reversed C-order i = nditer(a.reshape(2, 3, 2)[::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10]) i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7]) i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5]) # 3D reversed Fortran-order i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10]) i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7]) i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1], ['f_index'], [['readonly']]) assert_equal(iter_indices(i), [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5]) def test_iter_no_inner_full_coalesce(): # Check no_inner iterators which coalesce into a single inner loop for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]: size = np.prod(shape) a = arange(size) # Test each combination of forward and backwards indexing for dirs in range(2**len(shape)): dirs_index = [slice(None)]*len(shape) for bit in range(len(shape)): if ((2**bit) & dirs): dirs_index[bit] = slice(None, None, -1) dirs_index = tuple(dirs_index) aview = a.reshape(shape)[dirs_index] # C-order i = nditer(aview, ['external_loop'], [['readonly']]) assert_equal(i.ndim, 1) assert_equal(i[0].shape, (size,)) # Fortran-order i = nditer(aview.T, ['external_loop'], [['readonly']]) assert_equal(i.ndim, 1) assert_equal(i[0].shape, (size,)) # Other order if len(shape) > 2: i = nditer(aview.swapaxes(0, 1), ['external_loop'], [['readonly']]) assert_equal(i.ndim, 1) assert_equal(i[0].shape, (size,)) def test_iter_no_inner_dim_coalescing(): # Check no_inner iterators whose dimensions may not coalesce completely # Skipping the last element in a dimension prevents coalescing # with the next-bigger dimension a = arange(24).reshape(2, 3, 4)[:,:, :-1] i = nditer(a, ['external_loop'], [['readonly']]) assert_equal(i.ndim, 2) assert_equal(i[0].shape, (3,)) a = arange(24).reshape(2, 3, 4)[:, :-1,:] i = nditer(a, ['external_loop'], [['readonly']]) assert_equal(i.ndim, 2) assert_equal(i[0].shape, (8,)) a = arange(24).reshape(2, 3, 4)[:-1,:,:] i = nditer(a, ['external_loop'], [['readonly']]) assert_equal(i.ndim, 1) assert_equal(i[0].shape, (12,)) # Even with lots of 1-sized dimensions, should still coalesce a = arange(24).reshape(1, 1, 2, 1, 1, 3, 1, 1, 4, 1, 1) i = nditer(a, ['external_loop'], [['readonly']]) assert_equal(i.ndim, 1) assert_equal(i[0].shape, (24,)) def test_iter_dim_coalescing(): # Check that the correct number of dimensions are coalesced # Tracking a multi-index disables coalescing a = arange(24).reshape(2, 3, 4) i = nditer(a, ['multi_index'], [['readonly']]) assert_equal(i.ndim, 3) # A tracked index can allow coalescing if it's compatible with the array a3d = arange(24).reshape(2, 3, 4) i = nditer(a3d, ['c_index'], [['readonly']]) assert_equal(i.ndim, 1) i = nditer(a3d.swapaxes(0, 1), ['c_index'], [['readonly']]) assert_equal(i.ndim, 3) i = nditer(a3d.T, ['c_index'], [['readonly']]) assert_equal(i.ndim, 3) i = nditer(a3d.T, ['f_index'], [['readonly']]) assert_equal(i.ndim, 1) i = nditer(a3d.T.swapaxes(0, 1), ['f_index'], [['readonly']]) assert_equal(i.ndim, 3) # When C or F order is forced, coalescing may still occur a3d = arange(24).reshape(2, 3, 4) i = nditer(a3d, order='C') assert_equal(i.ndim, 1) i = nditer(a3d.T, order='C') assert_equal(i.ndim, 3) i = nditer(a3d, order='F') assert_equal(i.ndim, 3) i = nditer(a3d.T, order='F') assert_equal(i.ndim, 1) i = nditer(a3d, order='A') assert_equal(i.ndim, 1) i = nditer(a3d.T, order='A') assert_equal(i.ndim, 1) def test_iter_broadcasting(): # Standard NumPy broadcasting rules # 1D with scalar i = nditer([arange(6), np.int32(2)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 6) assert_equal(i.shape, (6,)) # 2D with scalar i = nditer([arange(6).reshape(2, 3), np.int32(2)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 6) assert_equal(i.shape, (2, 3)) # 2D with 1D i = nditer([arange(6).reshape(2, 3), arange(3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 6) assert_equal(i.shape, (2, 3)) i = nditer([arange(2).reshape(2, 1), arange(3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 6) assert_equal(i.shape, (2, 3)) # 2D with 2D i = nditer([arange(2).reshape(2, 1), arange(3).reshape(1, 3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 6) assert_equal(i.shape, (2, 3)) # 3D with scalar i = nditer([np.int32(2), arange(24).reshape(4, 2, 3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) # 3D with 1D i = nditer([arange(3), arange(24).reshape(4, 2, 3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) i = nditer([arange(3), arange(8).reshape(4, 2, 1)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) # 3D with 2D i = nditer([arange(6).reshape(2, 3), arange(24).reshape(4, 2, 3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) i = nditer([arange(2).reshape(2, 1), arange(24).reshape(4, 2, 3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) i = nditer([arange(3).reshape(1, 3), arange(8).reshape(4, 2, 1)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) # 3D with 3D i = nditer([arange(2).reshape(1, 2, 1), arange(3).reshape(1, 1, 3), arange(4).reshape(4, 1, 1)], ['multi_index'], [['readonly']]*3) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) i = nditer([arange(6).reshape(1, 2, 3), arange(4).reshape(4, 1, 1)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) i = nditer([arange(24).reshape(4, 2, 3), arange(12).reshape(4, 1, 3)], ['multi_index'], [['readonly']]*2) assert_equal(i.itersize, 24) assert_equal(i.shape, (4, 2, 3)) def test_iter_itershape(): # Check that allocated outputs work with a specified shape a = np.arange(6, dtype='i2').reshape(2, 3) i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], op_axes=[[0, 1, None], None], itershape=(-1, -1, 4)) assert_equal(i.operands[1].shape, (2, 3, 4)) assert_equal(i.operands[1].strides, (24, 8, 2)) i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']], op_axes=[[0, 1, None], None], itershape=(-1, -1, 4)) assert_equal(i.operands[1].shape, (3, 2, 4)) assert_equal(i.operands[1].strides, (8, 24, 2)) i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']], order='F', op_axes=[[0, 1, None], None], itershape=(-1, -1, 4)) assert_equal(i.operands[1].shape, (3, 2, 4)) assert_equal(i.operands[1].strides, (2, 6, 12)) # If we specify 1 in the itershape, it shouldn't allow broadcasting # of that dimension to a bigger value assert_raises(ValueError, nditer, [a, None], [], [['readonly'], ['writeonly', 'allocate']], op_axes=[[0, 1, None], None], itershape=(-1, 1, 4)) # Test bug that for no op_axes but itershape, they are NULLed correctly i = np.nditer([np.ones(2), None, None], itershape=(2,)) def test_iter_broadcasting_errors(): # Check that errors are thrown for bad broadcasting shapes # 1D with 1D assert_raises(ValueError, nditer, [arange(2), arange(3)], [], [['readonly']]*2) # 2D with 1D assert_raises(ValueError, nditer, [arange(6).reshape(2, 3), arange(2)], [], [['readonly']]*2) # 2D with 2D assert_raises(ValueError, nditer, [arange(6).reshape(2, 3), arange(9).reshape(3, 3)], [], [['readonly']]*2) assert_raises(ValueError, nditer, [arange(6).reshape(2, 3), arange(4).reshape(2, 2)], [], [['readonly']]*2) # 3D with 3D assert_raises(ValueError, nditer, [arange(36).reshape(3, 3, 4), arange(24).reshape(2, 3, 4)], [], [['readonly']]*2) assert_raises(ValueError, nditer, [arange(8).reshape(2, 4, 1), arange(24).reshape(2, 3, 4)], [], [['readonly']]*2) # Verify that the error message mentions the right shapes try: nditer([arange(2).reshape(1, 2, 1), arange(3).reshape(1, 3), arange(6).reshape(2, 3)], [], [['readonly'], ['readonly'], ['writeonly', 'no_broadcast']]) raise AssertionError('Should have raised a broadcast error') except ValueError as e: msg = str(e) # The message should contain the shape of the 3rd operand assert_(msg.find('(2,3)') >= 0, 'Message "%s" doesn\'t contain operand shape (2,3)' % msg) # The message should contain the broadcast shape assert_(msg.find('(1,2,3)') >= 0, 'Message "%s" doesn\'t contain broadcast shape (1,2,3)' % msg) try: nditer([arange(6).reshape(2, 3), arange(2)], [], [['readonly'], ['readonly']], op_axes=[[0, 1], [0, np.newaxis]], itershape=(4, 3)) raise AssertionError('Should have raised a broadcast error') except ValueError as e: msg = str(e) # The message should contain "shape->remappedshape" for each operand assert_(msg.find('(2,3)->(2,3)') >= 0, 'Message "%s" doesn\'t contain operand shape (2,3)->(2,3)' % msg) assert_(msg.find('(2,)->(2,newaxis)') >= 0, ('Message "%s" doesn\'t contain remapped operand shape' + '(2,)->(2,newaxis)') % msg) # The message should contain the itershape parameter assert_(msg.find('(4,3)') >= 0, 'Message "%s" doesn\'t contain itershape parameter (4,3)' % msg) try: nditer([np.zeros((2, 1, 1)), np.zeros((2,))], [], [['writeonly', 'no_broadcast'], ['readonly']]) raise AssertionError('Should have raised a broadcast error') except ValueError as e: msg = str(e) # The message should contain the shape of the bad operand assert_(msg.find('(2,1,1)') >= 0, 'Message "%s" doesn\'t contain operand shape (2,1,1)' % msg) # The message should contain the broadcast shape assert_(msg.find('(2,1,2)') >= 0, 'Message "%s" doesn\'t contain the broadcast shape (2,1,2)' % msg) def test_iter_flags_errors(): # Check that bad combinations of flags produce errors a = arange(6) # Not enough operands assert_raises(ValueError, nditer, [], [], []) # Too many operands assert_raises(ValueError, nditer, [a]*100, [], [['readonly']]*100) # Bad global flag assert_raises(ValueError, nditer, [a], ['bad flag'], [['readonly']]) # Bad op flag assert_raises(ValueError, nditer, [a], [], [['readonly', 'bad flag']]) # Bad order parameter assert_raises(ValueError, nditer, [a], [], [['readonly']], order='G') # Bad casting parameter assert_raises(ValueError, nditer, [a], [], [['readonly']], casting='noon') # op_flags must match ops assert_raises(ValueError, nditer, [a]*3, [], [['readonly']]*2) # Cannot track both a C and an F index assert_raises(ValueError, nditer, a, ['c_index', 'f_index'], [['readonly']]) # Inner iteration and multi-indices/indices are incompatible assert_raises(ValueError, nditer, a, ['external_loop', 'multi_index'], [['readonly']]) assert_raises(ValueError, nditer, a, ['external_loop', 'c_index'], [['readonly']]) assert_raises(ValueError, nditer, a, ['external_loop', 'f_index'], [['readonly']]) # Must specify exactly one of readwrite/readonly/writeonly per operand assert_raises(ValueError, nditer, a, [], [[]]) assert_raises(ValueError, nditer, a, [], [['readonly', 'writeonly']]) assert_raises(ValueError, nditer, a, [], [['readonly', 'readwrite']]) assert_raises(ValueError, nditer, a, [], [['writeonly', 'readwrite']]) assert_raises(ValueError, nditer, a, [], [['readonly', 'writeonly', 'readwrite']]) # Python scalars are always readonly assert_raises(TypeError, nditer, 1.5, [], [['writeonly']]) assert_raises(TypeError, nditer, 1.5, [], [['readwrite']]) # Array scalars are always readonly assert_raises(TypeError, nditer, np.int32(1), [], [['writeonly']]) assert_raises(TypeError, nditer, np.int32(1), [], [['readwrite']]) # Check readonly array a.flags.writeable = False assert_raises(ValueError, nditer, a, [], [['writeonly']]) assert_raises(ValueError, nditer, a, [], [['readwrite']]) a.flags.writeable = True # Multi-indices available only with the multi_index flag i = nditer(arange(6), [], [['readonly']]) assert_raises(ValueError, lambda i:i.multi_index, i) # Index available only with an index flag assert_raises(ValueError, lambda i:i.index, i) # GotoCoords and GotoIndex incompatible with buffering or no_inner def assign_multi_index(i): i.multi_index = (0,) def assign_index(i): i.index = 0 def assign_iterindex(i): i.iterindex = 0 def assign_iterrange(i): i.iterrange = (0, 1) i = nditer(arange(6), ['external_loop']) assert_raises(ValueError, assign_multi_index, i) assert_raises(ValueError, assign_index, i) assert_raises(ValueError, assign_iterindex, i) assert_raises(ValueError, assign_iterrange, i) i = nditer(arange(6), ['buffered']) assert_raises(ValueError, assign_multi_index, i) assert_raises(ValueError, assign_index, i) assert_raises(ValueError, assign_iterrange, i) # Can't iterate if size is zero assert_raises(ValueError, nditer, np.array([])) def test_iter_slice(): a, b, c = np.arange(3), np.arange(3), np.arange(3.) i = nditer([a, b, c], [], ['readwrite']) with i: i[0:2] = (3, 3) assert_equal(a, [3, 1, 2]) assert_equal(b, [3, 1, 2]) assert_equal(c, [0, 1, 2]) i[1] = 12 assert_equal(i[0:2], [3, 12]) def test_iter_assign_mapping(): a = np.arange(24, dtype='f8').reshape(2, 3, 4).T it = np.nditer(a, [], [['readwrite', 'updateifcopy']], casting='same_kind', op_dtypes=[np.dtype('f4')]) with it: it.operands[0][...] = 3 it.operands[0][...] = 14 assert_equal(a, 14) it = np.nditer(a, [], [['readwrite', 'updateifcopy']], casting='same_kind', op_dtypes=[np.dtype('f4')]) with it: x = it.operands[0][-1:1] x[...] = 14 it.operands[0][...] = -1234 assert_equal(a, -1234) # check for no warnings on dealloc x = None it = None def test_iter_nbo_align_contig(): # Check that byte order, alignment, and contig changes work # Byte order change by requesting a specific dtype a = np.arange(6, dtype='f4') au = a.byteswap().newbyteorder() assert_(a.dtype.byteorder != au.dtype.byteorder) i = nditer(au, [], [['readwrite', 'updateifcopy']], casting='equiv', op_dtypes=[np.dtype('f4')]) with i: # context manager triggers WRITEBACKIFCOPY on i at exit assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder) assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder) assert_equal(i.operands[0], a) i.operands[0][:] = 2 assert_equal(au, [2]*6) del i # should not raise a warning # Byte order change by requesting NBO a = np.arange(6, dtype='f4') au = a.byteswap().newbyteorder() assert_(a.dtype.byteorder != au.dtype.byteorder) with nditer(au, [], [['readwrite', 'updateifcopy', 'nbo']], casting='equiv') as i: # context manager triggers UPDATEIFCOPY on i at exit assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder) assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder) assert_equal(i.operands[0], a) i.operands[0][:] = 12345 i.operands[0][:] = 2 assert_equal(au, [2]*6) # Unaligned input a = np.zeros((6*4+1,), dtype='i1')[1:] a.dtype = 'f4' a[:] = np.arange(6, dtype='f4') assert_(not a.flags.aligned) # Without 'aligned', shouldn't copy i = nditer(a, [], [['readonly']]) assert_(not i.operands[0].flags.aligned) assert_equal(i.operands[0], a) # With 'aligned', should make a copy with nditer(a, [], [['readwrite', 'updateifcopy', 'aligned']]) as i: assert_(i.operands[0].flags.aligned) # context manager triggers UPDATEIFCOPY on i at exit assert_equal(i.operands[0], a) i.operands[0][:] = 3 assert_equal(a, [3]*6) # Discontiguous input a = arange(12) # If it is contiguous, shouldn't copy i = nditer(a[:6], [], [['readonly']]) assert_(i.operands[0].flags.contiguous) assert_equal(i.operands[0], a[:6]) # If it isn't contiguous, should buffer i = nditer(a[::2], ['buffered', 'external_loop'], [['readonly', 'contig']], buffersize=10) assert_(i[0].flags.contiguous) assert_equal(i[0], a[::2]) def test_iter_array_cast(): # Check that arrays are cast as requested # No cast 'f4' -> 'f4' a = np.arange(6, dtype='f4').reshape(2, 3) i = nditer(a, [], [['readwrite']], op_dtypes=[np.dtype('f4')]) with i: assert_equal(i.operands[0], a) assert_equal(i.operands[0].dtype, np.dtype('f4')) # Byte-order cast ' '>f4' a = np.arange(6, dtype='f4')]) as i: assert_equal(i.operands[0], a) assert_equal(i.operands[0].dtype, np.dtype('>f4')) # Safe case 'f4' -> 'f8' a = np.arange(24, dtype='f4').reshape(2, 3, 4).swapaxes(1, 2) i = nditer(a, [], [['readonly', 'copy']], casting='safe', op_dtypes=[np.dtype('f8')]) assert_equal(i.operands[0], a) assert_equal(i.operands[0].dtype, np.dtype('f8')) # The memory layout of the temporary should match a (a is (48,4,16)) # except negative strides get flipped to positive strides. assert_equal(i.operands[0].strides, (96, 8, 32)) a = a[::-1,:, ::-1] i = nditer(a, [], [['readonly', 'copy']], casting='safe', op_dtypes=[np.dtype('f8')]) assert_equal(i.operands[0], a) assert_equal(i.operands[0].dtype, np.dtype('f8')) assert_equal(i.operands[0].strides, (96, 8, 32)) # Same-kind cast 'f8' -> 'f4' -> 'f8' a = np.arange(24, dtype='f8').reshape(2, 3, 4).T with nditer(a, [], [['readwrite', 'updateifcopy']], casting='same_kind', op_dtypes=[np.dtype('f4')]) as i: assert_equal(i.operands[0], a) assert_equal(i.operands[0].dtype, np.dtype('f4')) assert_equal(i.operands[0].strides, (4, 16, 48)) # Check that WRITEBACKIFCOPY is activated at exit i.operands[0][2, 1, 1] = -12.5 assert_(a[2, 1, 1] != -12.5) assert_equal(a[2, 1, 1], -12.5) a = np.arange(6, dtype='i4')[::-2] with nditer(a, [], [['writeonly', 'updateifcopy']], casting='unsafe', op_dtypes=[np.dtype('f4')]) as i: assert_equal(i.operands[0].dtype, np.dtype('f4')) # Even though the stride was negative in 'a', it # becomes positive in the temporary assert_equal(i.operands[0].strides, (4,)) i.operands[0][:] = [1, 2, 3] assert_equal(a, [1, 2, 3]) def test_iter_array_cast_errors(): # Check that invalid casts are caught # Need to enable copying for casts to occur assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], [['readonly']], op_dtypes=[np.dtype('f8')]) # Also need to allow casting for casts to occur assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], [['readonly', 'copy']], casting='no', op_dtypes=[np.dtype('f8')]) assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], [['readonly', 'copy']], casting='equiv', op_dtypes=[np.dtype('f8')]) assert_raises(TypeError, nditer, arange(2, dtype='f8'), [], [['writeonly', 'updateifcopy']], casting='no', op_dtypes=[np.dtype('f4')]) assert_raises(TypeError, nditer, arange(2, dtype='f8'), [], [['writeonly', 'updateifcopy']], casting='equiv', op_dtypes=[np.dtype('f4')]) # ' '>f4' should not work with casting='no' assert_raises(TypeError, nditer, arange(2, dtype='f4')]) # 'f4' -> 'f8' is a safe cast, but 'f8' -> 'f4' isn't assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], [['readwrite', 'updateifcopy']], casting='safe', op_dtypes=[np.dtype('f8')]) assert_raises(TypeError, nditer, arange(2, dtype='f8'), [], [['readwrite', 'updateifcopy']], casting='safe', op_dtypes=[np.dtype('f4')]) # 'f4' -> 'i4' is neither a safe nor a same-kind cast assert_raises(TypeError, nditer, arange(2, dtype='f4'), [], [['readonly', 'copy']], casting='same_kind', op_dtypes=[np.dtype('i4')]) assert_raises(TypeError, nditer, arange(2, dtype='i4'), [], [['writeonly', 'updateifcopy']], casting='same_kind', op_dtypes=[np.dtype('f4')]) def test_iter_scalar_cast(): # Check that scalars are cast as requested # No cast 'f4' -> 'f4' i = nditer(np.float32(2.5), [], [['readonly']], op_dtypes=[np.dtype('f4')]) assert_equal(i.dtypes[0], np.dtype('f4')) assert_equal(i.value.dtype, np.dtype('f4')) assert_equal(i.value, 2.5) # Safe cast 'f4' -> 'f8' i = nditer(np.float32(2.5), [], [['readonly', 'copy']], casting='safe', op_dtypes=[np.dtype('f8')]) assert_equal(i.dtypes[0], np.dtype('f8')) assert_equal(i.value.dtype, np.dtype('f8')) assert_equal(i.value, 2.5) # Same-kind cast 'f8' -> 'f4' i = nditer(np.float64(2.5), [], [['readonly', 'copy']], casting='same_kind', op_dtypes=[np.dtype('f4')]) assert_equal(i.dtypes[0], np.dtype('f4')) assert_equal(i.value.dtype, np.dtype('f4')) assert_equal(i.value, 2.5) # Unsafe cast 'f8' -> 'i4' i = nditer(np.float64(3.0), [], [['readonly', 'copy']], casting='unsafe', op_dtypes=[np.dtype('i4')]) assert_equal(i.dtypes[0], np.dtype('i4')) assert_equal(i.value.dtype, np.dtype('i4')) assert_equal(i.value, 3) # Readonly scalars may be cast even without setting COPY or BUFFERED i = nditer(3, [], [['readonly']], op_dtypes=[np.dtype('f8')]) assert_equal(i[0].dtype, np.dtype('f8')) assert_equal(i[0], 3.) def test_iter_scalar_cast_errors(): # Check that invalid casts are caught # Need to allow copying/buffering for write casts of scalars to occur assert_raises(TypeError, nditer, np.float32(2), [], [['readwrite']], op_dtypes=[np.dtype('f8')]) assert_raises(TypeError, nditer, 2.5, [], [['readwrite']], op_dtypes=[np.dtype('f4')]) # 'f8' -> 'f4' isn't a safe cast if the value would overflow assert_raises(TypeError, nditer, np.float64(1e60), [], [['readonly']], casting='safe', op_dtypes=[np.dtype('f4')]) # 'f4' -> 'i4' is neither a safe nor a same-kind cast assert_raises(TypeError, nditer, np.float32(2), [], [['readonly']], casting='same_kind', op_dtypes=[np.dtype('i4')]) def test_iter_object_arrays_basic(): # Check that object arrays work obj = {'a':3,'b':'d'} a = np.array([[1, 2, 3], None, obj, None], dtype='O') if HAS_REFCOUNT: rc = sys.getrefcount(obj) # Need to allow references for object arrays assert_raises(TypeError, nditer, a) if HAS_REFCOUNT: assert_equal(sys.getrefcount(obj), rc) i = nditer(a, ['refs_ok'], ['readonly']) vals = [x_[()] for x_ in i] assert_equal(np.array(vals, dtype='O'), a) vals, i, x = [None]*3 if HAS_REFCOUNT: assert_equal(sys.getrefcount(obj), rc) i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'], ['readonly'], order='C') assert_(i.iterationneedsapi) vals = [x_[()] for x_ in i] assert_equal(np.array(vals, dtype='O'), a.reshape(2, 2).ravel(order='F')) vals, i, x = [None]*3 if HAS_REFCOUNT: assert_equal(sys.getrefcount(obj), rc) i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'], ['readwrite'], order='C') with i: for x in i: x[...] = None vals, i, x = [None]*3 if HAS_REFCOUNT: assert_(sys.getrefcount(obj) == rc-1) assert_equal(a, np.array([None]*4, dtype='O')) def test_iter_object_arrays_conversions(): # Conversions to/from objects a = np.arange(6, dtype='O') i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], casting='unsafe', op_dtypes='i4') with i: for x in i: x[...] += 1 assert_equal(a, np.arange(6)+1) a = np.arange(6, dtype='i4') i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], casting='unsafe', op_dtypes='O') with i: for x in i: x[...] += 1 assert_equal(a, np.arange(6)+1) # Non-contiguous object array a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'O')]) a = a['a'] a[:] = np.arange(6) i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], casting='unsafe', op_dtypes='i4') with i: for x in i: x[...] += 1 assert_equal(a, np.arange(6)+1) #Non-contiguous value array a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'i4')]) a = a['a'] a[:] = np.arange(6) + 98172488 i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'], casting='unsafe', op_dtypes='O') with i: ob = i[0][()] if HAS_REFCOUNT: rc = sys.getrefcount(ob) for x in i: x[...] += 1 if HAS_REFCOUNT: assert_(sys.getrefcount(ob) == rc-1) assert_equal(a, np.arange(6)+98172489) def test_iter_common_dtype(): # Check that the iterator finds a common data type correctly i = nditer([array([3], dtype='f4'), array([0], dtype='f8')], ['common_dtype'], [['readonly', 'copy']]*2, casting='safe') assert_equal(i.dtypes[0], np.dtype('f8')) assert_equal(i.dtypes[1], np.dtype('f8')) i = nditer([array([3], dtype='i4'), array([0], dtype='f4')], ['common_dtype'], [['readonly', 'copy']]*2, casting='safe') assert_equal(i.dtypes[0], np.dtype('f8')) assert_equal(i.dtypes[1], np.dtype('f8')) i = nditer([array([3], dtype='f4'), array(0, dtype='f8')], ['common_dtype'], [['readonly', 'copy']]*2, casting='same_kind') assert_equal(i.dtypes[0], np.dtype('f4')) assert_equal(i.dtypes[1], np.dtype('f4')) i = nditer([array([3], dtype='u4'), array(0, dtype='i4')], ['common_dtype'], [['readonly', 'copy']]*2, casting='safe') assert_equal(i.dtypes[0], np.dtype('u4')) assert_equal(i.dtypes[1], np.dtype('u4')) i = nditer([array([3], dtype='u4'), array(-12, dtype='i4')], ['common_dtype'], [['readonly', 'copy']]*2, casting='safe') assert_equal(i.dtypes[0], np.dtype('i8')) assert_equal(i.dtypes[1], np.dtype('i8')) i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'), array([2j], dtype='c8'), array([9], dtype='f8')], ['common_dtype'], [['readonly', 'copy']]*4, casting='safe') assert_equal(i.dtypes[0], np.dtype('c16')) assert_equal(i.dtypes[1], np.dtype('c16')) assert_equal(i.dtypes[2], np.dtype('c16')) assert_equal(i.dtypes[3], np.dtype('c16')) assert_equal(i.value, (3, -12, 2j, 9)) # When allocating outputs, other outputs aren't factored in i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')], [], [['readonly', 'copy'], ['writeonly', 'allocate'], ['writeonly']], casting='safe') assert_equal(i.dtypes[0], np.dtype('i4')) assert_equal(i.dtypes[1], np.dtype('i4')) assert_equal(i.dtypes[2], np.dtype('c16')) # But, if common data types are requested, they are i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')], ['common_dtype'], [['readonly', 'copy'], ['writeonly', 'allocate'], ['writeonly']], casting='safe') assert_equal(i.dtypes[0], np.dtype('c16')) assert_equal(i.dtypes[1], np.dtype('c16')) assert_equal(i.dtypes[2], np.dtype('c16')) def test_iter_copy_if_overlap(): # Ensure the iterator makes copies on read/write overlap, if requested # Copy not needed, 1 op for flag in ['readonly', 'writeonly', 'readwrite']: a = arange(10) i = nditer([a], ['copy_if_overlap'], [[flag]]) with i: assert_(i.operands[0] is a) # Copy needed, 2 ops, read-write overlap x = arange(10) a = x[1:] b = x[:-1] with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i: assert_(not np.shares_memory(*i.operands)) # Copy not needed with elementwise, 2 ops, exactly same arrays x = arange(10) a = x b = x i = nditer([a, b], ['copy_if_overlap'], [['readonly', 'overlap_assume_elementwise'], ['readwrite', 'overlap_assume_elementwise']]) with i: assert_(i.operands[0] is a and i.operands[1] is b) with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i: assert_(i.operands[0] is a and not np.shares_memory(i.operands[1], b)) # Copy not needed, 2 ops, no overlap x = arange(10) a = x[::2] b = x[1::2] i = nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']]) assert_(i.operands[0] is a and i.operands[1] is b) # Copy needed, 2 ops, read-write overlap x = arange(4, dtype=np.int8) a = x[3:] b = x.view(np.int32)[:1] with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']]) as i: assert_(not np.shares_memory(*i.operands)) # Copy needed, 3 ops, read-write overlap for flag in ['writeonly', 'readwrite']: x = np.ones([10, 10]) a = x b = x.T c = x with nditer([a, b, c], ['copy_if_overlap'], [['readonly'], ['readonly'], [flag]]) as i: a2, b2, c2 = i.operands assert_(not np.shares_memory(a2, c2)) assert_(not np.shares_memory(b2, c2)) # Copy not needed, 3 ops, read-only overlap x = np.ones([10, 10]) a = x b = x.T c = x i = nditer([a, b, c], ['copy_if_overlap'], [['readonly'], ['readonly'], ['readonly']]) a2, b2, c2 = i.operands assert_(a is a2) assert_(b is b2) assert_(c is c2) # Copy not needed, 3 ops, read-only overlap x = np.ones([10, 10]) a = x b = np.ones([10, 10]) c = x.T i = nditer([a, b, c], ['copy_if_overlap'], [['readonly'], ['writeonly'], ['readonly']]) a2, b2, c2 = i.operands assert_(a is a2) assert_(b is b2) assert_(c is c2) # Copy not needed, 3 ops, write-only overlap x = np.arange(7) a = x[:3] b = x[3:6] c = x[4:7] i = nditer([a, b, c], ['copy_if_overlap'], [['readonly'], ['writeonly'], ['writeonly']]) a2, b2, c2 = i.operands assert_(a is a2) assert_(b is b2) assert_(c is c2) def test_iter_op_axes(): # Check that custom axes work # Reverse the axes a = arange(6).reshape(2, 3) i = nditer([a, a.T], [], [['readonly']]*2, op_axes=[[0, 1], [1, 0]]) assert_(all([x == y for (x, y) in i])) a = arange(24).reshape(2, 3, 4) i = nditer([a.T, a], [], [['readonly']]*2, op_axes=[[2, 1, 0], None]) assert_(all([x == y for (x, y) in i])) # Broadcast 1D to any dimension a = arange(1, 31).reshape(2, 3, 5) b = arange(1, 3) i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [0, -1, -1]]) assert_equal([x*y for (x, y) in i], (a*b.reshape(2, 1, 1)).ravel()) b = arange(1, 4) i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [-1, 0, -1]]) assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 3, 1)).ravel()) b = arange(1, 6) i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [np.newaxis, np.newaxis, 0]]) assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 1, 5)).ravel()) # Inner product-style broadcasting a = arange(24).reshape(2, 3, 4) b = arange(40).reshape(5, 2, 4) i = nditer([a, b], ['multi_index'], [['readonly']]*2, op_axes=[[0, 1, -1, -1], [-1, -1, 0, 1]]) assert_equal(i.shape, (2, 3, 5, 2)) # Matrix product-style broadcasting a = arange(12).reshape(3, 4) b = arange(20).reshape(4, 5) i = nditer([a, b], ['multi_index'], [['readonly']]*2, op_axes=[[0, -1], [-1, 1]]) assert_equal(i.shape, (3, 5)) def test_iter_op_axes_errors(): # Check that custom axes throws errors for bad inputs # Wrong number of items in op_axes a = arange(6).reshape(2, 3) assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, op_axes=[[0], [1], [0]]) # Out of bounds items in op_axes assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, op_axes=[[2, 1], [0, 1]]) assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, op_axes=[[0, 1], [2, -1]]) # Duplicate items in op_axes assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, op_axes=[[0, 0], [0, 1]]) assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, op_axes=[[0, 1], [1, 1]]) # Different sized arrays in op_axes assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, op_axes=[[0, 1], [0, 1, 0]]) # Non-broadcastable dimensions in the result assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2, op_axes=[[0, 1], [1, 0]]) def test_iter_copy(): # Check that copying the iterator works correctly a = arange(24).reshape(2, 3, 4) # Simple iterator i = nditer(a) j = i.copy() assert_equal([x[()] for x in i], [x[()] for x in j]) i.iterindex = 3 j = i.copy() assert_equal([x[()] for x in i], [x[()] for x in j]) # Buffered iterator i = nditer(a, ['buffered', 'ranged'], order='F', buffersize=3) j = i.copy() assert_equal([x[()] for x in i], [x[()] for x in j]) i.iterindex = 3 j = i.copy() assert_equal([x[()] for x in i], [x[()] for x in j]) i.iterrange = (3, 9) j = i.copy() assert_equal([x[()] for x in i], [x[()] for x in j]) i.iterrange = (2, 18) next(i) next(i) j = i.copy() assert_equal([x[()] for x in i], [x[()] for x in j]) # Casting iterator with nditer(a, ['buffered'], order='F', casting='unsafe', op_dtypes='f8', buffersize=5) as i: j = i.copy() assert_equal([x[()] for x in j], a.ravel(order='F')) a = arange(24, dtype=' unstructured (any to object), and many other # casts, which cause this to require all steps in the casting machinery # one level down as well as the iterator copy (which uses NpyAuxData clone) in_dtype = np.dtype([("a", np.dtype("i,")), ("b", np.dtype(">i,d,S17,>d,(3)f,O,i1"))]) out_dtype = np.dtype([("a", np.dtype("O")), ("b", np.dtype(">i,>i,S17,>d,>U3,(3)d,i1,O"))]) arr = np.ones(1000, dtype=in_dtype) it = np.nditer((arr,), ["buffered", "external_loop", "refs_ok"], op_dtypes=[out_dtype], casting="unsafe") it_copy = it.copy() res1 = next(it) del it res2 = next(it_copy) del it_copy expected = arr["a"].astype(out_dtype["a"]) assert_array_equal(res1["a"], expected) assert_array_equal(res2["a"], expected) for field in in_dtype["b"].names: # Note that the .base avoids the subarray field expected = arr["b"][field].astype(out_dtype["b"][field].base) assert_array_equal(res1["b"][field], expected) assert_array_equal(res2["b"][field], expected) def test_iter_copy_casts_structured2(): # Similar to the above, this is a fairly arcane test to cover internals in_dtype = np.dtype([("a", np.dtype("O,O")), ("b", np.dtype("(5)O,(3)O,(1,)O,(1,)i,(1,)O"))]) out_dtype = np.dtype([("a", np.dtype("O")), ("b", np.dtype("O,(3)i,(4)O,(4)O,(4)i"))]) arr = np.ones(1, dtype=in_dtype) it = np.nditer((arr,), ["buffered", "external_loop", "refs_ok"], op_dtypes=[out_dtype], casting="unsafe") it_copy = it.copy() res1 = next(it) del it res2 = next(it_copy) del it_copy # Array of two structured scalars: for res in res1, res2: # Cast to tuple by getitem, which may be weird and changable?: assert type(res["a"][0]) == tuple assert res["a"][0] == (1, 1) for res in res1, res2: assert_array_equal(res["b"]["f0"][0], np.ones(5, dtype=object)) assert_array_equal(res["b"]["f1"], np.ones((1, 3), dtype="i")) assert res["b"]["f2"].shape == (1, 4) assert_array_equal(res["b"]["f2"][0], np.ones(4, dtype=object)) assert_array_equal(res["b"]["f3"][0], np.ones(4, dtype=object)) assert_array_equal(res["b"]["f3"][0], np.ones(4, dtype="i")) def test_iter_allocate_output_simple(): # Check that the iterator will properly allocate outputs # Simple case a = arange(6) i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], op_dtypes=[None, np.dtype('f4')]) assert_equal(i.operands[1].shape, a.shape) assert_equal(i.operands[1].dtype, np.dtype('f4')) def test_iter_allocate_output_buffered_readwrite(): # Allocated output with buffering + delay_bufalloc a = arange(6) i = nditer([a, None], ['buffered', 'delay_bufalloc'], [['readonly'], ['allocate', 'readwrite']]) with i: i.operands[1][:] = 1 i.reset() for x in i: x[1][...] += x[0][...] assert_equal(i.operands[1], a+1) def test_iter_allocate_output_itorder(): # The allocated output should match the iteration order # C-order input, best iteration order a = arange(6, dtype='i4').reshape(2, 3) i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], op_dtypes=[None, np.dtype('f4')]) assert_equal(i.operands[1].shape, a.shape) assert_equal(i.operands[1].strides, a.strides) assert_equal(i.operands[1].dtype, np.dtype('f4')) # F-order input, best iteration order a = arange(24, dtype='i4').reshape(2, 3, 4).T i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], op_dtypes=[None, np.dtype('f4')]) assert_equal(i.operands[1].shape, a.shape) assert_equal(i.operands[1].strides, a.strides) assert_equal(i.operands[1].dtype, np.dtype('f4')) # Non-contiguous input, C iteration order a = arange(24, dtype='i4').reshape(2, 3, 4).swapaxes(0, 1) i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']], order='C', op_dtypes=[None, np.dtype('f4')]) assert_equal(i.operands[1].shape, a.shape) assert_equal(i.operands[1].strides, (32, 16, 4)) assert_equal(i.operands[1].dtype, np.dtype('f4')) def test_iter_allocate_output_opaxes(): # Specifying op_axes should work a = arange(24, dtype='i4').reshape(2, 3, 4) i = nditer([None, a], [], [['writeonly', 'allocate'], ['readonly']], op_dtypes=[np.dtype('u4'), None], op_axes=[[1, 2, 0], None]) assert_equal(i.operands[0].shape, (4, 2, 3)) assert_equal(i.operands[0].strides, (4, 48, 16)) assert_equal(i.operands[0].dtype, np.dtype('u4')) def test_iter_allocate_output_types_promotion(): # Check type promotion of automatic outputs i = nditer([array([3], dtype='f4'), array([0], dtype='f8'), None], [], [['readonly']]*2+[['writeonly', 'allocate']]) assert_equal(i.dtypes[2], np.dtype('f8')) i = nditer([array([3], dtype='i4'), array([0], dtype='f4'), None], [], [['readonly']]*2+[['writeonly', 'allocate']]) assert_equal(i.dtypes[2], np.dtype('f8')) i = nditer([array([3], dtype='f4'), array(0, dtype='f8'), None], [], [['readonly']]*2+[['writeonly', 'allocate']]) assert_equal(i.dtypes[2], np.dtype('f4')) i = nditer([array([3], dtype='u4'), array(0, dtype='i4'), None], [], [['readonly']]*2+[['writeonly', 'allocate']]) assert_equal(i.dtypes[2], np.dtype('u4')) i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'), None], [], [['readonly']]*2+[['writeonly', 'allocate']]) assert_equal(i.dtypes[2], np.dtype('i8')) def test_iter_allocate_output_types_byte_order(): # Verify the rules for byte order changes # When there's just one input, the output type exactly matches a = array([3], dtype='u4').newbyteorder() i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']]) assert_equal(i.dtypes[0], i.dtypes[1]) # With two or more inputs, the output type is in native byte order i = nditer([a, a, None], [], [['readonly'], ['readonly'], ['writeonly', 'allocate']]) assert_(i.dtypes[0] != i.dtypes[2]) assert_equal(i.dtypes[0].newbyteorder('='), i.dtypes[2]) def test_iter_allocate_output_types_scalar(): # If the inputs are all scalars, the output should be a scalar i = nditer([None, 1, 2.3, np.float32(12), np.complex128(3)], [], [['writeonly', 'allocate']] + [['readonly']]*4) assert_equal(i.operands[0].dtype, np.dtype('complex128')) assert_equal(i.operands[0].ndim, 0) def test_iter_allocate_output_subtype(): # Make sure that the subtype with priority wins class MyNDArray(np.ndarray): __array_priority__ = 15 # subclass vs ndarray a = np.array([[1, 2], [3, 4]]).view(MyNDArray) b = np.arange(4).reshape(2, 2).T i = nditer([a, b, None], [], [['readonly'], ['readonly'], ['writeonly', 'allocate']]) assert_equal(type(a), type(i.operands[2])) assert_(type(b) is not type(i.operands[2])) assert_equal(i.operands[2].shape, (2, 2)) # If subtypes are disabled, we should get back an ndarray. i = nditer([a, b, None], [], [['readonly'], ['readonly'], ['writeonly', 'allocate', 'no_subtype']]) assert_equal(type(b), type(i.operands[2])) assert_(type(a) is not type(i.operands[2])) assert_equal(i.operands[2].shape, (2, 2)) def test_iter_allocate_output_errors(): # Check that the iterator will throw errors for bad output allocations # Need an input if no output data type is specified a = arange(6) assert_raises(TypeError, nditer, [a, None], [], [['writeonly'], ['writeonly', 'allocate']]) # Allocated output should be flagged for writing assert_raises(ValueError, nditer, [a, None], [], [['readonly'], ['allocate', 'readonly']]) # Allocated output can't have buffering without delayed bufalloc assert_raises(ValueError, nditer, [a, None], ['buffered'], ['allocate', 'readwrite']) # Must specify dtype if there are no inputs (cannot promote existing ones; # maybe this should use the 'f4' here, but it does not historically.) assert_raises(TypeError, nditer, [None, None], [], [['writeonly', 'allocate'], ['writeonly', 'allocate']], op_dtypes=[None, np.dtype('f4')]) # If using op_axes, must specify all the axes a = arange(24, dtype='i4').reshape(2, 3, 4) assert_raises(ValueError, nditer, [a, None], [], [['readonly'], ['writeonly', 'allocate']], op_dtypes=[None, np.dtype('f4')], op_axes=[None, [0, np.newaxis, 1]]) # If using op_axes, the axes must be within bounds assert_raises(ValueError, nditer, [a, None], [], [['readonly'], ['writeonly', 'allocate']], op_dtypes=[None, np.dtype('f4')], op_axes=[None, [0, 3, 1]]) # If using op_axes, there can't be duplicates assert_raises(ValueError, nditer, [a, None], [], [['readonly'], ['writeonly', 'allocate']], op_dtypes=[None, np.dtype('f4')], op_axes=[None, [0, 2, 1, 0]]) # Not all axes may be specified if a reduction. If there is a hole # in op_axes, this is an error. a = arange(24, dtype='i4').reshape(2, 3, 4) assert_raises(ValueError, nditer, [a, None], ["reduce_ok"], [['readonly'], ['readwrite', 'allocate']], op_dtypes=[None, np.dtype('f4')], op_axes=[None, [0, np.newaxis, 2]]) def test_all_allocated(): # When no output and no shape is given, `()` is used as shape. i = np.nditer([None], op_dtypes=["int64"]) assert i.operands[0].shape == () assert i.dtypes == (np.dtype("int64"),) i = np.nditer([None], op_dtypes=["int64"], itershape=(2, 3, 4)) assert i.operands[0].shape == (2, 3, 4) def test_iter_remove_axis(): a = arange(24).reshape(2, 3, 4) i = nditer(a, ['multi_index']) i.remove_axis(1) assert_equal([x for x in i], a[:, 0,:].ravel()) a = a[::-1,:,:] i = nditer(a, ['multi_index']) i.remove_axis(0) assert_equal([x for x in i], a[0,:,:].ravel()) def test_iter_remove_multi_index_inner_loop(): # Check that removing multi-index support works a = arange(24).reshape(2, 3, 4) i = nditer(a, ['multi_index']) assert_equal(i.ndim, 3) assert_equal(i.shape, (2, 3, 4)) assert_equal(i.itviews[0].shape, (2, 3, 4)) # Removing the multi-index tracking causes all dimensions to coalesce before = [x for x in i] i.remove_multi_index() after = [x for x in i] assert_equal(before, after) assert_equal(i.ndim, 1) assert_raises(ValueError, lambda i:i.shape, i) assert_equal(i.itviews[0].shape, (24,)) # Removing the inner loop means there's just one iteration i.reset() assert_equal(i.itersize, 24) assert_equal(i[0].shape, tuple()) i.enable_external_loop() assert_equal(i.itersize, 24) assert_equal(i[0].shape, (24,)) assert_equal(i.value, arange(24)) def test_iter_iterindex(): # Make sure iterindex works buffersize = 5 a = arange(24).reshape(4, 3, 2) for flags in ([], ['buffered']): i = nditer(a, flags, buffersize=buffersize) assert_equal(iter_iterindices(i), list(range(24))) i.iterindex = 2 assert_equal(iter_iterindices(i), list(range(2, 24))) i = nditer(a, flags, order='F', buffersize=buffersize) assert_equal(iter_iterindices(i), list(range(24))) i.iterindex = 5 assert_equal(iter_iterindices(i), list(range(5, 24))) i = nditer(a[::-1], flags, order='F', buffersize=buffersize) assert_equal(iter_iterindices(i), list(range(24))) i.iterindex = 9 assert_equal(iter_iterindices(i), list(range(9, 24))) i = nditer(a[::-1, ::-1], flags, order='C', buffersize=buffersize) assert_equal(iter_iterindices(i), list(range(24))) i.iterindex = 13 assert_equal(iter_iterindices(i), list(range(13, 24))) i = nditer(a[::1, ::-1], flags, buffersize=buffersize) assert_equal(iter_iterindices(i), list(range(24))) i.iterindex = 23 assert_equal(iter_iterindices(i), list(range(23, 24))) i.reset() i.iterindex = 2 assert_equal(iter_iterindices(i), list(range(2, 24))) def test_iter_iterrange(): # Make sure getting and resetting the iterrange works buffersize = 5 a = arange(24, dtype='i4').reshape(4, 3, 2) a_fort = a.ravel(order='F') i = nditer(a, ['ranged'], ['readonly'], order='F', buffersize=buffersize) assert_equal(i.iterrange, (0, 24)) assert_equal([x[()] for x in i], a_fort) for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]: i.iterrange = r assert_equal(i.iterrange, r) assert_equal([x[()] for x in i], a_fort[r[0]:r[1]]) i = nditer(a, ['ranged', 'buffered'], ['readonly'], order='F', op_dtypes='f8', buffersize=buffersize) assert_equal(i.iterrange, (0, 24)) assert_equal([x[()] for x in i], a_fort) for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]: i.iterrange = r assert_equal(i.iterrange, r) assert_equal([x[()] for x in i], a_fort[r[0]:r[1]]) def get_array(i): val = np.array([], dtype='f8') for x in i: val = np.concatenate((val, x)) return val i = nditer(a, ['ranged', 'buffered', 'external_loop'], ['readonly'], order='F', op_dtypes='f8', buffersize=buffersize) assert_equal(i.iterrange, (0, 24)) assert_equal(get_array(i), a_fort) for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]: i.iterrange = r assert_equal(i.iterrange, r) assert_equal(get_array(i), a_fort[r[0]:r[1]]) def test_iter_buffering(): # Test buffering with several buffer sizes and types arrays = [] # F-order swapped array arrays.append(np.arange(24, dtype='c16').reshape(2, 3, 4).T.newbyteorder().byteswap()) # Contiguous 1-dimensional array arrays.append(np.arange(10, dtype='f4')) # Unaligned array a = np.zeros((4*16+1,), dtype='i1')[1:] a.dtype = 'i4' a[:] = np.arange(16, dtype='i4') arrays.append(a) # 4-D F-order array arrays.append(np.arange(120, dtype='i4').reshape(5, 3, 2, 4).T) for a in arrays: for buffersize in (1, 2, 3, 5, 8, 11, 16, 1024): vals = [] i = nditer(a, ['buffered', 'external_loop'], [['readonly', 'nbo', 'aligned']], order='C', casting='equiv', buffersize=buffersize) while not i.finished: assert_(i[0].size <= buffersize) vals.append(i[0].copy()) i.iternext() assert_equal(np.concatenate(vals), a.ravel(order='C')) def test_iter_write_buffering(): # Test that buffering of writes is working # F-order swapped array a = np.arange(24).reshape(2, 3, 4).T.newbyteorder().byteswap() i = nditer(a, ['buffered'], [['readwrite', 'nbo', 'aligned']], casting='equiv', order='C', buffersize=16) x = 0 with i: while not i.finished: i[0] = x x += 1 i.iternext() assert_equal(a.ravel(order='C'), np.arange(24)) def test_iter_buffering_delayed_alloc(): # Test that delaying buffer allocation works a = np.arange(6) b = np.arange(1, dtype='f4') i = nditer([a, b], ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok'], ['readwrite'], casting='unsafe', op_dtypes='f4') assert_(i.has_delayed_bufalloc) assert_raises(ValueError, lambda i:i.multi_index, i) assert_raises(ValueError, lambda i:i[0], i) assert_raises(ValueError, lambda i:i[0:2], i) def assign_iter(i): i[0] = 0 assert_raises(ValueError, assign_iter, i) i.reset() assert_(not i.has_delayed_bufalloc) assert_equal(i.multi_index, (0,)) with i: assert_equal(i[0], 0) i[1] = 1 assert_equal(i[0:2], [0, 1]) assert_equal([[x[0][()], x[1][()]] for x in i], list(zip(range(6), [1]*6))) def test_iter_buffered_cast_simple(): # Test that buffering can handle a simple cast a = np.arange(10, dtype='f4') i = nditer(a, ['buffered', 'external_loop'], [['readwrite', 'nbo', 'aligned']], casting='same_kind', op_dtypes=[np.dtype('f8')], buffersize=3) with i: for v in i: v[...] *= 2 assert_equal(a, 2*np.arange(10, dtype='f4')) def test_iter_buffered_cast_byteswapped(): # Test that buffering can handle a cast which requires swap->cast->swap a = np.arange(10, dtype='f4').newbyteorder().byteswap() i = nditer(a, ['buffered', 'external_loop'], [['readwrite', 'nbo', 'aligned']], casting='same_kind', op_dtypes=[np.dtype('f8').newbyteorder()], buffersize=3) with i: for v in i: v[...] *= 2 assert_equal(a, 2*np.arange(10, dtype='f4')) with suppress_warnings() as sup: sup.filter(np.ComplexWarning) a = np.arange(10, dtype='f8').newbyteorder().byteswap() i = nditer(a, ['buffered', 'external_loop'], [['readwrite', 'nbo', 'aligned']], casting='unsafe', op_dtypes=[np.dtype('c8').newbyteorder()], buffersize=3) with i: for v in i: v[...] *= 2 assert_equal(a, 2*np.arange(10, dtype='f8')) def test_iter_buffered_cast_byteswapped_complex(): # Test that buffering can handle a cast which requires swap->cast->copy a = np.arange(10, dtype='c8').newbyteorder().byteswap() a += 2j i = nditer(a, ['buffered', 'external_loop'], [['readwrite', 'nbo', 'aligned']], casting='same_kind', op_dtypes=[np.dtype('c16')], buffersize=3) with i: for v in i: v[...] *= 2 assert_equal(a, 2*np.arange(10, dtype='c8') + 4j) a = np.arange(10, dtype='c8') a += 2j i = nditer(a, ['buffered', 'external_loop'], [['readwrite', 'nbo', 'aligned']], casting='same_kind', op_dtypes=[np.dtype('c16').newbyteorder()], buffersize=3) with i: for v in i: v[...] *= 2 assert_equal(a, 2*np.arange(10, dtype='c8') + 4j) a = np.arange(10, dtype=np.clongdouble).newbyteorder().byteswap() a += 2j i = nditer(a, ['buffered', 'external_loop'], [['readwrite', 'nbo', 'aligned']], casting='same_kind', op_dtypes=[np.dtype('c16')], buffersize=3) with i: for v in i: v[...] *= 2 assert_equal(a, 2*np.arange(10, dtype=np.clongdouble) + 4j) a = np.arange(10, dtype=np.longdouble).newbyteorder().byteswap() i = nditer(a, ['buffered', 'external_loop'], [['readwrite', 'nbo', 'aligned']], casting='same_kind', op_dtypes=[np.dtype('f4')], buffersize=7) with i: for v in i: v[...] *= 2 assert_equal(a, 2*np.arange(10, dtype=np.longdouble)) def test_iter_buffered_cast_structured_type(): # Tests buffering of structured types # simple -> struct type (duplicates the value) sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')] a = np.arange(3, dtype='f4') + 0.5 i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt) vals = [np.array(x) for x in i] assert_equal(vals[0]['a'], 0.5) assert_equal(vals[0]['b'], 0) assert_equal(vals[0]['c'], [[(0.5)]*3]*2) assert_equal(vals[0]['d'], 0.5) assert_equal(vals[1]['a'], 1.5) assert_equal(vals[1]['b'], 1) assert_equal(vals[1]['c'], [[(1.5)]*3]*2) assert_equal(vals[1]['d'], 1.5) assert_equal(vals[0].dtype, np.dtype(sdt)) # object -> struct type sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')] a = np.zeros((3,), dtype='O') a[0] = (0.5, 0.5, [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]], 0.5) a[1] = (1.5, 1.5, [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]], 1.5) a[2] = (2.5, 2.5, [[2.5, 2.5, 2.5], [2.5, 2.5, 2.5]], 2.5) if HAS_REFCOUNT: rc = sys.getrefcount(a[0]) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt) vals = [x.copy() for x in i] assert_equal(vals[0]['a'], 0.5) assert_equal(vals[0]['b'], 0) assert_equal(vals[0]['c'], [[(0.5)]*3]*2) assert_equal(vals[0]['d'], 0.5) assert_equal(vals[1]['a'], 1.5) assert_equal(vals[1]['b'], 1) assert_equal(vals[1]['c'], [[(1.5)]*3]*2) assert_equal(vals[1]['d'], 1.5) assert_equal(vals[0].dtype, np.dtype(sdt)) vals, i, x = [None]*3 if HAS_REFCOUNT: assert_equal(sys.getrefcount(a[0]), rc) # single-field struct type -> simple sdt = [('a', 'f4')] a = np.array([(5.5,), (8,)], dtype=sdt) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes='i4') assert_equal([x_[()] for x_ in i], [5, 8]) # make sure multi-field struct type -> simple doesn't work sdt = [('a', 'f4'), ('b', 'i8'), ('d', 'O')] a = np.array([(5.5, 7, 'test'), (8, 10, 11)], dtype=sdt) assert_raises(TypeError, lambda: ( nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes='i4'))) # struct type -> struct type (field-wise copy) sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')] sdt2 = [('d', 'u2'), ('a', 'O'), ('b', 'f8')] a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) assert_equal([np.array(x_) for x_ in i], [np.array((1, 2, 3), dtype=sdt2), np.array((4, 5, 6), dtype=sdt2)]) def test_iter_buffered_cast_structured_type_failure_with_cleanup(): # make sure struct type -> struct type with different # number of fields fails sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')] sdt2 = [('b', 'O'), ('a', 'f8')] a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1) for intent in ["readwrite", "readonly", "writeonly"]: # This test was initially designed to test an error at a different # place, but will now raise earlier to to the cast not being possible: # `assert np.can_cast(a.dtype, sdt2, casting="unsafe")` fails. # Without a faulty DType, there is probably no reliable # way to get the initial tested behaviour. simple_arr = np.array([1, 2], dtype="i,i") # requires clean up with pytest.raises(TypeError): nditer((simple_arr, a), ['buffered', 'refs_ok'], [intent, intent], casting='unsafe', op_dtypes=["f,f", sdt2]) def test_buffered_cast_error_paths(): with pytest.raises(ValueError): # The input is cast into an `S3` buffer np.nditer((np.array("a", dtype="S1"),), op_dtypes=["i"], casting="unsafe", flags=["buffered"]) # The `M8[ns]` is cast into the `S3` output it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"], op_flags=["writeonly"], casting="unsafe", flags=["buffered"]) with pytest.raises(ValueError): with it: buf = next(it) buf[...] = "a" # cannot be converted to int. @pytest.mark.skipif(IS_WASM, reason="Cannot start subprocess") @pytest.mark.skipif(not HAS_REFCOUNT, reason="PyPy seems to not hit this.") def test_buffered_cast_error_paths_unraisable(): # The following gives an unraisable error. Pytest sometimes captures that # (depending python and/or pytest version). So with Python>=3.8 this can # probably be cleaned out in the future to check for # pytest.PytestUnraisableExceptionWarning: code = textwrap.dedent(""" import numpy as np it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"], op_flags=["writeonly"], casting="unsafe", flags=["buffered"]) buf = next(it) buf[...] = "a" del buf, it # Flushing only happens during deallocate right now. """) res = subprocess.check_output([sys.executable, "-c", code], stderr=subprocess.STDOUT, text=True) assert "ValueError" in res def test_iter_buffered_cast_subarray(): # Tests buffering of subarrays # one element -> many (copies it to all) sdt1 = [('a', 'f4')] sdt2 = [('a', 'f8', (3, 2, 2))] a = np.zeros((6,), dtype=sdt1) a['a'] = np.arange(6) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) for x, count in zip(i, list(range(6))): assert_(np.all(x['a'] == count)) # one element -> many -> back (copies it to all) sdt1 = [('a', 'O', (1, 1))] sdt2 = [('a', 'O', (3, 2, 2))] a = np.zeros((6,), dtype=sdt1) a['a'][:, 0, 0] = np.arange(6) i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'], casting='unsafe', op_dtypes=sdt2) with i: assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_(np.all(x['a'] == count)) x['a'][0] += 2 count += 1 assert_equal(a['a'], np.arange(6).reshape(6, 1, 1)+2) # many -> one element -> back (copies just element 0) sdt1 = [('a', 'O', (3, 2, 2))] sdt2 = [('a', 'O', (1,))] a = np.zeros((6,), dtype=sdt1) a['a'][:, 0, 0, 0] = np.arange(6) i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'], casting='unsafe', op_dtypes=sdt2) with i: assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'], count) x['a'] += 2 count += 1 assert_equal(a['a'], np.arange(6).reshape(6, 1, 1, 1)*np.ones((1, 3, 2, 2))+2) # many -> one element -> back (copies just element 0) sdt1 = [('a', 'f8', (3, 2, 2))] sdt2 = [('a', 'O', (1,))] a = np.zeros((6,), dtype=sdt1) a['a'][:, 0, 0, 0] = np.arange(6) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'], count) count += 1 # many -> one element (copies just element 0) sdt1 = [('a', 'O', (3, 2, 2))] sdt2 = [('a', 'f4', (1,))] a = np.zeros((6,), dtype=sdt1) a['a'][:, 0, 0, 0] = np.arange(6) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'], count) count += 1 # many -> matching shape (straightforward copy) sdt1 = [('a', 'O', (3, 2, 2))] sdt2 = [('a', 'f4', (3, 2, 2))] a = np.zeros((6,), dtype=sdt1) a['a'] = np.arange(6*3*2*2).reshape(6, 3, 2, 2) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'], a[count]['a']) count += 1 # vector -> smaller vector (truncates) sdt1 = [('a', 'f8', (6,))] sdt2 = [('a', 'f4', (2,))] a = np.zeros((6,), dtype=sdt1) a['a'] = np.arange(6*6).reshape(6, 6) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'], a[count]['a'][:2]) count += 1 # vector -> bigger vector (pads with zeros) sdt1 = [('a', 'f8', (2,))] sdt2 = [('a', 'f4', (6,))] a = np.zeros((6,), dtype=sdt1) a['a'] = np.arange(6*2).reshape(6, 2) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'][:2], a[count]['a']) assert_equal(x['a'][2:], [0, 0, 0, 0]) count += 1 # vector -> matrix (broadcasts) sdt1 = [('a', 'f8', (2,))] sdt2 = [('a', 'f4', (2, 2))] a = np.zeros((6,), dtype=sdt1) a['a'] = np.arange(6*2).reshape(6, 2) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'][0], a[count]['a']) assert_equal(x['a'][1], a[count]['a']) count += 1 # vector -> matrix (broadcasts and zero-pads) sdt1 = [('a', 'f8', (2, 1))] sdt2 = [('a', 'f4', (3, 2))] a = np.zeros((6,), dtype=sdt1) a['a'] = np.arange(6*2).reshape(6, 2, 1) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'][:2, 0], a[count]['a'][:, 0]) assert_equal(x['a'][:2, 1], a[count]['a'][:, 0]) assert_equal(x['a'][2,:], [0, 0]) count += 1 # matrix -> matrix (truncates and zero-pads) sdt1 = [('a', 'f8', (2, 3))] sdt2 = [('a', 'f4', (3, 2))] a = np.zeros((6,), dtype=sdt1) a['a'] = np.arange(6*2*3).reshape(6, 2, 3) i = nditer(a, ['buffered', 'refs_ok'], ['readonly'], casting='unsafe', op_dtypes=sdt2) assert_equal(i[0].dtype, np.dtype(sdt2)) count = 0 for x in i: assert_equal(x['a'][:2, 0], a[count]['a'][:, 0]) assert_equal(x['a'][:2, 1], a[count]['a'][:, 1]) assert_equal(x['a'][2,:], [0, 0]) count += 1 def test_iter_buffering_badwriteback(): # Writing back from a buffer cannot combine elements # a needs write buffering, but had a broadcast dimension a = np.arange(6).reshape(2, 3, 1) b = np.arange(12).reshape(2, 3, 2) assert_raises(ValueError, nditer, [a, b], ['buffered', 'external_loop'], [['readwrite'], ['writeonly']], order='C') # But if a is readonly, it's fine nditer([a, b], ['buffered', 'external_loop'], [['readonly'], ['writeonly']], order='C') # If a has just one element, it's fine too (constant 0 stride, a reduction) a = np.arange(1).reshape(1, 1, 1) nditer([a, b], ['buffered', 'external_loop', 'reduce_ok'], [['readwrite'], ['writeonly']], order='C') # check that it fails on other dimensions too a = np.arange(6).reshape(1, 3, 2) assert_raises(ValueError, nditer, [a, b], ['buffered', 'external_loop'], [['readwrite'], ['writeonly']], order='C') a = np.arange(4).reshape(2, 1, 2) assert_raises(ValueError, nditer, [a, b], ['buffered', 'external_loop'], [['readwrite'], ['writeonly']], order='C') def test_iter_buffering_string(): # Safe casting disallows shrinking strings a = np.array(['abc', 'a', 'abcd'], dtype=np.bytes_) assert_equal(a.dtype, np.dtype('S4')) assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'], op_dtypes='S2') i = nditer(a, ['buffered'], ['readonly'], op_dtypes='S6') assert_equal(i[0], b'abc') assert_equal(i[0].dtype, np.dtype('S6')) a = np.array(['abc', 'a', 'abcd'], dtype=np.str_) assert_equal(a.dtype, np.dtype('U4')) assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'], op_dtypes='U2') i = nditer(a, ['buffered'], ['readonly'], op_dtypes='U6') assert_equal(i[0], 'abc') assert_equal(i[0].dtype, np.dtype('U6')) def test_iter_buffering_growinner(): # Test that the inner loop grows when no buffering is needed a = np.arange(30) i = nditer(a, ['buffered', 'growinner', 'external_loop'], buffersize=5) # Should end up with just one inner loop here assert_equal(i[0].size, a.size) @pytest.mark.slow def test_iter_buffered_reduce_reuse(): # large enough array for all views, including negative strides. a = np.arange(2*3**5)[3**5:3**5+1] flags = ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok', 'refs_ok'] op_flags = [('readonly',), ('readwrite', 'allocate')] op_axes_list = [[(0, 1, 2), (0, 1, -1)], [(0, 1, 2), (0, -1, -1)]] # wrong dtype to force buffering op_dtypes = [float, a.dtype] def get_params(): for xs in range(-3**2, 3**2 + 1): for ys in range(xs, 3**2 + 1): for op_axes in op_axes_list: # last stride is reduced and because of that not # important for this test, as it is the inner stride. strides = (xs * a.itemsize, ys * a.itemsize, a.itemsize) arr = np.lib.stride_tricks.as_strided(a, (3, 3, 3), strides) for skip in [0, 1]: yield arr, op_axes, skip for arr, op_axes, skip in get_params(): nditer2 = np.nditer([arr.copy(), None], op_axes=op_axes, flags=flags, op_flags=op_flags, op_dtypes=op_dtypes) with nditer2: nditer2.operands[-1][...] = 0 nditer2.reset() nditer2.iterindex = skip for (a2_in, b2_in) in nditer2: b2_in += a2_in.astype(np.int_) comp_res = nditer2.operands[-1] for bufsize in range(0, 3**3): nditer1 = np.nditer([arr, None], op_axes=op_axes, flags=flags, op_flags=op_flags, buffersize=bufsize, op_dtypes=op_dtypes) with nditer1: nditer1.operands[-1][...] = 0 nditer1.reset() nditer1.iterindex = skip for (a1_in, b1_in) in nditer1: b1_in += a1_in.astype(np.int_) res = nditer1.operands[-1] assert_array_equal(res, comp_res) def test_iter_no_broadcast(): # Test that the no_broadcast flag works a = np.arange(24).reshape(2, 3, 4) b = np.arange(6).reshape(2, 3, 1) c = np.arange(12).reshape(3, 4) nditer([a, b, c], [], [['readonly', 'no_broadcast'], ['readonly'], ['readonly']]) assert_raises(ValueError, nditer, [a, b, c], [], [['readonly'], ['readonly', 'no_broadcast'], ['readonly']]) assert_raises(ValueError, nditer, [a, b, c], [], [['readonly'], ['readonly'], ['readonly', 'no_broadcast']]) class TestIterNested: def test_basic(self): # Test nested iteration basic usage a = arange(12).reshape(2, 3, 2) i, j = np.nested_iters(a, [[0], [1, 2]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]) i, j = np.nested_iters(a, [[0, 1], [2]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]]) i, j = np.nested_iters(a, [[0, 2], [1]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) def test_reorder(self): # Test nested iteration basic usage a = arange(12).reshape(2, 3, 2) # In 'K' order (default), it gets reordered i, j = np.nested_iters(a, [[0], [2, 1]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]) i, j = np.nested_iters(a, [[1, 0], [2]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]]) i, j = np.nested_iters(a, [[2, 0], [1]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) # In 'C' order, it doesn't i, j = np.nested_iters(a, [[0], [2, 1]], order='C') vals = [list(j) for _ in i] assert_equal(vals, [[0, 2, 4, 1, 3, 5], [6, 8, 10, 7, 9, 11]]) i, j = np.nested_iters(a, [[1, 0], [2]], order='C') vals = [list(j) for _ in i] assert_equal(vals, [[0, 1], [6, 7], [2, 3], [8, 9], [4, 5], [10, 11]]) i, j = np.nested_iters(a, [[2, 0], [1]], order='C') vals = [list(j) for _ in i] assert_equal(vals, [[0, 2, 4], [6, 8, 10], [1, 3, 5], [7, 9, 11]]) def test_flip_axes(self): # Test nested iteration with negative axes a = arange(12).reshape(2, 3, 2)[::-1, ::-1, ::-1] # In 'K' order (default), the axes all get flipped i, j = np.nested_iters(a, [[0], [1, 2]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]) i, j = np.nested_iters(a, [[0, 1], [2]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]]) i, j = np.nested_iters(a, [[0, 2], [1]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) # In 'C' order, flipping axes is disabled i, j = np.nested_iters(a, [[0], [1, 2]], order='C') vals = [list(j) for _ in i] assert_equal(vals, [[11, 10, 9, 8, 7, 6], [5, 4, 3, 2, 1, 0]]) i, j = np.nested_iters(a, [[0, 1], [2]], order='C') vals = [list(j) for _ in i] assert_equal(vals, [[11, 10], [9, 8], [7, 6], [5, 4], [3, 2], [1, 0]]) i, j = np.nested_iters(a, [[0, 2], [1]], order='C') vals = [list(j) for _ in i] assert_equal(vals, [[11, 9, 7], [10, 8, 6], [5, 3, 1], [4, 2, 0]]) def test_broadcast(self): # Test nested iteration with broadcasting a = arange(2).reshape(2, 1) b = arange(3).reshape(1, 3) i, j = np.nested_iters([a, b], [[0], [1]]) vals = [list(j) for _ in i] assert_equal(vals, [[[0, 0], [0, 1], [0, 2]], [[1, 0], [1, 1], [1, 2]]]) i, j = np.nested_iters([a, b], [[1], [0]]) vals = [list(j) for _ in i] assert_equal(vals, [[[0, 0], [1, 0]], [[0, 1], [1, 1]], [[0, 2], [1, 2]]]) def test_dtype_copy(self): # Test nested iteration with a copy to change dtype # copy a = arange(6, dtype='i4').reshape(2, 3) i, j = np.nested_iters(a, [[0], [1]], op_flags=['readonly', 'copy'], op_dtypes='f8') assert_equal(j[0].dtype, np.dtype('f8')) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1, 2], [3, 4, 5]]) vals = None # writebackifcopy - using context manager a = arange(6, dtype='f4').reshape(2, 3) i, j = np.nested_iters(a, [[0], [1]], op_flags=['readwrite', 'updateifcopy'], casting='same_kind', op_dtypes='f8') with i, j: assert_equal(j[0].dtype, np.dtype('f8')) for x in i: for y in j: y[...] += 1 assert_equal(a, [[0, 1, 2], [3, 4, 5]]) assert_equal(a, [[1, 2, 3], [4, 5, 6]]) # writebackifcopy - using close() a = arange(6, dtype='f4').reshape(2, 3) i, j = np.nested_iters(a, [[0], [1]], op_flags=['readwrite', 'updateifcopy'], casting='same_kind', op_dtypes='f8') assert_equal(j[0].dtype, np.dtype('f8')) for x in i: for y in j: y[...] += 1 assert_equal(a, [[0, 1, 2], [3, 4, 5]]) i.close() j.close() assert_equal(a, [[1, 2, 3], [4, 5, 6]]) def test_dtype_buffered(self): # Test nested iteration with buffering to change dtype a = arange(6, dtype='f4').reshape(2, 3) i, j = np.nested_iters(a, [[0], [1]], flags=['buffered'], op_flags=['readwrite'], casting='same_kind', op_dtypes='f8') assert_equal(j[0].dtype, np.dtype('f8')) for x in i: for y in j: y[...] += 1 assert_equal(a, [[1, 2, 3], [4, 5, 6]]) def test_0d(self): a = np.arange(12).reshape(2, 3, 2) i, j = np.nested_iters(a, [[], [1, 0, 2]]) vals = [list(j) for _ in i] assert_equal(vals, [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]]) i, j = np.nested_iters(a, [[1, 0, 2], []]) vals = [list(j) for _ in i] assert_equal(vals, [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11]]) i, j, k = np.nested_iters(a, [[2, 0], [], [1]]) vals = [] for x in i: for y in j: vals.append([z for z in k]) assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]]) def test_iter_nested_iters_dtype_buffered(self): # Test nested iteration with buffering to change dtype a = arange(6, dtype='f4').reshape(2, 3) i, j = np.nested_iters(a, [[0], [1]], flags=['buffered'], op_flags=['readwrite'], casting='same_kind', op_dtypes='f8') with i, j: assert_equal(j[0].dtype, np.dtype('f8')) for x in i: for y in j: y[...] += 1 assert_equal(a, [[1, 2, 3], [4, 5, 6]]) def test_iter_reduction_error(): a = np.arange(6) assert_raises(ValueError, nditer, [a, None], [], [['readonly'], ['readwrite', 'allocate']], op_axes=[[0], [-1]]) a = np.arange(6).reshape(2, 3) assert_raises(ValueError, nditer, [a, None], ['external_loop'], [['readonly'], ['readwrite', 'allocate']], op_axes=[[0, 1], [-1, -1]]) def test_iter_reduction(): # Test doing reductions with the iterator a = np.arange(6) i = nditer([a, None], ['reduce_ok'], [['readonly'], ['readwrite', 'allocate']], op_axes=[[0], [-1]]) # Need to initialize the output operand to the addition unit with i: i.operands[1][...] = 0 # Do the reduction for x, y in i: y[...] += x # Since no axes were specified, should have allocated a scalar assert_equal(i.operands[1].ndim, 0) assert_equal(i.operands[1], np.sum(a)) a = np.arange(6).reshape(2, 3) i = nditer([a, None], ['reduce_ok', 'external_loop'], [['readonly'], ['readwrite', 'allocate']], op_axes=[[0, 1], [-1, -1]]) # Need to initialize the output operand to the addition unit with i: i.operands[1][...] = 0 # Reduction shape/strides for the output assert_equal(i[1].shape, (6,)) assert_equal(i[1].strides, (0,)) # Do the reduction for x, y in i: # Use a for loop instead of ``y[...] += x`` # (equivalent to ``y[...] = y[...].copy() + x``), # because y has zero strides we use for the reduction for j in range(len(y)): y[j] += x[j] # Since no axes were specified, should have allocated a scalar assert_equal(i.operands[1].ndim, 0) assert_equal(i.operands[1], np.sum(a)) # This is a tricky reduction case for the buffering double loop # to handle a = np.ones((2, 3, 5)) it1 = nditer([a, None], ['reduce_ok', 'external_loop'], [['readonly'], ['readwrite', 'allocate']], op_axes=[None, [0, -1, 1]]) it2 = nditer([a, None], ['reduce_ok', 'external_loop', 'buffered', 'delay_bufalloc'], [['readonly'], ['readwrite', 'allocate']], op_axes=[None, [0, -1, 1]], buffersize=10) with it1, it2: it1.operands[1].fill(0) it2.operands[1].fill(0) it2.reset() for x in it1: x[1][...] += x[0] for x in it2: x[1][...] += x[0] assert_equal(it1.operands[1], it2.operands[1]) assert_equal(it2.operands[1].sum(), a.size) def test_iter_buffering_reduction(): # Test doing buffered reductions with the iterator a = np.arange(6) b = np.array(0., dtype='f8').byteswap().newbyteorder() i = nditer([a, b], ['reduce_ok', 'buffered'], [['readonly'], ['readwrite', 'nbo']], op_axes=[[0], [-1]]) with i: assert_equal(i[1].dtype, np.dtype('f8')) assert_(i[1].dtype != b.dtype) # Do the reduction for x, y in i: y[...] += x # Since no axes were specified, should have allocated a scalar assert_equal(b, np.sum(a)) a = np.arange(6).reshape(2, 3) b = np.array([0, 0], dtype='f8').byteswap().newbyteorder() i = nditer([a, b], ['reduce_ok', 'external_loop', 'buffered'], [['readonly'], ['readwrite', 'nbo']], op_axes=[[0, 1], [0, -1]]) # Reduction shape/strides for the output with i: assert_equal(i[1].shape, (3,)) assert_equal(i[1].strides, (0,)) # Do the reduction for x, y in i: # Use a for loop instead of ``y[...] += x`` # (equivalent to ``y[...] = y[...].copy() + x``), # because y has zero strides we use for the reduction for j in range(len(y)): y[j] += x[j] assert_equal(b, np.sum(a, axis=1)) # Iterator inner double loop was wrong on this one p = np.arange(2) + 1 it = np.nditer([p, None], ['delay_bufalloc', 'reduce_ok', 'buffered', 'external_loop'], [['readonly'], ['readwrite', 'allocate']], op_axes=[[-1, 0], [-1, -1]], itershape=(2, 2)) with it: it.operands[1].fill(0) it.reset() assert_equal(it[0], [1, 2, 1, 2]) # Iterator inner loop should take argument contiguity into account x = np.ones((7, 13, 8), np.int8)[4:6,1:11:6,1:5].transpose(1, 2, 0) x[...] = np.arange(x.size).reshape(x.shape) y_base = np.arange(4*4, dtype=np.int8).reshape(4, 4) y_base_copy = y_base.copy() y = y_base[::2,:,None] it = np.nditer([y, x], ['buffered', 'external_loop', 'reduce_ok'], [['readwrite'], ['readonly']]) with it: for a, b in it: a.fill(2) assert_equal(y_base[1::2], y_base_copy[1::2]) assert_equal(y_base[::2], 2) def test_iter_buffering_reduction_reuse_reduce_loops(): # There was a bug triggering reuse of the reduce loop inappropriately, # which caused processing to happen in unnecessarily small chunks # and overran the buffer. a = np.zeros((2, 7)) b = np.zeros((1, 7)) it = np.nditer([a, b], flags=['reduce_ok', 'external_loop', 'buffered'], op_flags=[['readonly'], ['readwrite']], buffersize=5) with it: bufsizes = [x.shape[0] for x, y in it] assert_equal(bufsizes, [5, 2, 5, 2]) assert_equal(sum(bufsizes), a.size) def test_iter_writemasked_badinput(): a = np.zeros((2, 3)) b = np.zeros((3,)) m = np.array([[True, True, False], [False, True, False]]) m2 = np.array([True, True, False]) m3 = np.array([0, 1, 1], dtype='u1') mbad1 = np.array([0, 1, 1], dtype='i1') mbad2 = np.array([0, 1, 1], dtype='f4') # Need an 'arraymask' if any operand is 'writemasked' assert_raises(ValueError, nditer, [a, m], [], [['readwrite', 'writemasked'], ['readonly']]) # A 'writemasked' operand must not be readonly assert_raises(ValueError, nditer, [a, m], [], [['readonly', 'writemasked'], ['readonly', 'arraymask']]) # 'writemasked' and 'arraymask' may not be used together assert_raises(ValueError, nditer, [a, m], [], [['readonly'], ['readwrite', 'arraymask', 'writemasked']]) # 'arraymask' may only be specified once assert_raises(ValueError, nditer, [a, m, m2], [], [['readwrite', 'writemasked'], ['readonly', 'arraymask'], ['readonly', 'arraymask']]) # An 'arraymask' with nothing 'writemasked' also doesn't make sense assert_raises(ValueError, nditer, [a, m], [], [['readwrite'], ['readonly', 'arraymask']]) # A writemasked reduction requires a similarly smaller mask assert_raises(ValueError, nditer, [a, b, m], ['reduce_ok'], [['readonly'], ['readwrite', 'writemasked'], ['readonly', 'arraymask']]) # But this should work with a smaller/equal mask to the reduction operand np.nditer([a, b, m2], ['reduce_ok'], [['readonly'], ['readwrite', 'writemasked'], ['readonly', 'arraymask']]) # The arraymask itself cannot be a reduction assert_raises(ValueError, nditer, [a, b, m2], ['reduce_ok'], [['readonly'], ['readwrite', 'writemasked'], ['readwrite', 'arraymask']]) # A uint8 mask is ok too np.nditer([a, m3], ['buffered'], [['readwrite', 'writemasked'], ['readonly', 'arraymask']], op_dtypes=['f4', None], casting='same_kind') # An int8 mask isn't ok assert_raises(TypeError, np.nditer, [a, mbad1], ['buffered'], [['readwrite', 'writemasked'], ['readonly', 'arraymask']], op_dtypes=['f4', None], casting='same_kind') # A float32 mask isn't ok assert_raises(TypeError, np.nditer, [a, mbad2], ['buffered'], [['readwrite', 'writemasked'], ['readonly', 'arraymask']], op_dtypes=['f4', None], casting='same_kind') def _is_buffered(iterator): try: iterator.itviews except ValueError: return True return False @pytest.mark.parametrize("a", [np.zeros((3,), dtype='f8'), np.zeros((9876, 3*5), dtype='f8')[::2, :], np.zeros((4, 312, 124, 3), dtype='f8')[::2, :, ::2, :], # Also test with the last dimension strided (so it does not fit if # there is repeated access) np.zeros((9,), dtype='f8')[::3], np.zeros((9876, 3*10), dtype='f8')[::2, ::5], np.zeros((4, 312, 124, 3), dtype='f8')[::2, :, ::2, ::-1]]) def test_iter_writemasked(a): # Note, the slicing above is to ensure that nditer cannot combine multiple # axes into one. The repetition is just to make things a bit more # interesting. shape = a.shape reps = shape[-1] // 3 msk = np.empty(shape, dtype=bool) msk[...] = [True, True, False] * reps # When buffering is unused, 'writemasked' effectively does nothing. # It's up to the user of the iterator to obey the requested semantics. it = np.nditer([a, msk], [], [['readwrite', 'writemasked'], ['readonly', 'arraymask']]) with it: for x, m in it: x[...] = 1 # Because we violated the semantics, all the values became 1 assert_equal(a, np.broadcast_to([1, 1, 1] * reps, shape)) # Even if buffering is enabled, we still may be accessing the array # directly. it = np.nditer([a, msk], ['buffered'], [['readwrite', 'writemasked'], ['readonly', 'arraymask']]) # @seberg: I honestly don't currently understand why a "buffered" iterator # would end up not using a buffer for the small array here at least when # "writemasked" is used, that seems confusing... Check by testing for # actual memory overlap! is_buffered = True with it: for x, m in it: x[...] = 2.5 if np.may_share_memory(x, a): is_buffered = False if not is_buffered: # Because we violated the semantics, all the values became 2.5 assert_equal(a, np.broadcast_to([2.5, 2.5, 2.5] * reps, shape)) else: # For large sizes, the iterator may be buffered: assert_equal(a, np.broadcast_to([2.5, 2.5, 1] * reps, shape)) a[...] = 2.5 # If buffering will definitely happening, for instance because of # a cast, only the items selected by the mask will be copied back from # the buffer. it = np.nditer([a, msk], ['buffered'], [['readwrite', 'writemasked'], ['readonly', 'arraymask']], op_dtypes=['i8', None], casting='unsafe') with it: for x, m in it: x[...] = 3 # Even though we violated the semantics, only the selected values # were copied back assert_equal(a, np.broadcast_to([3, 3, 2.5] * reps, shape)) @pytest.mark.parametrize(["mask", "mask_axes"], [ # Allocated operand (only broadcasts with -1) (None, [-1, 0]), # Reduction along the first dimension (with and without op_axes) (np.zeros((1, 4), dtype="bool"), [0, 1]), (np.zeros((1, 4), dtype="bool"), None), # Test 0-D and -1 op_axes (np.zeros(4, dtype="bool"), [-1, 0]), (np.zeros((), dtype="bool"), [-1, -1]), (np.zeros((), dtype="bool"), None)]) def test_iter_writemasked_broadcast_error(mask, mask_axes): # This assumes that a readwrite mask makes sense. This is likely not the # case and should simply be deprecated. arr = np.zeros((3, 4)) itflags = ["reduce_ok"] mask_flags = ["arraymask", "readwrite", "allocate"] a_flags = ["writeonly", "writemasked"] if mask_axes is None: op_axes = None else: op_axes = [mask_axes, [0, 1]] with assert_raises(ValueError): np.nditer((mask, arr), flags=itflags, op_flags=[mask_flags, a_flags], op_axes=op_axes) def test_iter_writemasked_decref(): # force casting (to make it interesting) by using a structured dtype. arr = np.arange(10000).astype(">i,O") original = arr.copy() mask = np.random.randint(0, 2, size=10000).astype(bool) it = np.nditer([arr, mask], ['buffered', "refs_ok"], [['readwrite', 'writemasked'], ['readonly', 'arraymask']], op_dtypes=["