np.meshgrid的循环计算可以矢量化吗

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循环怎么办？

编辑历史：

1. 注意，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)

但是比你花更多的时间。