[Numpy-discussion] Is this the optimal way to take index along a single axis?

josef.pktd@gmai... josef.pktd@gmai...
Tue Mar 8 14:39:18 CST 2011


On Tue, Mar 8, 2011 at 3:03 PM, Jonathan Taylor
<jonathan.taylor@utoronto.ca> wrote:
> I am wanting to use an array b to index into an array x with dimension
> bigger by 1 where the element of b indicates what value to extract
> along a certain direction.  For example, b = x.argmin(axis=1).
> Perhaps I want to use b to create x.min(axis=1) but also to index
> perhaps another array of the same size.
>
> I had a difficult time finding a way to do this with np.take easily
> and even with fancy indexing the resulting line is very complicated:
>
> In [322]: x.shape
> Out[322]: (2, 3, 4)
>
> In [323]: x.min(axis=1)
> Out[323]:
> array([[ 2,  1,  7,  4],
>       [ 8,  0, 15, 12]])
>
> In [324]: x[np.arange(x.shape[0])[:,np.newaxis,np.newaxis],
> idx[:,np.newaxis,:], np.arange(x.shape[2])]
> Out[324]:
> array([[[ 2,  1,  7,  4]],
>
>       [[ 8,  0, 15, 12]]])
>
> In any case I wrote myself my own function for doing this (below) and
> am wondering if this is the best way to do this or if there is
> something else in numpy that I should be using? -- I figure that this
> is a relatively common usecase.
>
> Thanks,
> Jon.
>
> def mytake(A, b, axis):
>    assert len(A.shape) == len(b.shape)+1
>
>    idx = []
>    for i in range(len(A.shape)):
>        if i == axis:
>            temp = b.copy()
>            shapey = list(temp.shape)
>            shapey.insert(i,1)
>        else:
>            temp = np.arange(A.shape[i])
>            shapey = [1]*len(b.shape)
>            shapey.insert(i,A.shape[i])
>        shapey = tuple(shapey)
>        temp = temp.reshape(shapey)
>        idx += [temp]
>
>    return A[tuple(idx)].squeeze()
>
>
> In [319]: util.mytake(x,x.argmin(axis=1), 1)
> Out[319]:
> array([[ 2,  1,  7,  4],
>       [ 8,  0, 15, 12]])
>
> In [320]: x.min(axis=1)
> Out[320]:
> array([[ 2,  1,  7,  4],
>       [ 8,  0, 15, 12]])

fewer lines but essentially the same thing and no shortcuts, I think

>>> x= np.random.randint(5, size=(2, 3, 4))
>>> x
array([[[3, 1, 0, 1],
        [4, 2, 2, 1],
        [2, 3, 2, 2]],

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

>>> idx = [np.arange(i) for i in x.shape]
>>> idx = list(np.ix_(*idx))
>>> idx[axis]=np.expand_dims(x.argmin(axis),axis)
>>> x[idx]
array([[[2, 1, 0, 1]],

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

>>> np.squeeze(x[idx])
array([[2, 1, 0, 1],
       [0, 1, 0, 1]])

>>> mytake(x,x.argmin(axis=1), 1)
array([[2, 1, 0, 1],
       [0, 1, 0, 1]])

Josef

> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion@scipy.org
> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>


More information about the NumPy-Discussion mailing list