将 numpy 花式索引与省略号混合

Question

如何在 numpy 中将花式索引与 ... 相结合？

考虑一些测试数据：

# Some test data of high dimensionality 
shape = (3,5,2,8)
data = np.arange(np.prod(shape)).reshape(shape)

也许我有一个索引元组。我无法在索引期间解压缩这些值：

# Arbitrary indices
indices = (0, 4)
# Try unpacking
data[*indices]
# >>> SyntaxError: invalid syntax

不用担心，如果我单独传递索引或将索引作为元组传递，我会得到相同的结果 — 很好！

np.all(data[indices[0], indices[1]] == data[indices])
# >>> True

所以不需要拆包，对吧？相关问题似乎暗示如此。

Python's indexing already has built-in support for tuples in general [...] there is no need for iterator unpacking in indexing. Simply convert your iterator to a tuple

乍一看，它甚至似乎可以与奇特的索引一起使用：

# New arbitrary indices
new_indices = (np.array([0,1]), np.array([4,4]))
np.all(data[new_indices] == data[new_indices[0], new_indices[1]])
# >>> True

但是如果我尝试使用 ellipsis 它就会崩溃。例如，如果我想在尾随轴上进行花式索引：

# Passing unpacked values produces the desired result
data[..., new_indices[0], new_indices[1]] # shape=(3,5,2)
# Passing as a tuple does NOT produce the desired result
data[..., new_indices] # shape=(3,5,2,2,2)

当然如果我的数据是变维的，我就不能写..., new_indices[0], new_indices[1]。所以看来在这种情况下实际上需要拆包。

好吧，也许我可以这样做：

data.T[indices[::-1]].T.shape # shape=(3,5,2)

感觉有点老套，但确实有效。除非我们想在引导轴和尾随轴上使用 ... 之间的奇特索引，否则我们似乎真的不走运：

# Some test data of even higher dimensionality 
data_shape = (3,5,2,8,12,25)
data = np.arange(np.prod(data_shape)).reshape(data_shape)

index_A = (np.array([0,0,1]), np.array([0,2,1]))
index_B = (np.array([0,0,0]), np.array([1,1,1]))

# Desired result, shape=(3,2)
data[index_A[0], index_A[1], ..., index_B[0], index_B[1]]
# Without unpacking, shape=(2,3,5,2,8)
data[index_A, ..., index_B]

也许这是一个边缘案例，但似乎应该有办法做到这一点。怎么样？

P.S。无论如何，为什么在这种情况下可迭代解包会引发 SyntaxError？

Answer 1

tuple 加入就可以了：

In [142]: data[(...,)+indices]
Out[142]: 
array([[  4,  20,  36,  52,  68],
       [ 84, 100, 116, 132, 148],
       [164, 180, 196, 212, 228]])
In [143]: data[...,0,4]
Out[143]: 
array([[  4,  20,  36,  52,  68],
       [ 84, 100, 116, 132, 148],
       [164, 180, 196, 212, 228]])

In [144]: (...,)+indices
Out[144]: (Ellipsis, 0, 4)

在元组内解包也可以：

data[(...,*indices)]

将 numpy 花式索引与省略号混合

Mixing numpy fancy indexing with ellipsis

python

arrays

indexing

numpy

ellipsis