Numpy ndarray的动态轴索引

时间:2015-06-27 23:35:45

标签: numpy indexing

我想在3D数组的给定方向上获取2D切片,其中direction(或将从中提取切片的轴)由另一个变量给出。

假设{1}是3D数组中2D切片的索引,idx获取该2D切片的轴,则初始方法为:

direction

我非常确定必须有一种方法可以在不执行条件的情况下执行此操作,或者至少不是在原始python中。 numpy有这个捷径吗?

我到目前为止找到的更好的解决方案(动态执行)依赖于转置运算符:

if direction == 0:
    return A[idx, :, :]
elif direction == 1:
    return A[:, idx, :]
else:
    return A[:, :, idx]

但我想知道是否有更好/更容易/更快的功能,因为对于3D,转置代码几乎看起来比3 if / elif更糟糕。它对ND的推广更为普遍,N越大,代码相比就越漂亮,但对于3D来说却是完全相同的。

2 个答案:

答案 0 :(得分:5)

转置很便宜(按时间)。有numpy个函数用它来将操作轴(或轴)移动到已知位置 - 通常是形状列表的前端或末端。 tensordot是我想到的。

其他函数构造索引元组。它们可以以列表或数组开头以便于操作,然后将其转换为应用程序的元组。例如

I = [slice(None)]*A.ndim
I[axis] = idx
A[tuple(I)]

np.apply_along_axis做了类似的事情。查看类似函数的代码是有益的。

我认为numpy函数的编写者最担心的是它是否有效,其次是关于速度,最后是否看起来很漂亮。你可以把各种丑陋的代码埋在一个函数中!。

tensordot

结尾
at = a.transpose(newaxes_a).reshape(newshape_a)
bt = b.transpose(newaxes_b).reshape(newshape_b)
res = dot(at, bt)
return res.reshape(olda + oldb)

前面的代码计算newaxes_..newshape...

apply_along_axis构造一个(0...,:,0...)索引元组

i = zeros(nd, 'O')
i[axis] = slice(None, None)
i.put(indlist, ind)
....arr[tuple(i.tolist())]

答案 1 :(得分:-1)

要动态索引维度,您可以使用swapaxes,如下所示:

a = np.arange(7 * 8 * 9).reshape((7, 8, 9))

axis = 1
idx = 2

np.swapaxes(a, 0, axis)[idx]

运行时比较

自然方法(非动态):

%timeit a[:, idx, :]
300 ns ± 1.58 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

swapaxes:

%timeit np.swapaxes(a, 0, axis)[idx]
752 ns ± 4.54 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

带列表理解的索引:

%timeit a[[idx if i==axis else slice(None) for i in range(a.ndim)]]