为多维 numpy 数组结合切片和广播索引答案

【问题标题】：Combining slicing and broadcasted indexing for multi-dimensional numpy arrays为多维 numpy 数组结合切片和广播索引
【发布时间】：2015-02-02 12:49:37
【问题描述】：

我有一个 ND numpy 数组（比如说 3x3x3），我想从中提取一个子数组，结合切片和索引数组。例如：

import numpy as np  
A = np.arange(3*3*3).reshape((3,3,3))
i0, i1, i2 = ([0,1], [0,1,2], [0,2])
ind1 = j0, j1, j2 = np.ix_(i0, i1, i2)
ind2 = (j0, slice(None), j2)
B1 = A[ind1]
B2 = A[ind2]

我希望 B1 == B2，但实际上，形状不同

>>> B1.shape
(2, 3, 2)
>>> B2.shape
(2, 1, 2, 3)
>>> B1
array([[[ 0,  2],
        [ 3,  5],
        [ 6,  8]],

       [[ 9, 11],
        [12, 14],
        [15, 17]]])
>>> B2
array([[[[ 0,  3,  6],
         [ 2,  5,  8]]],

       [[[ 9, 12, 15],
         [11, 14, 17]]]])

有人明白为什么吗？知道如何通过仅操作“A”和“ind2”对象来获得“B1”吗？目标是它适用于任何 nD 数组，并且我不必寻找我想要完全保留的维度的形状（希望我足够清楚:)）。谢谢！！
---编辑---
更清楚地说，我想要一个“有趣”的功能，这样

A[fun(ind2)] == B1

【问题讨论】：

我认为这是你不想要的：ind2 = (i0, np.arange(A.shape[1]).reshape(-1,1), i2)
B2 = select(A,"0,1;:;0,2") 够用吗？如果您愿意，我可以发布答案。
理想情况下，我更喜欢保留符号 B2 = A[ind3]，否则，这意味着我必须替换一个很长的程序中的每个索引......但不幸的是，我不确定我的问题有一个简单的解决方案。

标签： python arrays numpy multidimensional-array slice

【解决方案1】：

这是我越接近你的规格，我无法设计出一个可以在不知道A（或者更准确地说，它的形状......）的情况下计算正确索引的解决方案。

import numpy as np  

def index(A, s):
    ind = []
    groups = s.split(';')
    for i, group in enumerate(groups):
        if group == ":":
            ind.append(range(A.shape[i]))
        else:
            ind.append([int(n) for n in group.split(',')])
    return np.ix_(*ind)

A = np.arange(3*3*3).reshape((3,3,3))

ind2 = index(A,"0,1;:;0,2")
print A[ind2]

更短的版本

def index2(A,s):return np.ix_(*[range(A.shape[i])if g==":"else[int(n)for n in g.split(',')]for i,g in enumerate(s.split(';'))])

ind3 = index2(A,"0,1;:;0,2")
print A[ind3]

【讨论】：

感谢您的关心和回答。我想在不知道 A 的形状的情况下，确实没有办法得到我想要的东西。
在我看来，我认为这是 numpy 的一个（小）缺陷（在许多许多很棒的东西中），我们无法做到，例如 A[[1,5],:,[1, 2,4]] 对数组进行切片，就像我们在 matlab 中所做的那样。

【解决方案2】：

ind1 的索引子空间是 (2,),(3,),(2,)，结果 B 是 (2,3,2)。这是高级索引的一个简单案例。

ind2 是（高级）部分索引的一种情况。有 2 个索引数组和 1 个切片。高级索引文档指出：

如果索引子空间是分开的（通过切片对象），那么广播的索引空间是第一个，然后是x的切片子空间。

在这种情况下，高级索引构造了一个 (2,2) 数组（从第一个和第三个索引开始），并在末尾附加切片维度，从而生成一个 (2,2,3) 数组。

我在https://stackoverflow.com/a/27097133/901925中更详细地解释了推理

修复像ind2 这样的元组的一种方法是将每个切片扩展为一个数组。我最近在np.insert 看到了这个。

np.arange(*ind2[1].indices(3))

将: 扩展为[0,1,2]。但替换必须具有正确的形状。

ind=list(ind2)
ind[1]=np.arange(*ind2[1].indices(3)).reshape(1,-1,1)
A[ind]

我将省略确定哪个术语是切片、其维度以及相关重塑的细节。目标是重现i1。

如果索引是由 ix_ 以外的其他东西生成的，那么重塑这个切片可能会更加困难。例如

A[np.array([0,1])[None,:,None],:,np.array([0,2])[None,None,:]] # (1,2,2,3)
A[np.array([0,1])[None,:,None],np.array([0,1,2])[:,None,None],np.array([0,2])[None,None,:]]
# (3,2,2)

扩展的切片必须与广播中的其他数组兼容。

索引后交换轴是另一种选择。不过，逻辑可能更复杂。但在某些情况下，转置实际上可能更简单：

A[np.array([0,1])[:,None],:,np.array([0,2])[None,:]].transpose(2,0,1)
# (3,2,2)
A[np.array([0,1])[:,None],:,np.array([0,2])[None,:]].transpose(0,2,1)
# (2, 3, 2)

【讨论】：

感谢您的回答。我确实注意到最后添加了切片维度。但是当切片维度是第一维度时有一个例外。例如，回到我的例子： ind3 = (slice(None),[[0],[1]],[[0,1]]) B3 = A[ind3] 那么 B3.shape 是 (3, 2,2)
在您的解决方案中，我遇到的问题是您需要知道切片维度的大小（此处为 3）。在这里，我希望保留符号 B=A[ind]，其中 ind 的定义独立于 A 的形状（或者更准确地说，独立于被切片的维度的大小）。
我知道的两个解决方案（但我并不十分满意）是： - 正如@gboffi 所建议的，将每次出现的索引 B=A[ind] 替换为函数 select (A,ind) 在一个很长的程序中 - 而不是声明 ndarrays，而是编写从 ndarray 继承的子类 MyNdarray，并覆盖 getitem 方法，以便 A[ind] 按我的意愿工作是的。
是的，我的解决方案需要A[fun(ind2, A.shape)]。
AxisConcatenator in np.lib.index_tricks 是具有自定义（相当长）__getitem__ 方法的类的示例。它是方便函数r_ 和s_ 的基础。

【解决方案3】：

在这种使用ix_ 的受限索引情况下，可以在连续的步骤中进行索引。

A[ind1]

与

相同

A[i1][:,i2][:,:,i3]

由于i2 是完整范围，

A[i1][...,i3]

如果您只有ind2 可用

A[ind2[0].flatten()][[ind2[2].flatten()]

在更一般的情况下，您必须知道j0,j1,j2 是如何相互广播的，但是当它们由ix_ 生成时，关系就很简单了。

我可以想象分配A1 = A[i1] 会很方便的情况，然后是涉及A1 的各种操作，包括但不限于A1[...,i3]。您必须注意A1 何时是视图，何时是副本。

另一个索引工具是take：

A.take(i0,axis=0).take(i2,axis=2)

【讨论】：