NumPy：在矢量化赋值期间评估索引数组答案

【问题标题】：NumPy: Evaulate index array during vectorized assignmentNumPy：在矢量化赋值期间评估索引数组
【发布时间】：2014-08-01 00:13:57
【问题描述】：

我想对这个 NumPy 操作进行矢量化：

for j in range(yt):
    for i in range(xt):
        y[j, i] = x[idx[j, i], j, i]

其中idx 包含指向x 切片的轴0 索引。有没有一些简单的方法可以做到这一点？

【问题讨论】：

标签： python arrays numpy vectorization

【解决方案1】：

你可以使用：

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)

【讨论】：

【解决方案2】：

这是使用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

【讨论】：