Python：在任意维度中使用冒号运算符索引数组

Question

我有一个numpy的多维数组。我的任务简单来说就是从每个轴上取一个向量。为了说明这个问题：

import numpy
x = numpy.array(range(24)).reshape((2,3,4))
x0 = x[0,0,:]
x1 = x[0,:,0]
x2 = x[:,0,0]

但是我并不一定知道这个数组有多少个维度。所以挑战在于如何把冒号 : 这个索引操作符放在一个可变的位置。下面是这种语法可能的样子：

n = x.ndim
ind = list(np.zeros(n))
dim = 0
ind[dim] = ':'
y = x[ind]

或者

y = indexer.index(x,ind)

用于某个模块的索引器。我可以自己写这个，但我觉得这个问题应该已经有人解决过了，我不可能是唯一想这样做的人。例如，在MATLAB中，你可以用subsref()函数来实现。

在python / numpy / 其他模块中，有没有类似的构造？

Answer 1

创建一个叫做 subref 的函数（就像你在 MATLAB 中写的那样）

def subref(arr, ind): 
    return arr[tuple([slice(None) if i == ind else 0 for i in range(arr.ndim)])]

Answer 2

你可以用代码来组合一个字符串，选择你想要的维度，然后用eval来执行这个字符串里的代码。

一个开始的例子是：

n = 2
sel = "0,"*(n-1) + ":"
eval('x[' + sel + ']')

如果你想得到你想要的结果，事情会稍微复杂一点（但也不算太难）：

ind = 2
n = 3
sel = "".join([ ("0" if i != ind else ":") + ("," if i < n-1 else "") for i in xrange(n)])
eval('x[' + sel + ']')

这和动态SQL使用的策略是一样的。

Answer 3

根据 numpy 的文档，关于 索引 的内容，你可以使用 slice 这个内置函数和元组拼接来创建可变的索引。

其实，子脚本中的 : 只是 slice 的一种简单表示方式。

特别地，: 相当于 slice(None)（而 slice(None) 本身又等于 slice(None, None, None)，其中的参数分别是 start、stop 和 step）。

举个例子：

a[(0,) * N + (slice(None),)]

这和下面的内容是一样的：

a[0, 0, ..., 0, :]   # with N zeros

注意，: 这种表示法只能直接在子脚本中使用。例如，下面这个写法是错误的：

In [10]: a[(0,0,:)]
  File "<ipython-input-10-f41b33bd742f>", line 1
    a[(0,0,:)]
           ^
SyntaxError: invalid syntax

为了从任意维度的数组中提取切片，你可以写一个简单的函数，比如：

def make_index(num_dimension, slice_pos):
    zeros = [0] * num_dimension
    zeros[slice_pos] = slice(None)
    return tuple(zeros)

然后可以这样使用它：

In [3]: a = np.array(range(24)).reshape((2, 3, 4))

In [4]: a[make_index(3, 2)]
Out[4]: array([0, 1, 2, 3])

In [5]: a[make_index(3, 1)]
Out[5]: array([0, 4, 8])

In [6]: a[make_index(3, 0)]
Out[6]: array([ 0, 12])

你可以把 make_index 函数扩展到做任何事情。重要的是要记住，最后它应该返回一个包含整数或 slice 的 tuple。