np.meshgrid 的这个 for 循环计算可以向量化吗？

Question

我的代码中有一个部分，我在其中通过 for 循环，将应用于 my_array 切片的 np.meshgrid 的结果分配到占位符 res_array如下：

p = 360; q = 50; r = 50
my_array   =  np.random.rand(p, q, 2)
res_array  = np.zeros((p,q,r), dtype=np.float32)
for i in range(p):
   x, x_ = np.meshgrid(my_array[i,:,0], my_array[i,:,0])
   res_array[i] = (x_-x)

我的代码中还有一些类似的操作，我不得不运行我的代码数千次。因此速度成为一个问题。我在想，如果我可以不用 for 循环，我可以在这里节省一些时间，但是我自己做不到。

没有 for 循环怎么办？

编辑历史：

请注意，my_array[i,:,1] 未在上面显示的代码段中使用。同样，我在另一个代码块中使用 my_array[i,:,1] （此处未显示，但类似循环）。这就是我所说的“我的代码中还有一些类似的操作......”

Answer 1

你可以这样做：

my_array   =  np.random.rand(p, q, 2)
res_array = np.array([x[1]-x[0] for x in [np.meshgrid(my_array[i,:,0], my_array[i,:,0]) for i in range(p)]], np.float32).reshape(p,q,r)

但是比你花的时间多

Answer 2

您可以将 np.meshgrid 重写为更快的低级 numpy 操作：

p = 360; q = 50; r = 50
my_array   =  np.random.rand(p, q, 2)
res_array  = np.zeros((p,q,r), dtype=np.float32)
for i in range(p):
    x = my_array[None,i,:,0].repeat(q, axis=0).reshape(q, q)
    y = my_array[None,i,:,0].repeat(q, axis=1).reshape(q, q)
    res_array[i] = y - x

这段代码在我的机器上快了 2 倍。 Numba 的 @njit 可用于加速上述代码（同样快 3 倍），但使用它可以实现更高效的实现。这是：

@njit(parallel=True)
def fasterImpl(my_array, p, q, r):
    res_array = np.zeros((p, q, r))
    for i in prange(p):
        for j in range(q):
            for k in range(r):
                res_array[i,j,k] = my_array[i,j,0] - my_array[i,k,0]
    return res_array

p = 360; q = 50; r = 50
my_array   = np.random.rand(p, q, 2)
res_array  = fasterImpl(my_array, p, q, r)

这个最终实现比我机器上的原始实现快 29 倍！

np.meshgrid 的这个 for 循环计算可以向量化吗？

Can this for-loop calculation of np.meshgrid be vectorized?

python

arrays

performance

for-loop

vectorization