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

Question

我有一个numpy nd数组.我的任务的简化版本是沿每个轴取一个向量.为了显示:

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

但是,我不一定知道x将具有的维数.因此,挑战在于如何将冒号:索引运算符放在可变位置.这种语法的示例如下:

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

要么

y = indexer.index(x,ind)

对于某些模块索引器.我可以写它,但我觉得这必须已经解决,我不能成为唯一一个想要这样做的人.例如,在MATLAB中,您可以使用subsref()函数执行此操作.

python/numpy/other模块中是否存在任何此类构造？

Answer 1

正如numpy关于索引的文档所示,您可以使用slice内置函数和元组串联来创建变量索引.

事实上:,下标只是文字的字面符号slice.

特别:是相当于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地做任何事情.要记住的重要一点是,它最终应该返回一个tuple包含整数或slices的东西.