[英]NumPy: Evaulate index array during vectorized assignment
I would like to vectorize this NumPy operation: 我想对这个NumPy操作进行矢量化:
for j in range(yt):
for i in range(xt):
y[j, i] = x[idx[j, i], j, i]
where idx
contains axis-0 index to an x
slice. 其中
idx
包含x
切片的axis-0索引。 Is there some simple way to do this? 有一些简单的方法可以做到这一点吗?
You can use: 您可以使用:
J, I = np.ogrid[:yt, :xt]
x[idx, J, I]
Here is the test: 这是测试:
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)
Here's one approach using linear indexing
- 这是使用
linear indexing
的一种方法 -
zt,yt,xt = x.shape
out = x.reshape(zt,-1)[idx.ravel(),np.arange(yt*xt)].reshape(-1,xt)
Runtime tests & verify output 运行时测试并验证输出
This section compares the proposed approach in this post and the other orgid based solution
on performance and also verifies the outputs. 本节比较了本文中提出的方法和
other orgid based solution
性能other orgid based solution
,并验证了输出。
Function definitions - 功能定义 -
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)
Setup inputs - 设置输入 -
In [56]: # Inputs
...: zt,yt,xt = 100,100,100
...: x = np.random.rand(zt,yt,xt)
...: idx = np.random.randint(0,zt,(yt,xt))
...:
Verify outputs - 验证输出 -
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
Timings - 计时 -
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
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.