Controlling iteration order#

There are times when it is important to visit the elements of an arrayin a specific order, irrespective of the layout of the elements in memory.Thenditer object provides anorder parameter to control thisaspect of iteration. The default, having the behavior described above,is order=’K’ to keep the existing order. This can be overridden withorder=’C’ for C order and order=’F’ for Fortran order.

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)>>>forxinnp.nditer(a,order='F'):...print(x,end=' ')...0 3 1 4 2 5>>>forxinnp.nditer(a.T,order='C'):...print(x,end=' ')...0 3 1 4 2 5

Modifying array values#

By default, thenditer treats the input operand as a read-onlyobject. To be able to modify the array elements, you must specify eitherread-write or write-only mode using the‘readwrite’ or‘writeonly’per-operand flags.

The nditer will then yield writeable buffer arrays which you may modify. However,because the nditer must copy this buffer data back to the original array onceiteration is finished, you must signal when the iteration is ended, by one of twomethods. You may either:

• used the nditer as a context manager using thewith statement, andthe temporary data will be written back when the context is exited.

• call the iterator’sclose method once finished iterating, which will triggerthe write-back.

The nditer can no longer be iterated once eitherclose is called or itscontext is exited.

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)>>>aarray([[0, 1, 2],       [3, 4, 5]])>>>withnp.nditer(a,op_flags=['readwrite'])asit:...forxinit:...x[...]=2*x...>>>aarray([[ 0,  2,  4],       [ 6,  8, 10]])

If you are writing code that needs to support older versions of numpy,note that prior to 1.15,nditer was not a context manager anddid not have aclose method. Instead it relied on the destructor toinitiate the writeback of the buffer.

Using an external loop#

In all the examples so far, the elements ofa are provided by theiterator one at a time, because all the looping logic is internal to theiterator. While this is simple and convenient, it is not very efficient.A better approach is to move the one-dimensional innermost loop into yourcode, external to the iterator. This way, NumPy’s vectorized operationscan be used on larger chunks of the elements being visited.

Thenditer will try to provide chunks that areas large as possible to the inner loop. By forcing ‘C’ and ‘F’ order,we get different external loop sizes. This mode is enabled by specifyingan iterator flag.

Observe that with the default of keeping native memory order, theiterator is able to provide a single one-dimensional chunk, whereaswhen forcing Fortran order, it has to provide three chunks of twoelements each.

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)>>>forxinnp.nditer(a,flags=['external_loop']):...print(x,end=' ')...[0 1 2 3 4 5]

>>>forxinnp.nditer(a,flags=['external_loop'],order='F'):...print(x,end=' ')...[0 3] [1 4] [2 5]

Tracking an index or multi-index#

During iteration, you may want to use the index of the currentelement in a computation. For example, you may want to visit theelements of an array in memory order, but use a C-order, Fortran-order,or multidimensional index to look up values in a different array.

The index is tracked by the iterator object itself, and accessiblethrough theindex ormulti_index properties, depending on what wasrequested. The examples below show printouts demonstrating theprogression of the index:

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)>>>it=np.nditer(a,flags=['f_index'])>>>forxinit:...print("%d <%d>"%(x,it.index),end=' ')...0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>

>>>it=np.nditer(a,flags=['multi_index'])>>>forxinit:...print("%d <%s>"%(x,it.multi_index),end=' ')...0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>

>>>withnp.nditer(a,flags=['multi_index'],op_flags=['writeonly'])asit:...forxinit:...x[...]=it.multi_index[1]-it.multi_index[0]...>>>aarray([[ 0,  1,  2],       [-1,  0,  1]])

Tracking an index or multi-index is incompatible with using an externalloop, because it requires a different index value per element. Ifyou try to combine these flags, thenditer object willraise an exception.

>>>importnumpyasnp

>>>a=np.zeros((2,3))>>>it=np.nditer(a,flags=['c_index','external_loop'])Traceback (most recent call last):  File"<stdin>", line1, in<module>ValueError:Iterator flag EXTERNAL_LOOP cannot be used if an index or multi-index is being tracked

Alternative looping and element access#

To make its properties more readily accessible during iteration,nditer has an alternative syntax for iterating, which worksexplicitly with the iterator object itself. With this looping construct,the current value is accessible by indexing into the iterator. Otherproperties, such as tracked indices remain as before. The examples belowproduce identical results to the ones in the previous section.

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)>>>it=np.nditer(a,flags=['f_index'])>>>whilenotit.finished:...print("%d <%d>"%(it[0],it.index),end=' ')...is_not_finished=it.iternext()...0 <0> 1 <2> 2 <4> 3 <1> 4 <3> 5 <5>

>>>it=np.nditer(a,flags=['multi_index'])>>>whilenotit.finished:...print("%d <%s>"%(it[0],it.multi_index),end=' ')...is_not_finished=it.iternext()...0 <(0, 0)> 1 <(0, 1)> 2 <(0, 2)> 3 <(1, 0)> 4 <(1, 1)> 5 <(1, 2)>

>>>withnp.nditer(a,flags=['multi_index'],op_flags=['writeonly'])asit:...whilenotit.finished:...it[0]=it.multi_index[1]-it.multi_index[0]...is_not_finished=it.iternext()...>>>aarray([[ 0,  1,  2],       [-1,  0,  1]])

Buffering the array elements#

When forcing an iteration order, we observed that the external loopoption may provide the elements in smaller chunks because the elementscan’t be visited in the appropriate order with a constant stride.When writing C code, this is generally fine, however in pure Python codethis can cause a significant reduction in performance.

By enabling buffering mode, the chunks provided by the iterator tothe inner loop can be made larger, significantly reducing the overheadof the Python interpreter. In the example forcing Fortran iteration order,the inner loop gets to see all the elements in one go when bufferingis enabled.

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)>>>forxinnp.nditer(a,flags=['external_loop'],order='F'):...print(x,end=' ')...[0 3] [1 4] [2 5]

>>>forxinnp.nditer(a,flags=['external_loop','buffered'],order='F'):...print(x,end=' ')...[0 3 1 4 2 5]

Iterating as a specific data type#

There are times when it is necessary to treat an array as a differentdata type than it is stored as. For instance, one may want to do allcomputations on 64-bit floats, even if the arrays being manipulatedare 32-bit floats. Except when writing low-level C code, it’s generallybetter to let the iterator handle the copying or buffering insteadof casting the data type yourself in the inner loop.

There are two mechanisms which allow this to be done, temporary copiesand buffering mode. With temporary copies, a copy of the entire array ismade with the new data type, then iteration is done in the copy. Writeaccess is permitted through a mode which updates the original array afterall the iteration is complete. The major drawback of temporary copies isthat the temporary copy may consume a large amount of memory, particularlyif the iteration data type has a larger itemsize than the original one.

Buffering mode mitigates the memory usage issue and is more cache-friendlythan making temporary copies. Except for special cases, where the wholearray is needed at once outside the iterator, buffering is recommendedover temporary copying. Within NumPy, buffering is used by the ufuncs andother functions to support flexible inputs with minimal memory overhead.

In our examples, we will treat the input array with a complex data type,so that we can take square roots of negative numbers. Without enablingcopies or buffering mode, the iterator will raise an exception if thedata type doesn’t match precisely.

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)-3>>>forxinnp.nditer(a,op_dtypes=['complex128']):...print(np.sqrt(x),end=' ')...Traceback (most recent call last):  File"<stdin>", line1, in<module>TypeError:Iterator operand required copying or buffering, but neither copying nor buffering was enabled

In copying mode, ‘copy’ is specified as a per-operand flag. This isdone to provide control in a per-operand fashion. Buffering mode isspecified as an iterator flag.

>>>importnumpyasnp

>>>a=np.arange(6).reshape(2,3)-3>>>forxinnp.nditer(a,op_flags=['readonly','copy'],...op_dtypes=['complex128']):...print(np.sqrt(x),end=' ')...1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)

>>>forxinnp.nditer(a,flags=['buffered'],op_dtypes=['complex128']):...print(np.sqrt(x),end=' ')...1.7320508075688772j 1.4142135623730951j 1j 0j (1+0j) (1.4142135623730951+0j)

The iterator uses NumPy’s casting rules to determine whether a specificconversion is permitted. By default, it enforces ‘safe’ casting. This means,for example, that it will raise an exception if you try to treat a64-bit float array as a 32-bit float array. In many cases, the rule‘same_kind’ is the most reasonable rule to use, since it will allowconversion from 64 to 32-bit float, but not from float to int or fromcomplex to float.

>>>importnumpyasnp

>>>a=np.arange(6.)>>>forxinnp.nditer(a,flags=['buffered'],op_dtypes=['float32']):...print(x,end=' ')...Traceback (most recent call last):  File"<stdin>", line1, in<module>TypeError:Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('float32') according to the rule 'safe'

>>>forxinnp.nditer(a,flags=['buffered'],op_dtypes=['float32'],...casting='same_kind'):...print(x,end=' ')...0.0 1.0 2.0 3.0 4.0 5.0

>>>forxinnp.nditer(a,flags=['buffered'],op_dtypes=['int32'],casting='same_kind'):...print(x,end=' ')...Traceback (most recent call last):  File"<stdin>", line1, in<module>TypeError:Iterator operand 0 dtype could not be cast from dtype('float64') to dtype('int32') according to the rule 'same_kind'

One thing to watch out for is conversions back to the original datatype when using a read-write or write-only operand. A common case isto implement the inner loop in terms of 64-bit floats, and use ‘same_kind’casting to allow the other floating-point types to be processed as well.While in read-only mode, an integer array could be provided, read-writemode will raise an exception because conversion back to the arraywould violate the casting rule.

>>>importnumpyasnp

>>>a=np.arange(6)>>>forxinnp.nditer(a,flags=['buffered'],op_flags=['readwrite'],...op_dtypes=['float64'],casting='same_kind'):...x[...]=x/2.0...Traceback (most recent call last):  File"<stdin>", line2, in<module>TypeError:Iterator requested dtype could not be cast from dtype('float64') to dtype('int64'), the operand 0 dtype, according to the rule 'same_kind'

Iterator-allocated output arrays#

A common case in NumPy functions is to have outputs allocated basedon the broadcasting of the input, and additionally have an optionalparameter called ‘out’ where the result will be placed when it isprovided. Thenditer object provides a convenient idiom thatmakes it very easy to support this mechanism.

We’ll show how this works by creating a functionsquare which squaresits input. Let’s start with a minimal function definition excluding ‘out’parameter support.

>>>importnumpyasnp

>>>defsquare(a):...withnp.nditer([a,None])asit:...forx,yinit:...y[...]=x*x...returnit.operands[1]...>>>square([1,2,3])array([1, 4, 9])

By default, thenditer uses the flags ‘allocate’ and ‘writeonly’for operands that are passed in as None. This means we were able to providejust the two operands to the iterator, and it handled the rest.

When adding the ‘out’ parameter, we have to explicitly provide those flags,because if someone passes in an array as ‘out’, the iterator will defaultto ‘readonly’, and our inner loop would fail. The reason ‘readonly’ isthe default for input arrays is to prevent confusion about unintentionallytriggering a reduction operation. If the default were ‘readwrite’, anybroadcasting operation would also trigger a reduction, a topicwhich is covered later in this document.

While we’re at it, let’s also introduce the ‘no_broadcast’ flag, whichwill prevent the output from being broadcast. This is important, becausewe only want one input value for each output. Aggregating more than oneinput value is a reduction operation which requires special handling.It would already raise an error because reductions must be explicitlyenabled in an iterator flag, but the error message that results fromdisabling broadcasting is much more understandable for end-users.To see how to generalize the square function to a reduction, lookat the sum of squares function in the section about Cython.

For completeness, we’ll also add the ‘external_loop’ and ‘buffered’flags, as these are what you will typically want for performancereasons.

>>>importnumpyasnp

>>>defsquare(a,out=None):...it=np.nditer([a,out],...flags=['external_loop','buffered'],...op_flags=[['readonly'],...['writeonly','allocate','no_broadcast']])...withit:...forx,yinit:...y[...]=x*x...returnit.operands[1]...

>>>square([1,2,3])array([1, 4, 9])

>>>b=np.zeros((3,))>>>square([1,2,3],out=b)array([1.,  4.,  9.])>>>barray([1.,  4.,  9.])

>>>square(np.arange(6).reshape(2,3),out=b)Traceback (most recent call last):...ValueError:non-broadcastable output operand with shape (3,) doesn'tmatch the broadcast shape (2,3)

Outer product iteration#

Any binary operation can be extended to an array operation in an outerproduct fashion like inouter, and thenditer objectprovides a way to accomplish this by explicitly mapping the axes ofthe operands. It is also possible to do this withnewaxisindexing, but we will show you how to directly use the nditerop_axesparameter to accomplish this with no intermediate views.

We’ll do a simple outer product, placing the dimensions of the firstoperand before the dimensions of the second operand. Theop_axesparameter needs one list of axes for each operand, and provides a mappingfrom the iterator’s axes to the axes of the operand.

Suppose the first operand is one dimensional and the second operand istwo dimensional. The iterator will have three dimensions, soop_axeswill have two 3-element lists. The first list picks out the oneaxis of the first operand, and is -1 for the rest of the iterator axes,with a final result of [0, -1, -1]. The second list picks out the twoaxes of the second operand, but shouldn’t overlap with the axes pickedout in the first operand. Its list is [-1, 0, 1]. The output operandmaps onto the iterator axes in the standard manner, so we can provideNone instead of constructing another list.

The operation in the inner loop is a straightforward multiplication.Everything to do with the outer product is handled by the iterator setup.

>>>importnumpyasnp

>>>a=np.arange(3)>>>b=np.arange(8).reshape(2,4)>>>it=np.nditer([a,b,None],flags=['external_loop'],...op_axes=[[0,-1,-1],[-1,0,1],None])>>>withit:...forx,y,zinit:...z[...]=x*y...result=it.operands[2]# same as z...>>>resultarray([[[ 0,  0,  0,  0],        [ 0,  0,  0,  0]],        [[ 0,  1,  2,  3],        [ 4,  5,  6,  7]],        [[ 0,  2,  4,  6],        [ 8, 10, 12, 14]]])

Note that once the iterator is closed we can not accessoperandsand must use a reference created inside the context manager.

Reduction iteration#

Whenever a writeable operand has fewer elements than the full iteration space,that operand is undergoing a reduction. Thenditer object requiresthat any reduction operand be flagged as read-write, and only allowsreductions when ‘reduce_ok’ is provided as an iterator flag.

For a simple example, consider taking the sum of all elements in an array.

>>>importnumpyasnp

>>>a=np.arange(24).reshape(2,3,4)>>>b=np.array(0)>>>withnp.nditer([a,b],flags=['reduce_ok'],...op_flags=[['readonly'],['readwrite']])asit:...forx,yinit:...y[...]+=x...>>>barray(276)>>>np.sum(a)276

Things are a little bit more tricky when combining reduction and allocatedoperands. Before iteration is started, any reduction operand must beinitialized to its starting values. Here’s how we can do this, takingsums along the last axis ofa.

>>>importnumpyasnp

>>>a=np.arange(24).reshape(2,3,4)>>>it=np.nditer([a,None],flags=['reduce_ok'],...op_flags=[['readonly'],['readwrite','allocate']],...op_axes=[None,[0,1,-1]])>>>withit:...it.operands[1][...]=0...forx,yinit:...y[...]+=x...result=it.operands[1]...>>>resultarray([[ 6, 22, 38],       [54, 70, 86]])>>>np.sum(a,axis=2)array([[ 6, 22, 38],       [54, 70, 86]])

To do buffered reduction requires yet another adjustment during thesetup. Normally the iterator construction involves copying the firstbuffer of data from the readable arrays into the buffer. Any reductionoperand is readable, so it may be read into a buffer. Unfortunately,initialization of the operand after this buffering operation is completewill not be reflected in the buffer that the iteration starts with, andgarbage results will be produced.

The iterator flag “delay_bufalloc” is there to allowiterator-allocated reduction operands to exist together with buffering.When this flag is set, the iterator will leave its buffers uninitializeduntil it receives a reset, after which it will be ready for regulariteration. Here’s how the previous example looks if we also enablebuffering.

>>>importnumpyasnp

>>>a=np.arange(24).reshape(2,3,4)>>>it=np.nditer([a,None],flags=['reduce_ok',...'buffered','delay_bufalloc'],...op_flags=[['readonly'],['readwrite','allocate']],...op_axes=[None,[0,1,-1]])>>>withit:...it.operands[1][...]=0...it.reset()...forx,yinit:...y[...]+=x...result=it.operands[1]...>>>resultarray([[ 6, 22, 38],       [54, 70, 86]])