Numpy切片与数组作为索引

Question

我试图将全套索引提取到一个N维立方体中，似乎np.mgrid正是我需要的。 例如， np.mgrid[0:4,0:4]生成一个4乘4的矩阵，其中包含所有索引到相同形状的数组中。

问题是我想根据另一个数组的形状在任意数量的维度中执行此操作。 也就是说，如果我有一个数组a任意尺寸的，我想这样做idx = np.mgrid[0:a.shape]但语法是不允许的。

是否有可能构建我需要np.mgrid工作的切片？ 或者是否有其他一些优雅的方式呢？ 以下表达式可以满足我的需要，但它相当复杂，可能效率不高：

np.reshape(np.array(list(np.ndindex(a.shape))),list(a.shape)+[len(a.shape)])

Answer 1

我通常使用np.indices ：

>>> a = np.arange(2*3).reshape(2,3)
>>> np.mgrid[:2, :3]
array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])
>>> np.indices(a.shape)
array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])
>>> a = np.arange(2*3*5).reshape(2,3,5)
>>> (np.mgrid[:2, :3, :5] == np.indices(a.shape)).all()
True

Answer 2

我相信以下内容符合您的要求：

>>> a = np.random.random((1, 2, 3))
>>> np.mgrid[map(slice, a.shape)]
array([[[[0, 0, 0],
         [0, 0, 0]]],


       [[[0, 0, 0],
         [1, 1, 1]]],


       [[[0, 1, 2],
         [0, 1, 2]]]])

它产生与np.mgrid[0:1,0:2,0:3]完全相同的结果，除了它使用a是形状而不是硬编码的尺寸。