Numpy ndarray的动态轴索引

Question

我想在3D阵列的给定方向上获得2D切片,其中direction(或者将要从中提取切片的轴)由另一个变量给出.

假设idx3D阵列中的2D切片的索引以及direction获得该2D切片的轴,初始方法将是:

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

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

到目前为止我发现的更好的解决方案(动态地做),依赖于转置运算符:

# for 3 dimensions [0,1,2] and direction == 1 --> [1, 0, 2]
tr = [direction] + range(A.ndim)
del tr[direction+1]

return np.transpose(A, tr)[idx]

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

Answer 1

转置很便宜(按时间).有一些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())]