CODEWITHSUNDEEP
pythonarraysnumpy
Prettymath77
Prettymath77

Reputation: 25

Two for loops on Numpy array

Goodmorning. Suppose I have a two-dimensional array (call it MAT(x,y)) created with numpy. On this array, I have to perform some operations. How can I rewrite the following 2 for loops, for example using np.nditer() or something else that uses numpy method? Thank you.

    for i in range(x): 
        for j in range(y): 

            if i == 0: MAT[i][j] = j   
            elif j == 0: MAT[i][j] = i

Upvotes: 0

Views: 643

Answers (4)

norok2
norok2

Reputation: 26886

You can basically assign np.arange() to an appropriate slicing of your input.

You can do this either hard-coding the 2D nature of your input (foo2()), or, for arbitrary dimensions using a dynamically defined slicing (foon()):

import numpy as np


def foo(arr):
    ii, jj = arr.shape
    for i in range(ii): 
        for j in range(jj): 
            if i == 0:
                arr[i, j] = j   
            elif j == 0:
                arr[i, j] = i
    return arr


def foo2(arr):
    ii, jj = arr.shape
    arr[:, 0] = np.arange(ii)
    arr[0, :] = np.arange(jj)
    return arr


def foon(arr):
    for i, d in enumerate(arr.shape):
        slicing = tuple(slice(None) if j == i else 0 for j in range(arr.ndim))
        arr[slicing] = np.arange(d)
    return arr


arr = np.zeros((3, 4))

print(foo(arr))
# [[0. 1. 2. 3.]
#  [1. 0. 0. 0.]
#  [2. 0. 0. 0.]]

print(foo2(arr))
# [[0. 1. 2. 3.]
#  [1. 0. 0. 0.]
#  [2. 0. 0. 0.]]

print(foon(arr))
# [[0. 1. 2. 3.]
#  [1. 0. 0. 0.]
#  [2. 0. 0. 0.]]

Note that double slicing (e.g. MAT[i][j]), while working, is not as efficient as slicing using a tuple (e.g. MAT[i, j]).

Finally, the nesting of the loops is largely unused in your code, and you could rewrite it with the two loops being separated (which is much more efficient):

def fool(arr):
    ii, jj = arr.shape
    for i in range(ii):
        arr[i, 0] = i
    for j in range(jj):
        arr[0, j] = j
    return arr

This is interesting because if we accelerate the code using Numba:

fool_nb = nb.jit(fool)
fool_nb.__name__ = 'fool_nb'

This results in the fastest approach:

funcs = foo, foo2, foon, fool, fool_nb


shape = 300, 400
for func in funcs:
    arr = np.zeros(shape)
    print(func.__name__)
    %timeit func(arr)
    print()

# foo
# 100 loops, best of 3: 6.53 ms per loop

# foo2
# 100000 loops, best of 3: 4.28 µs per loop

# foon
# 100000 loops, best of 3: 6.99 µs per loop

# fool
# 10000 loops, best of 3: 89.8 µs per loop

# fool_nb
# 1000000 loops, best of 3: 1.01 µs per loop

Upvotes: 0

Gerardo Zinno
Gerardo Zinno

Reputation: 1712

Since you are simply assigning the values 0,1,... to the first row and column there is no need for a double loop with if conditions inside.

Just assign the values to the first row:

MAT[0] = np.arange(len(MAT[0])

and to the first column:

MAT[:,0] = np.arange(len(MAT[:,0]))

Upvotes: 0

blue note
blue note

Reputation: 29071

You don't need a loop, just assign to slice

MAT[:, 0] = np.arange(x)
MAT[0, :] = np.arange(y)

Upvotes: 0

maciek97x
maciek97x

Reputation: 7072

You can simply set first row and first column like this

mat[:,0] = np.arange(0, mat.shape[0])
mat[0,:] = np.arange(0, mat.shape[1])

Example result

array([[0.        , 1.        , 2.        , 3.        , 4.        ],
       [1.        , 0.30487009, 0.97179858, 0.08143348, 0.99363866],
       [2.        , 0.69357714, 0.98421733, 0.42032313, 0.81041628]])

Upvotes: 1

Related Questions