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.

[green-cabbage]

Sorry about bringing up this thread from the grave, but I was wondering about ak.pad_none generating a int64 index of "None" value as @agoose77 mentioned previously.

Even if it is used at a record array level, wouldn't it be more memory efficient to have int32 index for "None" rather than int64?

Since I am processing my samples with millions of events (at least before any selection), even padding just one or two each row with int64 variable could add up in memory. Therefore, I was wondering if there is an easy way to enforce padded "None" values with int32 instead?

Thanks in advance!

[agoose77]

Even if it is used at a record array level, wouldn't it be more memory efficient to have int32 index for "None" rather than int64?

Yes, although in practice probably not because we use special compiled kernels that perform the high-performance lookups, and these are often specialized to int64. You can certainly look to see what happens!

[green-cabbage]

Sweet, so what could I do to set "None" indices in int32? Do I do something to the effect of : ak.pad_none(array, attrs={"type": np.int32}) ?

[jpivarski]

There isn't a high-level way to do this.

But there are some low-level ways, and they involve accessing some internal attributes that are unlikely (but possible!) to change. The following should not be used in production code, or anything in which it will matter if we change the interface and break it in the future. (That is, you can use it for quick tests and debugging.)

Staring from an array with missing values created by the ak.pad_none function,

>>> a = ak.Array([[0, 1, 2], [], [3, 4], [5], [6, 7, 8, 9]])
>>> b = ak.pad_none(a, target=4, clip=True)
>>> b.show()
[[0, 1, 2, None],
 [None, None, None, None],
 [3, 4, None, None],
 [5, None, None, None],
 [6, 7, 8, 9]]

you saw that the indexes are 64-bit:

>>> b.layout
<RegularArray size='4' len='5'>
    <content><IndexedOptionArray len='20'>
        <index><Index dtype='int64' len='20'>
            [ 0  1  2 -1 -1 -1 -1 -1  3  4 -1 -1  5 -1 -1 -1  6  7  8  9]
        </Index></index>
        <content><NumpyArray dtype='int64' len='10'>[0 1 2 3 4 5 6 7 8 9]</NumpyArray></content>
    </IndexedOptionArray></content>
</RegularArray>

By replacing the b.layout.content._index attribute, you can replace the 64-bit index with a 32-bit index. Since this attribute starts with an underscore, we can change its name without warning (but probably won't; this one is old and unlikely to change).

>>> b.layout.content._index = ak.index.Index32(b.layout.content._index.data)
>>> b.layout
<RegularArray size='4' len='5'>
    <content><IndexedOptionArray len='20'>
        <index><Index dtype='int32' len='20'>
            [ 0  1  2 -1 -1 -1 -1 -1  3  4 -1 -1  5 -1 -1 -1  6  7  8  9]
        </Index></index>
        <content><NumpyArray dtype='int64' len='10'>[0 1 2 3 4 5 6 7 8 9]</NumpyArray></content>
    </IndexedOptionArray></content>
</RegularArray>

But more than that, you can replace the IndexedOptionArray, which is now using 4 bytes per item in the index, with a ByteMaskedArray, which uses 1 byte per item.

>>> b.layout._content = b.layout._content.to_ByteMaskedArray(valid_when=True)
>>> b.layout
<RegularArray size='4' len='5'>
    <content><ByteMaskedArray valid_when='true' len='20'>
        <mask><Index dtype='int8' len='20'>
            [1 1 1 0 0 0 0 0 1 1 0 0 1 0 0 0 1 1 1 1]
        </Index></mask>
        <content><NumpyArray dtype='int64' len='20'>
            [0 1 2 9 9 9 9 9 3 4 9 9 5 9 9 9 6 7 8 9]
        </NumpyArray></content>
    </ByteMaskedArray></content>
</RegularArray>

It has the same high-level meaning:

>>> b.show()
[[0, 1, 2, None],
 [None, None, None, None],
 [3, 4, None, None],
 [5, None, None, None],
 [6, 7, 8, 9]]

However, notice that the nested NumpyArray is now padded with numbers. (The fact that it's 9 is very much an implementation detail!) ByteMaskedArray has a hidden cost: missing values have to be instantiated, so it's really 1 + itemsize bytes per value. The cost for IndexedOptionArray depends on how many items are missing, since they don't have to appear in the content. This can be especially bad if the content is a very large type, like some kind of record.

There's also a BitMaskedArray, which uses a single bit for each mask, but don't be tempted: nearly every operation immediately converts it into a ByteMaskedArray. It might be a useful type for saving to disk, but most of the file types do some sort of packing and compression anyway. (BitMaskedArray only exists so that Arrow masks can be viewed by Awkward Array without copying.)

[agoose77]

Actually @jpivarski I think it is pretty easy to change an int64 mask to an int32 using the public API:

def transformer(layout, continuation, **kwargs):
    if layout.is_indexed and layout.is_option:
        next_layout = continuation()
        index_32_bit = ak.index.Index32(layout.index.data)
        return next_layout.copy(index=index_32_bit)

array_with_32_bit_index = ak.transform(transformer, array)

The same approach can be used to convert options into ByteMaskedArrays. Of course, if you want to prevent the 64-bit index in the first place ... that's more tricky.

[jpivarski]

Oh, that's good! That's a great idea! It's not at all obvious (you need to know about layouts and how ak.transform's transformer functions work), but it does avoid the need to access a hidden attribute.

Also, @agoose77's solution will work for IndexedOptionArrays (layout.is_indexed and layout.is_option) nested within any number of lists, records, etc., whereas the one I wrote will only work for this exact layout.