NumPy-涉及范围迭代器的向量化循环

Question

有没有任何方法可以使此工作不带for循环？

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)

我尝试调用function（x [None ,:]，y [:, None]），其中：

function(i, j):

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

但是它需要列出.any或.all方法。我正在寻找专门的无功能方法（没有功能和向量化）。 非常感谢！

Answer 1

使用开放式网格的矢量化解决方案

我们可以为N使用两个 open 范围/网格数组，以模拟与迭代器相同的行为-

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

对于两个循环的通用范围

对于一般情况，我们将遍历两个循环直到分别说M和N，我们可以使用np.ogrid来创建那些开放式网格，然后在相同的行-

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

有关通用循环次数

对于一般数量的循环，只需创建与循环数一样多的变量即可。因此，对于三个循环：

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

，我们会：

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

，然后分别使用I,J,K而不是i,j,k进行像我们这里这样的元素操作。

---

在这种情况下替代最后一步

在phi_0部分设置为mask的情况下，通过将else缩放到0s，也可以通过元素乘法实现最后一步。

out = mask*phi_0

Answer 2

如果要使用行数和列数进行计算，则需要循环。 您可以使用一个循环。 Numpy具有在矩阵上迭代的ndenumerate属性。

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)

Answer 3

如果您的真正目标不是避免循环，而是获得良好的性能（在这种情况下为670倍加速），则一种简单的方法是使用编译器。在此示例中，我使用Numba，但您也可以使用Cython，这需要做更多的工作（类型声明，...）

示例

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)

时间

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