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

[green-cabbage]

Hello, I have a question about my specific use case of obtaining leading and subleading values from each columns of ak array. Let's assume that we have an array of jet pt values from a coffea.nanoevent :

import awkward as ak
jet_pts = ak.Array(
    [
        [3,2,1,0],
        [1,],
        [5,4]
    ]
)
pt_jec = ak.Array(
    [
        [0,4,0,0],
        [0,],
        [-1,2]
    ]
)
jet_pts = jet_pts + pt_jec
jet_selection = ak.Array(
    [
        [True,True,True,False],
        [False],
        [True,False]
    ]
)
good_jet_pts = jet_pts[jet_selection]

We can notice that the original jet_pts already lists its pt values in descending order on each row, which is great, but applying jet energy correction (JEC) can some times make the second column value to be now bigger than the first column vlaue. At the end of this dummy code, printing good_jet_pts gives me this value:

[[3, 6, 1],
 [],
 [4]]
---------------------
type: 3 * var * int64

Now, I need to get the top two leading jet pts seperately, and I do this by filling none the first two columns and then using ak where:

padded_good_jet_pts = ak.pad_none(good_jet_pts, 2)
pt_first_col = padded_good_jet_pts[:,0]
pt_second_col = padded_good_jet_pts[:,1]
jet_flip = pt_first_col < pt_second_col
jet_flip = ak.fill_none(jet_flip, value = False)
leading_jet_pt = ak.where(jet_flip, pt_second_col, pt_first_col)
subleading_jet_pt = ak.where(jet_flip, pt_first_col, pt_second_col)
print(f"leading_jet_pt : {leading_jet_pt}")
print(f"subleading_jet_pt : {subleading_jet_pt}")

This gives me

leading_jet_pt : [6, None, 4]
subleading_jet_pt : [3, None, None]

Which is exactly what I want. However, I realize that padding ak arrays with None values can bloat the memory usage considerably, so I would like to know if there's a way to do this operation without ak.fill_none(). My current plan B is to use range slicing (probably not the right name for it haha), where I do instead:

pt_first_col = good_jet_pts[:,:1]
pt_second_col = good_jet_pts[:,1:2]
print(f"pt_first_col : {pt_first_col}")
print(f"pt_second_col : {pt_second_col}")

And this prints pt_first_col : [[3], [], [4]] and pt_second_col : [[6], [], []] . This is great, but then when I try to obtain the jet_flip boolean array like:

jet_flip = pt_first_col < pt_second_col

I get this error message:

ValueError: cannot broadcast nested list

This error occurred while calling

    numpy.less.__call__(
        <Array [[3], [], [4]] type='3 * var * int64'>
        <Array [[6], [], []] type='3 * var * int64'>
    )

This is the point that I am stuck on. If anybody knows a way around this issue or a completely different way to go about this, I would be obliged! Thank you!

[agoose77]

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],
    target=2,
    clip=True
)

leading_jet = ordered_2_jets[:, 0]
subleading_jet = ordered_2_jets[:, 1]

As of the time of writing, we don't have an implementation of argpartition. Fundamentally, we want a selection algorithm that also (ideally) sorts the "larger" region (in our case), so we don't need a follow-up sort. I recall having this conversation with @jpivarski on Slack in October of 2023 (date for his benefit). I think Numba supports this, so perhaps we can implement something right now.

It should go without saying, all of this depends upon how many jets you have per event. If the answer is "not many", then this is all probably noise in the wind vs a simple sort-and-slice.

I started playing around with this in a notebook. I can't make an strong insights into it yet - I ran out of time. I assume @jpivarski might weigh in here too, so it might be useful to save him some time.

https://gist.github.com/agoose77/cd77e3b230b2132cbedff05c162eb222

What I can say is that from the findings above, is that for small numbers of jets, just sort the array and slice it. For larger jet counts, you'll perhaps want to perform some kind of partition.

[green-cabbage]

Thank you for this in depth answer. I took a quick look, sadly, I can't really understand what's going on with function definitions such as nlargest_last() in https://gist.github.com/agoose77/cd77e3b230b2132cbedff05c162eb222, but I get the big points.

As for your question of how many top N jets we are looking to obtain, the answer is indeed not many. Currently, just two leading jets, maybe in some fringe case, three or four in the future, but definitely a small number. The major concern was definitely a way to sort these arrays, and ak.sort() definitely fits the bill, though I wonder if this is applicable with coffea.nanoevents... in any circumstances, thank you for the input!

[NJManganelli]

Use a single ak.argsort result as index, then try range-slice (plus pad_none and fill_none where necessary to fill defaults). This is a typical pattern in coffea.

We do not have that many jets usually: after the first 2-10 jets (hard process + pileup) each additional one incurs a rate penalty (very roughly) ~1/10 of previous multiplicity.

[green-cabbage]

I see; that's great to hear. May I ask why ak.argsort is recommended over just ak.sort ? Or are they both just as good?

Also, this is a bit of a tangent question, but for cpu/memory performance perspective, is it recommended to ak.fill_none() at a coffea.nanoevent level, or at kinematic variable level individually (where the actual ak arrays lie)? In the sense that:
nanoevent level:

Jets = Array(
    [
        [Jet, Jet],
        [Jet],
        []
        ]
)
padded_jets = ak.pad_none(Jets, target=2)
jet1_pt = padded_jets.pt[:,0]
jet1_eta = padded_jets.eta[:,0]
jet1_phi = padded_jets.phi[:,0]

kinematic variable level:

Jets = Array(
    [
        [Jet, Jet],
        [Jet],
        []
        ]
)
padded_jet_pts = ak.pad_none(Jets.pt, target=2)
jet1_pt = padded_jet_pts[:,0]
padded_jet_etas = ak.pad_none(Jets.eta, target=2)
jet1_eta = padded_jet_etas[:,0]
padded_jet_phis = ak.pad_none(Jets.phi, target=2)
jet1_eta = padded_jet_phis[:,0]

Conceptually, it's more elegant if I ak.pad_none at the nanoevent level, because that gets applied to all kinematics (pt, eta and phi) in one stroke, but I am also curious it this gives a performance penalty, since operating on nanoevents could apply the same operation in ALL the fields in a nanoevent.

[agoose77]

Applying ak.pad_none is not particularly expensive compared to other array operations. If it is introducing a None type, then it adds a new int64 index, which is small compared to the number of records in your array (probably). If you apply pad_none to a record array, it will be more efficient in memory than if you apply it to every leaf (record). In terms of compute performance, it should be more efficient there too in some cases, during certain transforms. In general, though, don't worry too much about that.

[green-cabbage]

Also, just to answer my own question of "why not just use ak.sort onto a coffea.nanoevents.methods.vector (ie. events.Jet)?", if one simply applies ak.sort(jets), then all the kinematic/leaf variables (ie pt, eta, phi, etc) would be sorted regardless if the values correspond to the same jet variable. This would result in the first column Jet variable to have the highest pt, mass, phi and eta, which physically wouldn't really make sense. A quick ak.sort example would be:

jets = events.Jet
njets = ak.num(jets, axis=1)
jet_selection = njets == 2
good_jets = jets[jet_selection]
num_events = ak.num(good_jets, axis=0)
pt_override = np.zeros((num_events,2))
pt_override[:,1] = 10000
good_jets["pt"] = pt_override
mass_override = np.zeros((num_events,2))
mass_override[:,0] = 10000
good_jets["mass"] = mass_override
print(f"good jets pt b4 sort {good_jets.pt}")
print(f"good jets mass b4 sort {good_jets.mass}")
sorted_jets = ak.sort(good_jets, ascending=False)
print(f"good jets pt after sort {sorted_jets.pt}")
print(f"good jets mass after sort {sorted_jets.mass}")

This would print

good jets pt b4 sort [[0, 1e+04], [0, 1e+04], [0, 1e+04], ..., [0, ...], [0, 1e+04], [0, 1e+04]]
good jets mass b4 sort [[1e+04, 0], [1e+04, 0], [1e+04, 0], ..., [1e+04, ...], [1e+04, 0], [1e+04, 0]]
good jets pt after sort [[1e+04, 0], [1e+04, 0], [1e+04, 0], ..., [1e+04, ...], [1e+04, 0], [1e+04, 0]]
good jets mass after sort [[1e+04, 0], [1e+04, 0], [1e+04, 0], ..., [1e+04, ...], [1e+04, 0], [1e+04, 0]]

Where the first column jets have both pt and mass of 10,000, which didn't exist in the original good_jets array. For maintaining particle consistency of the jets array, we can use ak.argsort(), for example, if we would like to re-order the jet array in the descending order of pt, we can do:

ptsorted_args = ak.argsort(good_jets.pt, ascending=False)
ptsorted_jets = good_jets[ptsorted_args]
print(f"good jets pt after pt sort {ptsorted_jets.pt}")
print(f"good jets mass after pt sort {ptsorted_jets.mass}")

And this prints

good jets pt after pt sort [[1e+04, 0], [1e+04, 0], [1e+04, 0], ..., [1e+04, ...], [1e+04, 0], [1e+04, 0]]
good jets mass after pt sort [[0, 1e+04], [0, 1e+04], [0, 1e+04], ..., [0, ...], [0, 1e+04], [0, 1e+04]]

Which is what we want, since the jet arrays are sorted in the pt descending order, while its other kinematic variables follow the sorted jet vectors as well.