[Numpy-discussion] Fancy indexing for

Daran Rife drife@ucar....
Tue Sep 22 12:16:39 CDT 2009

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]

As you know, this is reasonably fast for modest-sized arrays,
but is far more expensive for large arrays.

Thanks in advance for your help.


Daran Rife

More information about the NumPy-Discussion mailing list