How to iterate 1d NumPy array with index and value
There are a few alternatives. The below assumes you are iterating over a 1d NumPy array.
Iterate with range
for j in range(theta.shape[0]): # or range(len(theta))
some_function(j, theta[j], theta)
Note this is the only of the 3 solutions which will work with numba
. This is noteworthy since iterating over a NumPy array explicitly is usually only efficient when combined with numba
or another means of pre-compilation.
Iterate with enumerate
for idx, j in enumerate(theta):
some_function(idx, j, theta)
The most efficient of the 3 solutions for 1d arrays. See benchmarking below.
Iterate with np.ndenumerate
for idx, j in np.ndenumerate(theta):
some_function(idx[0], j, theta)
Notice the additional indexing step in idx[0]
. This is necessary since the index (like shape
) of a 1d NumPy array is given as a singleton tuple. For a 1d array, np.ndenumerate
is inefficient; its benefits only show for multi-dimensional arrays.
Performance benchmarking
# Python 3.7, NumPy 1.14.3
np.random.seed(0)
arr = np.random.random(10**6)
def enumerater(arr):
for index, value in enumerate(arr):
index, value
pass
def ranger(arr):
for index in range(len(arr)):
index, arr[index]
pass
def ndenumerater(arr):
for index, value in np.ndenumerate(arr):
index[0], value
pass
%timeit enumerater(arr) # 131 ms
%timeit ranger(arr) # 171 ms
%timeit ndenumerater(arr) # 579 ms
You can use numpy.ndenumerate
for example
import numpy as np
test_array = np.arange(2, 3, 0.1)
for index, value in np.ndenumerate(test_array):
print(index[0], value)
For more information refer to https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndenumerate.html