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