Andrew Schwartz - 1 year ago 116

Python Question

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 Source

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