python - 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

但它需要 list .any 或 .all 方法。我正在寻找特定的无功能方法(没有 fromfunction 和矢量化)。非常感谢!

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 进行元素操作，就像我们这里的那样。

---

替代此特定案例的最后一步

最后一步也可以通过按元素乘法实现，方法是使用 mask 缩放到 phi_0，因为 else 部分设置为 0s -

out = mask*phi_0