Python NumPy:在矢量化赋值期间计算索引数组
我想将此NumPy操作矢量化:Python NumPy:在矢量化赋值期间计算索引数组,python,arrays,numpy,vectorization,Python,Arrays,Numpy,Vectorization,我想将此NumPy操作矢量化: for j in range(yt): for i in range(xt): y[j, i] = x[idx[j, i], j, i] 其中idx包含指向x切片的轴0索引。有什么简单的方法可以做到这一点吗?您可以使用: J, I = np.ogrid[:yt, :xt] x[idx, J, I] 以下是测试: import numpy as np yt, xt = 3, 5 x = np.random.rand(10, 6, 7)
for j in range(yt):
for i in range(xt):
y[j, i] = x[idx[j, i], j, i]
其中idx
包含指向x
切片的轴0索引。有什么简单的方法可以做到这一点吗?您可以使用:
J, I = np.ogrid[:yt, :xt]
x[idx, J, I]
以下是测试:
import numpy as np
yt, xt = 3, 5
x = np.random.rand(10, 6, 7)
y = np.zeros((yt, xt))
idx = np.random.randint(0, 10, (yt, xt))
for j in range(yt):
for i in range(xt):
y[j, i] = x[idx[j, i], j, i]
J, I = np.ogrid[:yt, :xt]
np.all(x[idx, J, I] == y)
这里有一种方法是使用-
运行时测试和验证输出 本节比较了本篇文章中提出的方法和绩效评估,并验证了结果 函数定义-
def original_app(x,idx):
_,yt,xt = x.shape
y = np.zeros((yt,xt))
for j in range(yt):
for i in range(xt):
y[j, i] = x[idx[j, i], j, i]
return y
def ogrid_based(x,idx):
_,yt,xt = x.shape
J, I = np.ogrid[:yt, :xt]
return x[idx, J, I]
def reshape_based(x,idx):
zt,yt,xt = x.shape
return x.reshape(zt,-1)[idx.ravel(),np.arange(yt*xt)].reshape(-1,xt)
设置输入-
In [56]: # Inputs
...: zt,yt,xt = 100,100,100
...: x = np.random.rand(zt,yt,xt)
...: idx = np.random.randint(0,zt,(yt,xt))
...:
验证输出-
In [57]: np.allclose(original_app(x,idx),ogrid_based(x,idx))
Out[57]: True
In [58]: np.allclose(original_app(x,idx),reshape_based(x,idx))
Out[58]: True
时间安排-
In [68]: %timeit original_app(x,idx)
100 loops, best of 3: 6.97 ms per loop
In [69]: %timeit ogrid_based(x,idx)
1000 loops, best of 3: 391 µs per loop
In [70]: %timeit reshape_based(x,idx)
1000 loops, best of 3: 230 µs per loop
In [68]: %timeit original_app(x,idx)
100 loops, best of 3: 6.97 ms per loop
In [69]: %timeit ogrid_based(x,idx)
1000 loops, best of 3: 391 µs per loop
In [70]: %timeit reshape_based(x,idx)
1000 loops, best of 3: 230 µs per loop