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

Jonathan Taylor jonathan.taylor@utoronto...
Tue Mar 8 14:03:28 CST 2011


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


More information about the NumPy-Discussion mailing list