[Numpy-discussion] Extracting values from one array corresponding to argmax elements in another array

[Vincent Schut]

On 04/06/2010 03:22 PM, Ken Basye wrote:
> From: Vincent Schut <schut at sarvision.nl>
>> On 04/05/2010 06:06 PM, Keith Goodman wrote:
>>
>>> On Mon, Apr 5, 2010 at 8:44 AM, Ken Basye<kbasye1 at jhu.edu>   wrote:
>>>
snip
>>>>> b[a.argmax(axis=0), range(3)]
>>>>>
>>>      array([0, 4, 5])
>>>
>>
>> Which does not work anymore when your arrays become more-dimensional
>> (like in my case: 4 or more) and the axis you want to select on is not
>> the first/last one. If I recall correctly, I needed to construct the
>> full index arrays for the other dimensions too (like with ogrid I
>> think). So: create the ogrid, replace the one for the dimensions you
>> want the argmax selection to take place on with the argmax parameter,
>> and use those index arrays to index your b array.
>> I'd need to look up my source code to be more sure/precise. If anyone
>> would like me to, please let me know. If anyone knows a less elaborate
>> way, also please let us know! :-)
>>
> Hi Vincent,
> I'd like to see more about your solution. For my present purposes,
> Keith's solution was sufficient, but I'm still very interested in a
> solution that's independent of dimension and axis.
> Thanks (and thanks, Keith),
> Ken
>

I've tracked it down. I have created the following quick'n dirty function:

def selectFromDimension(data, selection, dimension):
     ogridStatement = "numpy.ogrid["
     slices = []
     for d in data.shape:
         slices.append("0:" + str(d))
     ogridStatement += ",".join(slices)
     ogridStatement += "]"
     grid = eval(ogridStatement)
     grid[dimension] = selection
     result = data[grid]
     return result

which you call with the array to select from, the result of 
argmin/max(axis=...), and the axis to select from (usually the same as 
the axis argument to the argmin/max)

It uses eval to be able to create the ogrid independent of 
dimensionality. There's probably a much cleaner way, but this worked for 
me :-)

Now I think about it, probably another way could be to transpose your 
axes such that the axis to select on becomes the first axis. Then just 
index with the result of argmin/max, and transpose back if necessary. 
However, I have not tried this and I might miss something obvious here...

Vincent.