[Numpy-discussion] Add New Functionality for Indexing Along an Axis to Numpy?

[Jonathan Taylor]

Please excuse the double post as I suspect people who may have
thoughts on the inclusion of such functionality in numpy were not
following the discussion due to the old subject.  I am perfectly happy
keeping this functionality locally but some of my colleagues have also
indicated that they have resorted to loops in the past to solve this
not uncommon use case so perhaps it would be helpful to more people if
it (or something similar?) was included in numpy?

Jonathan.

On Thu, Mar 10, 2011 at 12:00 PM, Jonathan Taylor
<jonathan.taylor at utoronto.ca> wrote:
> I see.
>
> Should functionality like this be included in numpy?
>
> Jon.
>
>
> On Tue, Mar 8, 2011 at 3:39 PM,  <josef.pktd at gmail.com> wrote:
>> On Tue, Mar 8, 2011 at 3:03 PM, Jonathan Taylor
>> <jonathan.taylor at 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 at scipy.org
>>> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>>>
>> _______________________________________________
>> NumPy-Discussion mailing list
>> NumPy-Discussion at scipy.org
>> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>>
>