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

[David Warde-Farley]

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?

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