Skip to content

How to select the values for leading (max) and subleading (2nd max) varlues of an ak array in a memory efficient way? #3055

Answered by agoose77
green-cabbage asked this question in Q&A
Discussion options

You must be logged in to vote

My understanding of what we want to do here is something like np.argpartition, to extract the top-N largest values (in this case, 2), followed by a simple slice. If you always want both collections, it would be easier to perform a ragged-index followed by slice:

import awkward as ak

ix_jet_pt_2_largest = ak.argpartition(good_jet_pts, axis=1)
ordered_2_jets = good_jet_pts[ix_jet_pt_2_largest]

leading_jet = ak.firsts(ordered_2_jets[:, :1])
subleading_jet = ak.firsts(ordered_2_jets[:, 1:])

Although a pad_none might be faster at this point:

import awkward as ak

ix_jet_pt_2_largest = ak.argpartition(good_jet_pts, axis=1)
ordered_2_jets = ak.pad_none(
    good_jet_pts[ix_jet_pt_2_largest],
    

Replies: 2 comments 10 replies

Comment options

You must be logged in to vote
5 replies
@green-cabbage
Comment options

@NJManganelli
Comment options

@green-cabbage
Comment options

@agoose77
Comment options

@green-cabbage
Comment options

Answer selected by green-cabbage
Comment options

You must be logged in to vote
5 replies
@agoose77
Comment options

@green-cabbage
Comment options

@jpivarski
Comment options

@agoose77
Comment options

@jpivarski
Comment options

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Category
Q&A
Labels
None yet
4 participants