5
期间Evaulate索引阵列我想向量化这个NumPy的操作:NumPy的:矢量分配
for j in range(yt):
for i in range(xt):
y[j, i] = x[idx[j, i], j, i]
其中idx
包含轴-0索引到一个x
切片。有一些简单的方法来做到这一点?
期间Evaulate索引阵列我想向量化这个NumPy的操作:NumPy的:矢量分配
for j in range(yt):
for i in range(xt):
y[j, i] = x[idx[j, i], j, i]
其中idx
包含轴-0索引到一个x
切片。有一些简单的方法来做到这一点?
您可以使用:
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)
下面是使用linear indexing
一个方法 -
zt,yt,xt = x.shape
out = x.reshape(zt,-1)[idx.ravel(),np.arange(yt*xt)].reshape(-1,xt)
运行测试&验证输出
本节比较本文中的建议方法和other orgid based solution
的性能,并验证输出。
函数定义 -
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