NumPy - Vectorizing loops involving range iterators

问题

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!

回答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

回答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)

回答3:

If your real goal is not to avoid loops, but to get good performance (in this case 670x speedup) a simple approach would be to use a compiler. In this example I use Numba but you can also use Cython, which is a bit more work to to (type declaration,...)

Exampe

import numpy as np
import numba as nb
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

@nb.njit()
def create_arr(N,phi_0,c,r_0,dh2):
  arr = np.empty((N, N))
  for i in range(N):
    for j in range(N):
      if (i - c)**2 + (j - c)**2 < r_0**2/dh2:
        arr[i,j]=phi_0
      else:
        arr[i,j]=0.
  return arr

arr=create_arr(N,phi_0,c,r_0,dh2)

Timings

#Pure Python:   58 ms
#Numba version: 0.086 ms (the first call takes longer and isn't included in the timings)

来源：https://stackoverflow.com/questions/52704777/numpy-vectorizing-loops-involving-range-iterators

标签

python

numpy

vectorization