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)

我試過調用函數（x [None，：]，y [：，None]），其中：

function(i, j):

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

但它需要list .any或.all方法。 我正在尋找具體的無功能方法（沒有fromfunction和vectorization）。 十分感謝！

Answer 1

使用開放網格的矢量化解決方案

我們可以使用兩個開放的范圍/網格陣列N模擬相同的行為迭代器-

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分別用於像我們這里的元素操作。

---

替代替換此特定情況的最后一步

最后一步也可以通過元素乘法實現，通過使用mask縮放到phi_0 ，因為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，這是一個更多的工作（類型聲明，...）

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)

計時

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