Numpy slice with array as index

Question

I am trying to extract the full set of indices into an N-dimensional cube, and it seems like np.mgrid is just what I need for that. For example, np.mgrid[0:4,0:4] produces a 4 by 4 matrix containing all the indices into an array of the same shape.

The problem is that I want to do this in an arbitrary number of dimensions, based on the shape of another array. I.e. if I have an array a of arbitrary dimension, I want to do something like idx = np.mgrid[0:a.shape], but that syntax is not allowed.

Is it possible to construct the slice I need for np.mgrid to work? Or is there perhaps some other, elegant way of doing this? The following expression does what I need, but it is rather complicated and probably not very efficient:

np.reshape(np.array(list(np.ndindex(a.shape))),list(a.shape)+[len(a.shape)])

DSM · Accepted Answer · 2012-11-07 15:53:09Z

2

I usually use np.indices:

>>> a = np.arange(2*3).reshape(2,3)
>>> np.mgrid[:2, :3]
array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])
>>> np.indices(a.shape)
array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])
>>> a = np.arange(2*3*5).reshape(2,3,5)
>>> (np.mgrid[:2, :3, :5] == np.indices(a.shape)).all()
True

answered Nov 7, 2012 at 15:53

DSM

355k67 gold badges606 silver badges504 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

amaurea Over a year ago

Thanks, this is even better than the mgrid solution.

NPE · Accepted Answer · 2012-11-07 16:00:17Z

1

I believe the following does what you're asking:

>>> a = np.random.random((1, 2, 3))
>>> np.mgrid[map(slice, a.shape)]
array([[[[0, 0, 0],
         [0, 0, 0]]],


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


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

It produces exactly the same result as np.mgrid[0:1,0:2,0:3]except that it uses a's shape instead of hard-coded dimensions.

edited Nov 7, 2012 at 16:00

answered Nov 7, 2012 at 15:54

NPE

503k114 gold badges970 silver badges1k bronze badges

Collectives™ on Stack Overflow

Numpy slice with array as index

2 Answers 2

1 Comment

Comments

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Comments

Related