[Numpy-discussion] using numpy.argmax to index into another array

[Sebastian Berg]

Hey,

On Wed, 2012-10-31 at 20:22 -0400, David Warde-Farley wrote:
> On Wed, Oct 31, 2012 at 7:23 PM, Moroney, Catherine M (388D)
> <Catherine.M.Moroney at jpl.nasa.gov> wrote:
> > Hello Everybody,
> >
> > I have the following problem that I would be interested in finding an easy/elegant solution to.
> > I've got it working, but my solution is exceedingly clunky and I'm sure that there must be a
> > better way.
> >
> > Here is the (boiled-down) problem:
> >
> > I have 2 different 3-d array, shaped (4,4,4), "a" and "b"
> >
> > I can find the max values and location of those max values in "a" in the 0-th dimension
> > using max and argmax resulting in a 4x4 matrix.  So far, very easy.
> >
> > I then want to find the values in "b" that correspond to the maximum values in a.
> > This is where I got stuck.
> >
> > Below find the sample code I used (pretty clumsy stuff ...)
> > Can somebody show a better (less clumsy) way of finding bmax
> > in the code excerpt below?
> 

Its a bit off-label usage (as the documentation says), so maybe its
slower for larger arrays, but as long as you choose along axis 0, you
could also use:

np.choose(amax, b)

> Hi Catherine,
> 
> It's not a huge improvement, but one simple thing is that you can
> generate those index arrays easily and avoid the pesky reshaping at
> the end like so:
> 
> In [27]: idx2, idx3 = numpy.mgrid[0:4, 0:4]
> 
> In [28]: b[amax, idx2, idx3]
> Out[28]:
> array([[12,  1, 13,  3],
>        [ 8, 11,  9, 10],
>        [11, 11,  1, 10],
>        [ 5,  7,  9,  5]])
> 
> In [29]: b[amax, idx2, idx3] == bmax
> Out[29]:
> array([[ True,  True,  True,  True],
>        [ True,  True,  True,  True],
>        [ True,  True,  True,  True],
>        [ True,  True,  True,  True]], dtype=bool)
> 
> numpy.ogrid would work here too, and will use a bit less memory. mgrid
> and ogrid are special objects from the numpy.lib.index_tricks module
> that generate, respectively, a "mesh grid" and an "open mesh grid"
> (where the unchanging dimensions are of length 1) when indexed like
> so. In the latter case, broadcasting takes effect when you index with
> them.
> 
> Cheers,
> 
> David
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion at scipy.org
> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>