NumPy - Vectorizing loops involving range iterators

Question

Is there any way to make this work without for loops?

import import numpy as np
import matplotlib.pyplot as plt    

L = 1
N = 255
dh = 2*L/N
dh2 = dh*dh

phi_0 = 1
c = int(N/2)
r_0 = L/2

arr = np.empty((N, N))

for i in range(N):

    for j in range(N):

            arr[i, j] = phi_0 if (i - c)**2 + (j - c)**2 < r_0**2/dh2 else 0


plt.imshow(arr)

I've tried calling function(x[None,:], y[:, None]), where:

function(i, j):

    return phi_0 if (i - c)**2 + (j - c)**2 < r_0**2/dh2 else 0

but it requires list .any or .all methods. I'm looking for specifically functionless method (without fromfunction and vectorization). Big thanks!

Answer 1

Vectorized solution using open grids

We could use two open range/grid arrays for N simulating the same behavior as the iterators -

I = np.arange(N)
mask = (I[:,None] - c)**2 + (I - c)**2 < r_0**2/dh2
out = np.where(mask,phi_0,0)

For a generic range on the two loops

For the generic case where we would iterate through two loops that extend till say M and N respectively, we could make use of np.ogrid to create those open grids and then use on the same lines -

I,J = np.ogrid[:M,:N]
mask = (I - c)**2 + (J - c)**2 < r_0**2/dh2

For a generic number of loops

For a generic number of loops, simply create as many variables as the number of loops. Hence, for three loops :

for i in range(M):
    for j in range(N):
        for k in range(P):

, we would have :

I,J,K = np.ogrid[:M,:N,:P]

, then use I,J,K instead of i,j,k respectively for element-wise operations like we have here.

---

Alternative to replace last step for this specific case

Last step could also be implemented with elementwise multiplication by scaling to phi_0 with mask as the else part is setting to 0s -

out = mask*phi_0

Answer 2

If you want to use numbers of row and column to calculation, a loop is necessary. You can use one loop. Numpy has ndenumerate attribute that iterate over your matrix.

def function(i, j):
    return phi_0 if (i - c)**2 + (j - c)**2 < r_0**2/dh2 else 0

for (i,j), value in np.ndenumerate(arr):
    arr[i, j] = function(i, j)