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