[Numpy-discussion] Fancy indexing for

[Robert Kern]

On Tue, Sep 22, 2009 at 12:16, Daran Rife <drife at ucar.edu> wrote:
> Hello list,
>
> This didn't seem to get through last time round, and my
> first version was poorly written.
>
> I have a rather pedestrian question about fancy indexing
> for multi-dimensional arrays.
>
> Suppose I have two 3-D arrays, one named "A" and the other "B",
> where both arrays have identical dimensions of time, longitude,
> and latitude. I wish to use data from A to conditionally select
> values from array B. Specifically, I first find the time where
> the values at each point in A are at their maximum. This is
> accomplished with:
>
>  >>> tmax_idx = np.argsort(A, axis=0)
>
> I now wish to use this tmax_idx array to conditionally select
> the values from B. In essence, I want to pick values from B for
> times where the values at A are at their max. Can this be done
> with fancy indexing? Or is there a smarter way to do this? I've
> certainly done this sort of selection before, but the index
> selection array is 1D. I've carefully studied the excellent
> indexing documentation and examples on-line, but can't sort out
> whether what I want to do is even possible, without doing the
> brute force looping method, similar to:
>
> max_B = np.zeros((nlon, nlat), dtype=np.float32)
>
> for i in xrange(nlon):
>    for j in xrange(nlat):
>        max_B[i,j] = B[tmax_idx[i,j],i,j]

All of the index arrays need to be broadcastable to the same shape.
Thus, you want the "i" index array to be a column vector.

max_B = B[tmax_idx, np.arange(nlon)[:,np.newaxis], np.arange(nlat)]

-- 
Robert Kern

"I have come to believe that the whole world is an enigma, a harmless
enigma that is made terrible by our own mad attempt to interpret it as
though it had an underlying truth."
  -- Umberto Eco