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分别为元素方面的操作,比如我们这里.

替代替换此特定情况的最后一步

最后一步也可以通过按比例缩放来实现phi_0,mask因为else部分设置为0s-

out = mask*phi_0