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]]]])

但这样做开始变得有点麻烦，我不知道在更高维度的meshgrid输出中能否依赖这个顺序。numpy.mgrid给出了正确的顺序，但似乎不允许使用任意值，而我正需要这个功能。所以这归结为两个问题：

1）有没有更简单的方法，或者numpy里有没有我遗漏的函数，可以生成如上所述的坐标向量矩阵？

2）这种奇怪的顺序真的是我们对meshgrid的预期吗？有没有什么规范可以让我依赖这个点？

[编辑] 跟进Jaime的解决方案，这里有一个更通用的函数，可以更明确地构建它，供有兴趣的人参考：[编辑2，修复了一个bug，可能还有另一个，暂时没时间再花在这上面，这真的需要成为一个更常见的函数...]

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

我用的一个简单方法如下：

x,y = np.mgrid[-10:10:0.1, -10:10:0.1]
pos = np.empty(x.shape + (2,))
pos[:, :, 0] = x; pos[:, :, 1] = y
pos = np.reshape(pos, (x.shape[0]*x.shape[1], 2))

这里的 pos 是你需要的坐标数组。

Answer 2

这个问题最初是在六年前提出来的，但我发现自己一直在寻找解决这个问题的好方法，很多次都来到这里。最近我查看了Numpy的文档，整理出了一种直观且动态的方法来生成n维坐标，想在这里分享给那些还在寻找答案的人。

我们使用的是 numpy.ndindex()，它的作用是：

给定一个数组的形状，ndindex实例会遍历这个数组的N维索引。在每次遍历中，会返回一个索引的元组，最后一维会最先被遍历。

理解这个方法最好的方式是看一个例子：

In [100]: for index in np.ndindex(2,2,2):
              print(index)
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(0, 1, 1)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(1, 1, 1)

这正是我们想要的，那么我们如何将它转换成numpy数组格式或者列表呢？

如果我们想要将坐标作为列表，可以使用：

In [103]: coordList = [x for x in np.ndindex(2,2,2)]
In [104]: print(coordList)
[(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]

如果我们想要将坐标作为numpy数组，可以使用：

In [105]: coordArray = np.stack([x for x in np.ndindex(2,2,2)])
In [106]: print(coordArray)
[[0 0 0]
 [0 0 1]
 [0 1 0]
 [0 1 1]
 [1 0 0]
 [1 0 1]
 [1 1 0]
 [1 1 1]]

这种方法可以轻松适应不同的维度和大小，使用 numpy.reshape() 我们可以得到正好符合需求的格式：

In [117]: answer = np.stack([x for x in np.ndindex(2,2)]).reshape(2,2,2)
In [118]: print(answer)
[[[0 0]
  [0 1]]

 [[1 0]
  [1 1]]]

而且，这种方法也可以很容易地扩展到更高的维度：

In [120]: example = np.stack([x for x in np.ndindex(3,3,3)]).reshape(3,3,3,3)
In [121]: print(example)
[[[[0 0 0]
   [0 0 1]
   [0 0 2]]

  [[0 1 0]
   [0 1 1]
   [0 1 2]]

  [[0 2 0]
   [0 2 1]
   [0 2 2]]]


 [[[1 0 0]
   [1 0 1]
   [1 0 2]]

  [[1 1 0]
   [1 1 1]
   [1 1 2]]

  [[1 2 0]
   [1 2 1]
   [1 2 2]]]


 [[[2 0 0]
   [2 0 1]
   [2 0 2]]

  [[2 1 0]
   [2 1 1]
   [2 1 2]]

  [[2 2 0]
   [2 2 1]
   [2 2 2]]]]

Answer 3

给定一维坐标：

rows = np.arange(2)
cols = np.arange(3)

我原以为这样就能解决问题：

np.dstack((rows[:, None, None], cols[:, None]))

但显然，dstack 这类函数需要完全匹配的维度，它们不会自动调整尺寸，这让我觉得有点可惜。

所以这个替代方案虽然有点长，但明确总比模糊好，而且你可以把它们封装成一个小函数：

>>> coords = np.empty((len(rows), len(cols), 2), dtype=np.intp)
>>> coords[..., 0] = rows[:, None]
>>> coords[..., 1] = cols

>>> coords
array([[[0, 0],
        [0, 1],
        [0, 2]],

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

Answer 4

numpy库中的indices函数也可以用来实现这个功能，从名字上就能看出它的作用。

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

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

可以把它想象成一个2行3列的y坐标矩阵和一个2行3列的x坐标矩阵（y,x = np.indices((2,3))）。通过转置坐标轴，它可以变成Jaime所提到的形式：

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

它的功能和使用meshgrid的解决方案是一样的，只不过要用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那样的多维数组。通过在上面的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多，比如生成更一般的矩形网格而不是笛卡尔网格，所以在合适的情况下可以使用它们。在这种情况下，我会选择名字更具描述性的indices。

Answer 5

试试 np.meshgrid([0, 1], [0, 1], [0, 1], indexing="ij")。关于 meshgrid 的文档其实很明确地说明了，默认的 indexing="xy" 会产生一种奇怪的轴顺序，而非默认的 indexing="ij" 则不会，所以你可以去查看文档了解更多细节。（不过，他们没有很清楚地解释为什么会这样工作，真是可惜…）