NumPy - Iterating Over Array

NumPy - Iterating Over Arrays

Iterating over elements is a common requirement when processing arrays. While standard Python for loops work on NumPy arrays, they are often inefficient for large arrays, especially multi-dimensional ones. NumPy provides a highly optimized iterator object, numpy.nditer, for efficient multi-dimensional iteration.

Basic Iteration with np.nditer

np.nditer allows you to visit each element of an array using Python’s standard iterator protocol (for ... in ...). By default, it selects an iteration order that matches the array’s memory layout for better performance.

Example 1: Iterating a 2D Array

import numpy as np

a = np.arange(0, 60, 5).reshape(3, 4)

print(f"Original array:\n{a}\n")

print("Iterating elements using np.nditer:")
for x in np.nditer(a):
    print(x, end=' ')
print() # For newline

Output (Default order is usually C-style, row-major):

Original array:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

Iterating elements using np.nditer:
0 5 10 15 20 25 30 35 40 45 50 55

Example 2: Iteration Order and Transpose

The default iteration order depends on how the array is stored in memory. Iterating over a transposed array (a.T) might yield elements in a different order if the transpose changes the memory contiguity.

import numpy as np

a = np.arange(0, 60, 5).reshape(3, 4)

print(f"Original array (C-order layout):\n{a}\n")

print("Transpose of the array:")
b = a.T
print(f"{b}\n")

print("Iterating over the transpose 'b':") # Order might still be C-like based on underlying data
for x in np.nditer(b):
    print(x, end=' ')
print()

Output:

Original array (C-order layout):
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

Transpose of the array:
[[ 0 20 40]
 [ 5 25 45]
 [10 30 50]
 [15 35 55]]

Iterating over the transpose 'b':
0 5 10 15 20 25 30 35 40 45 50 55

Note: Even though b is the transpose, np.nditer might still choose the C-order traversal of the underlying data if it’s more efficient. The actual order depends on memory layout.

Controlling Iteration Order

You can force np.nditer to follow a specific iteration order (C-style/row-major or F-style/column-major) using the order parameter.

• order='C': Force C-order (row-major).
• order='F': Force Fortran-order (column-major).
• order='A': Iterate in Fortran order if the array is F-contiguous, C order otherwise.
• order='K': Iterate in the order the elements are laid out in memory (Keep order - default).

Example: C vs F Order Iteration

import numpy as np

a = np.arange(0, 60, 5).reshape(3, 4)
print(f"Original array:\n{a}\n")

print("Iterating in C-style order (row-major):")
for x in np.nditer(a, order='C'):
    print(x, end=' ')
print('\n')

print("Iterating in F-style order (column-major):")
for x in np.nditer(a, order='F'):
    print(x, end=' ')
print()

Output:

Original array:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

Iterating in C-style order (row-major):
0 5 10 15 20 25 30 35 40 45 50 55

Iterating in F-style order (column-major):
0 20 40 5 25 45 10 30 50 15 35 55

Example: Order based on Memory Layout

If an array is explicitly created or copied with Fortran order, np.nditer (with default ‘K’ or explicit ‘F’) will iterate column-wise efficiently.

import numpy as np

a = np.arange(0, 60, 5).reshape(3, 4)
b = a.T # Transpose

print("Iterating 'b' (transpose) with default ('K') order:")
# Still iterates C-style over the original data layout
for x in np.nditer(b):
    print(x, end=' ')
print('\n')

print("Creating a copy of 'b' in F-order ('c = b.copy(order='F')'):")
c = b.copy(order='F') # Now 'c' is genuinely F-contiguous
print(f"Is 'c' F-contiguous? {c.flags.f_contiguous}")

print("Iterating 'c' with default ('K') order (now efficient F-style):")
for x in np.nditer(c):
    print(x, end=' ')
print()

Output:

Iterating 'b' (transpose) with default ('K') order:
0 5 10 15 20 25 30 35 40 45 50 55

Creating a copy of 'b' in F-order ('c = b.copy(order='F')'):
Is 'c' F-contiguous? True
Iterating 'c' with default ('K') order (now efficient F-style):
0 20 40 5 25 45 10 30 50 15 35 55

Modifying Array Values During Iteration

By default, np.nditer treats the array elements as read-only. To modify elements during iteration, you need to specify the appropriate op_flags parameter. Common values are ['readonly'] (default), ['readwrite'], or ['writeonly'].

When modifying, you need to use the ellipsis (...) to assign to the yielded element: x[...] = value.

Example

import numpy as np

a = np.arange(0, 60, 5).reshape(3, 4)
print(f"Original array:\n{a}\n")

# Iterate and modify elements (multiply by 2)
for x in np.nditer(a, op_flags=['readwrite']):
    x[...] = 2 * x

print(f"Modified array:\n{a}")

Output:

Original array:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

Modified array:
[[  0  10  20  30]
 [ 40  50  60  70]
 [ 80  90 100 110]]

External Loop (flags=['external_loop'])

The flags=['external_loop'] option changes the iterator’s behavior. Instead of yielding individual elements (0-dimensional arrays), it yields 1-dimensional chunks of the array. This can be more efficient for some operations as it reduces the overhead of the Python loop.

The size of the chunk is determined by buffering, and the iteration follows the specified order.

Example

Traversing columns using external loop and F-order:

import numpy as np

a = np.arange(0, 60, 5).reshape(3, 4)

print(f"Original array:\n{a}\n")

print("Iterating with external_loop (F-order yields columns as 1D arrays):")
for x in np.nditer(a, flags=['external_loop'], order='F'):
    print(x, end=' ')
print()

Output:

Original array:
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

Iterating with external_loop (F-order yields columns as 1D arrays):
[ 0 20 40] [ 5 25 45] [10 30 50] [15 35 55]

Broadcasting Iteration

If you provide multiple arrays to np.nditer that are broadcastable to a common shape, the iterator can traverse them concurrently, yielding tuples of corresponding elements.

Example

Iterating over a 3x4 array a and a 1x4 array b simultaneously. b will be broadcast to match a’s shape during iteration.

import numpy as np

a = np.arange(0, 60, 5).reshape(3, 4)
print(f"First array 'a':\n{a}\n")

b = np.array([1, 2, 3, 4], dtype=int)
print(f"Second array 'b': {b}\n")

print("Iterating over 'a' and broadcast 'b' concurrently:")
for x, y in np.nditer([a, b]):
    print(f"{x}:{y}", end=' ')
print()

Output:

First array 'a':
[[ 0  5 10 15]
 [20 25 30 35]
 [40 45 50 55]]

Second array 'b': [1 2 3 4]

Iterating over 'a' and broadcast 'b' concurrently:
0:1 5:2 10:3 15:4 20:1 25:2 30:3 35:4 40:1 45:2 50:3 55:4