python - 坐标的 Numpy 矩阵

Question

我正在尝试获取坐标数组矩阵。这与 numpy.meshgrid 不同。例如，对于 2x2 尺寸，我想要 2x2x2 输出

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

作为一个 numpy 数组。这可能看起来和读起来更清晰的 2x2 元组矩阵:

[[(0,0),(0,1)],
 [(1,0),(1,1)]]

(除了我不认为你可以在 numpy 数组中有元组，这不是这里的重点)

这个简单的例子可以通过切换 numpy-meshgrid 的输出轴来完成(具体来说，将第一个轴移动到最后一个轴):

np.array(np.meshgrid([0,1],[0,1])).transpose([1,2,0])

这可以很容易地推广到任意维度，除了 meshgrid 的行为不符合我对超过 2 个输入的预期。具体来说，返回的矩阵具有沿轴以奇数顺序变化的坐标值:

In [627]: np.meshgrid([0,1],[0,1],[0,1])
Out[627]:
[array([[[0, 0],
        [1, 1]],

       [[0, 0],
        [1, 1]]]),
 array([[[0, 0],
        [0, 0]],

       [[1, 1],
        [1, 1]]]),
 array([[[0, 1],
        [0, 1]],

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

请注意，此输出的元素分别沿轴 1、0 和 2 变化。这将构建一个不正确的坐标矩阵；我需要输出按轴 0、1 和 2 的顺序变化。所以我可以做

In [642]: np.array(np.meshgrid([0,1],[0,1],[0,1])).swapaxes(1,2)
Out[642]:
array([[[[0, 0],
         [0, 0]],

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


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

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


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

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

但这开始变得非常棘手，我不知道我是否可以在更高维度的网格输出中依靠这个顺序。 numpy.mgrid 给出了正确的顺序，但似乎不允许我需要的任意值。所以这归结为两个问题:

1) 是否有更简洁的方法，也许我缺少 numpy 中的某些函数，它会生成所描述的坐标向量矩阵？2) 这种奇怪的顺序真的是我们对 meshgrid 的期望吗？在这一点上是否有我可以信赖的规范？

[编辑] 继 Jaime 的解决方案之后，这里有一个更通用的函数，可以为任何感兴趣的人更明确地构建它:[编辑 2，修复了一个错误，可能是另一个错误，不能花更多的时间在这个权利上现在，这确实需要一个更常见的功能...]

def build_coords(*vecs):
    coords = numpy.empty(map(len,vecs)+[len(vecs)])
    for ii in xrange(len(vecs)):
        s = np.hstack((len(vecs[ii]), np.ones(len(vecs)-ii-1)))
        v = vecs[ii].reshape(s)
        coords[...,ii] = v
    return coords

Answer 1

numpy 函数 indices也可以用于此效果，其功能也从其名称中一目了然。

>>> import numpy as np
>>> np.indices((2,3))
array([[[0, 0, 0],
        [1, 1, 1]],

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

它可以被认为是一个 2 x 3 的 y 坐标矩阵和一个 2 x 3 的 x 坐标矩阵 (y,x = np.indices((2,3)))。它可以通过转置轴重铸成 Jaime 提出的形式:

>>> np.indices((2,3)).transpose((1,2,0))

它在功能上等同于the meshgrid solution ，使用 indexing='ij'，但不需要您提供坐标数组，这在您有多个维度时会很有用。

>>> def f1(shape):
...     return np.array(np.meshgrid(*(np.arange(s) for s in shape), indexing='ij'))
...
>>> shape = (200, 31, 15, 4)
>>> np.all(f1(shape) == np.indices(shape))
True

在时间上，这些解决方案是相似的，当您考虑生成一维数组所需的时间时，meshgrid 在其上运行，但 meshgrid 返回一个列表(数组)，而不是像 indices 这样的 nd 数组。通过在上面的 f1 中添加对 np.array 的额外调用，indices 比 meshgrid 有明显的优势:

In [14]: %timeit f1(shape)
100 loops, best of 3: 14 ms per loop

In [15]: %timeit np.indices(shape)
100 loops, best of 3: 5.77 ms per loop

没有对 array 的额外调用:

In [16]: def f2(shape):
    return np.meshgrid(*(np.arange(s) for s in shape), indexing='ij')
   .....: 

In [17]: %timeit f2(shape)
100 loops, best of 3: 5.78 ms per loop

不过，在解释时间时始终要小心。这可能不会成为您解决任何问题的瓶颈。

无论如何，meshgrid 比indices 可以做更多的事情，比如生成一个更通用的 rectilinear grid而不是笛卡尔网格，所以在适当的时候使用它们。在这种情况下，我会使用命名方式更具描述性的 indices。