[Numpy-discussion] Extracting values from one array corresponding to argmax elements in another array
Wed Apr 7 07:24:43 CDT 2010
On 04/06/2010 03:22 PM, Ken Basye wrote:
> From: Vincent Schut <email@example.com>
>> On 04/05/2010 06:06 PM, Keith Goodman wrote:
>>> On Mon, Apr 5, 2010 at 8:44 AM, Ken Basye<firstname.lastname@example.org> wrote:
>>>>> 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),
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]
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
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...
More information about the NumPy-Discussion