Andrew Schwartz Andrew Schwartz - 25 days ago 6
Python Question

Numpy matrix of coordinates

I'm trying to get a matrix of coordinate-arrays. This is different from numpy.meshgrid. For example, for a 2x2 size I'd want the 2x2x2 output

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


as a numpy array. This probably looks and reads cleaner a 2x2 matrix of tuples:

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


(except I don't think you can have tuples in a numpy array, and it's not the point here)

This simple example can be done by switching the axes of numpy-meshgrid's output (specifically, moving the first axis to be last):

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


This could be easily generalized to arbitrary dimensions, except that meshgrid doesn't behave as I would expect for more than 2 inputs. Specifically, the returned matrices have coordinate values that vary along axes in an odd order:

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


Notice that the elements of this output vary along axes 1, 0, and 2, respectively. This will build an incorrect coordinate matrix; I would need the output to vary along axes 0, 1, and 2, in that order. So I could do

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


But this is starting to get really hacky and I don't know if I can count on this order in higher-dimension meshgrid outputs. numpy.mgrid gives the right order, but doesn't seem to allow arbitrary values, which I will need. So this boils down to two questions:

1) Is there a cleaner way, maybe some function in numpy I'm missing, that will generate a matrix of coordinate-vectors as described?
2) Is this odd ordering really what we expect from meshgrid? Is there a spec to this point that I can count on?

[EDIT] Following up on Jaime's solution, here's a more generalized function to build it a little more explicitly for anyone interested: [EDIT 2, fixed a bug, might be another, can't spend much more time on this right now, this really needs to be a more common function...]

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

Try np.meshgrid([0, 1], [0, 1], [0, 1], indexing="ij"). The meshgrid docs are actually pretty explicit about how the default indexing="xy" produces a funny axis ordering as compared to the non-default indexing="ij", so you can check that for more details. (They're not as clear on why it works this way, alas...)