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

Question

我的代码中有一个部分，我通过一个循环，将应用于一片的for结果分配到一个占位符中，如下所示：np.meshgridmy_arrayres_arrayp = 360; q = 50; r = 50my_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

您可以将 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 倍！

Answer 2

你可以这样做：

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)

但是比你花的时间多。