Numpy ndarray的动态轴索引

我想获得3D数组中给定方向上的2D切片，其中

方向

（或从中提取切片的轴）由另一个变量给定

假设

idx

3D数组中2D切片的索引，以及

方向

获取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来说，转置代码看起来几乎比3if/elif更糟糕。它对ND的泛化效果更好，与之相比，N越大，代码就越漂亮，但对于3D，情况完全相同。

转置便宜（时间方面）。有

numpy

函数使用它将操作轴（或多个轴）移动到已知位置-通常是形状列表的前端或末端

tensordot

是我想到的

其他函数构造索引元组。为了便于操作，它们可以从列表或数组开始，然后将其转换为应用程序的元组。比如说

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

np.沿轴应用。查看类似这样的函数的代码很有启发性

我想
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_沿_轴
构造一个
（0…，：，0…）
索引元组

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

要动态索引维度，可以使用，如下所示：

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

axis = 1
idx = 2

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

运行时比较
自然法（非动态）：

Swapax：

%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)]]

非常有趣的帖子谢谢！若我并没有错的话，transpose只是返回一个数组，其中的
步数
发生了变化，对吗？除非调用
.copy（）
，否则它不会对内存中的数据重新排序。根据两个
%timeit
测试，使用
a=np.random.randn（10001000）
和
B=a.T
，访问
a[：，idx]
中的随机列要比访问
B[idx]
中的相同数据慢。但是，如果我们完全索引
B[idx，：]
，访问
B
中的数据会变得更昂贵。是一个50纳秒（x1.25加速）的问题，所以不重要，但我很好奇。你知道这在numpy是怎么工作的吗？这些都是视图。它不是迭代数据或处理数据。因此，时间差异与解析表达式和创建新的数组对象（相同的数据指针）有关<无论我使用
B[100]
还是
A.T[100，：]
swapaxes
都不会返回空间一致的结果，code>标志和
数组接口看起来都是一样的。e、 g.对于
[100200300]
数组，如果我要索引最后一个轴（例如
数据[：，：，5]
），我希望返回
[100200]
数组。但是，如果我们执行
data.swapaxes（0,2）[5]
我们得到的是
[200,100]
数组，因为轴
0
现在位于末尾。它适用于
[0，1]
axis'，但仅适用于那些。没错，我忽略了这一点。选择相同的数据，但顺序可能会更改。
%timeit a[[idx if i==axis else slice(None) for i in range(a.ndim)]]