Store smallest deltaR result as new feature

[artlbv]

(I name this issue as the task I need to do, but there is probably a better more technical name for this)

I'm using awkward arrays as spit out by uproot4 from "flat" ROOT files.

The problem is the following:

• I have two object awkward arrays, e.g. for some selected gen-particles (genbs) and reco-jets (l1jets).
• I would like to store the smallest deltaR of the reco jets to the gen particles as another "property" of the reco jets

For the deltaR computation I use cell 82 from this tutorial notebook:

js, bs = ak.unzip(ak.cartesian((l1jets, genbs)))
dR = js.deltaR(bs)

A small "snippet" is in this screenshot (don't know how to insert that here otherwise)

Now what I would like to do is to create a new array/feauture e.g. min_dR in the l1jets that would correspond to the smallest dR of the jet to the gen particles in this event.

If I'd knew there are always e.g. 2 gen particles in an event (like in the selection above), I could in principle just slice the dR with a step of 2 and with offsets i for gen particle i = 0,1.

dR_b1 = dR[:,0::2]
dR_b2 = dR[:,1::2]

Given that I actually care about the dR being smaller than some threshold, I thought could then just do:

(dR_b1 < 0.4) & (dR_b2 < 0.4)

but that gives error ValueError: in ListOffsetArray64, cannot broadcast nested list.

This will in general not work when I have an a-priori unknown length of both of the input arrays.

So I was wondering whether it would be possible to somehow slice/reshape the arrays using the per-event number of input objects, etc to create a smallest deltaR array for the reco jets. Alternatively, the deltaR < thr result would also be fine to store.

Are there any built-in features that would allow this? Or an elegant way to achieve it without too much hacking.

[jpivarski]

A few weeks after that tutorial, I gave another tutorial that answered your question. (When writing it, I decided to go into more detail on these sorts of examples, but when I was presenting it, I didn't get to that part, so it's not in the video.)

This is the relevant section.

Something that I was going to do when presenting it live and didn't get a chance to: try swapping the abs(g.pdgId) == 13 with abs(g.pdgId) != 13 and see how the distribution of ΔR changes, once you have this notebook loaded and have evaluated cells up to this point.

The idea is that you pass nested=True when computing ak.cartesian, which creates a new level of nesting in which the combinations are grouped by the left item. That is, if you compute the Cartesian product of l1jets and genbs with nested=True, the new level of list-nesting will have all l1jets[..., 0] in the first new list, all l1jets[..., 1] in the second new list, and so on.

That extra level of nesting gives you structure to "reduce" over. In this case, you want to reduce by ak.min or ak.argmin; i.e. find the jet-gen pair with the smallest ΔR. If instead of finding the "best match," you were only interested in whether there was "any match within a ΔR cut," the reducer you'd use would be ak.any.

Reducing these nested lists, of which there is one for each of the l1jets, leaves you with an array that has one inner item per l1jet (each list gets reduced to one value). That gives you a "best gen" and "best ΔR" for each l1jet, and that's an array that you can use with the l1jets data. (Hint: if you have a new field with the right number of items at each level, you can assign it to an existing record: ak.Array.setitem.)

So that's the idea; hopefully, the notebook will be a help. If that's too difficult to follow, I also worked through the same question again using Numba instead of Awkward array-at-a-time functions. If might be easier to read the nested "for" loops in Numba, and since it's JIT-compiled, it's feasible to use at large scales.

[artlbv]

Worked like a charm, thanks a lot Jim!

jet_gen = ak.cartesian({"l1jets": l1jets, "gen": genjets}, nested=True)
js, gs = ak.unzip(jet_gen)
dR = js.deltaR(gs)

l1jets["min_dR_genjet"] = ak.min(dR, axis = -1)
best_dR = ak.argmin(dR, axis=-1, keepdims=True)
l1jets["partonFlavour"] = jet_gen[best_dR].gen.partonFlavour

Stupid me was for some reason watching and reading your tutorial from PyHEP2020