Search code examples
pythonnumpynumpy-ndarray

Looping through Numpy Array elements


Is there a more readable way to code a loop in Python that goes through each element of a Numpy array? I have come up with the following code, but it seems cumbersome & not very readable:

import numpy as np
arr01 = np.random.randint(1,10,(3,3))
for i in range(0,(np.shape(arr01[0])[0]+1)):
    for j in range(0,(np.shape(arr01[1])[0]+1)):
        print (arr01[i,j])

I could make it more explicit such as:

import numpy as np
arr01 = np.random.randint(1,10,(3,3))
rows = np.shape(arr01[0])[0]
cols = np.shape(arr01[1])[0]
for i in range(0, (rows + 1)):
    for j in range(0, (cols + 1)):
        print (arr01[i,j])

However, that still seems a bit more cumbersome, compared to other languages, i.e. an equivalent code in VBA could read (supposing the array had already been populated):

dim i, j as integer
for i = lbound(arr01,1) to ubound(arr01,1)
   for j = lbound(arr01,2) to ubound(arr01,2)
       msgBox arr01(i, j)
   next j
next i

Solution

  • Seems like you've skipped over some intro Python chapters. With a list there are several simple ways of iterating:

    In [1]: alist = ['a','b','c']                                                   
    In [2]: for i in alist: print(i)        # on the list itself                                                
    a
    b
    c
    In [3]: len(alist)                                                              
    Out[3]: 3
    In [4]: for i in range(len(alist)): print(i,alist[i])  # index is ok                 
    0 a
    1 b
    2 c
    In [5]: for i,v in enumerate(alist): print(i,v)   # but enumerate is simpler            
    0 a
    1 b
    2 c
    

    Note the indexes. range(3) is sufficient. alist[3] produces an error.

    In [6]: arr = np.arange(6).reshape(2,3)                                         
    In [7]: arr                                                                     
    Out[7]: 
    array([[0, 1, 2],
           [3, 4, 5]])
    In [8]: for row in arr: 
       ...:     for col in row: 
       ...:         print(row,col) 
       ...:                                                                         
    [0 1 2] 0
    [0 1 2] 1
    [0 1 2] 2
    [3 4 5] 3
    [3 4 5] 4
    [3 4 5] 5
    

    The shape is a tuple. The row count is then arr.shape[0], and columns arr.shape[1]. Or you can 'unpack' both at once:

    In [9]: arr.shape                                                               
    Out[9]: (2, 3)
    In [10]: n,m = arr.shape                                                        
    In [11]: [arr[i,j] for i in range(n) for j in range(m)]                         
    Out[11]: [0, 1, 2, 3, 4, 5]
    

    But we can get the same flat list of values with ravel and optional conversion to list:

    In [12]: arr.ravel()                                                            
    Out[12]: array([0, 1, 2, 3, 4, 5])
    In [13]: arr.ravel().tolist()                                                   
    Out[13]: [0, 1, 2, 3, 4, 5]
    

    But usually with numpy arrays, you shouldn't be iterating at all. Learn enough of the numpy basics so you can work with the whole array, not elements.

    nditer can be used, as the other answer shows, to iterate through an array in a flat manner, but there are a number of details about it that could easily confuse a beginner. There are a couple of intro pages to nditer, but they should be read in full. Usually I discourage its use.

    In [14]: for i in np.nditer(arr): 
        ...:     print(i, type(i), i.shape) 
        ...:                                                                        
    0 <class 'numpy.ndarray'> ()        # this element is a 0d array, not a scalar integer
    1 <class 'numpy.ndarray'> ()
    2 <class 'numpy.ndarray'> ()
    ... 
    

    Iterating with ndenumerate or on the tolist produce different types of elements. The type may matter if you try to do more than display the value, so be careful.

    In [15]: list(np.ndenumerate(arr))                                              
    Out[15]: [((0, 0), 0), ((0, 1), 1), ((0, 2), 2), ((1, 0), 3), ((1, 1), 4), ((1, 2), 5)]
    In [16]: for ij, v in np.ndenumerate(arr): 
        ...:     print(ij, v, type(v)) 
        ...:                                                                        
    (0, 0) 0 <class 'numpy.int64'>
    (0, 1) 1 <class 'numpy.int64'>
    ...
    In [17]: for i, v in enumerate(arr.ravel().tolist()): 
        ...:     print(i, v, type(v)) 
        ...:                                                                        
    0 0 <class 'int'>
    1 1 <class 'int'>
    ...