Numpy ndarray的动态轴索引 [英] Dynamic axis indexing of Numpy ndarray

问题描述

我想在3D数组的给定方向上获取2D切片，其中方向（或从中提取切片的轴）另一个变量。

I want to obtain the 2D slice in a given direction of a 3D array where the direction (or the axis from where the slice is going to be extracted) is given by another variable.

假设 idx 3D数组中2D切片的索引，并且 direction 获取2D切片的轴，初始方法是：

Assuming idx the index of the 2D slice in a 3D array, and direction the axis in which obtain that 2D slice, the initial approach would be:

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

我很确定必须有办法这样做没有条件，或者至少，不是在原始的python中。 numpy有这个捷径吗？

I'm pretty sure there must be a way of doing this without doing conditionals, or at least, not in raw python. Does numpy have a shortcut for this?

我到目前为止找到的更好的解决方案（动态地做），依赖于转置运算符：

The better solution I've found so far (for doing it dynamically), relies in the transpose operator:

# 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则完全相同。

But I wonder if there is any better/easier/faster function for this, since for 3D the transpose code almost looks more awful than the 3 if/elif. It generalizes better for ND and the larger the N the more beautiful the code gets in comparison, but for 3D is quite the same.

推荐答案

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

Transpose is cheap (timewise). There are numpy functions that use it to move the operational axis (or axes) to a known location - usually the front or end of the shape list. tensordot is one that comes to mind.

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

Other functions construct an indexing tuple. They may start with a list or array for ease of manipulation, and then turn it into a tuple for application. For example

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

np.apply_along_axis 做类似的事情。看看像这样的函数的代码是有益的。

np.apply_along_axis does something like that. It's instructive to look at the code for functions like this.

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

I imagine the writers of the numpy functions worried most about whether it works robustly, secondly about speed, and lastly whether it looks pretty. You can bury all kinds of ugly code in a function!.

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 ... 。

where the previous code calculated newaxes_.. and newshape....

apply_along_axis 构造一个（0 ...，：，0。 ..） index tuple

apply_along_axis constructs a (0...,:,0...) index tuple

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

这篇关于Numpy ndarray的动态轴索引的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！