[Numpy-discussion] Is this the optimal way to take index along a single axis?
Jonathan Taylor
jonathan.taylor at utoronto.ca
Thu Mar 10 12:00:05 EST 2011
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
>
More information about the NumPy-Discussion
mailing list