Calculating isolation

[cesarecazzaniga]

Dear experts,

I cannot find a straightforward way to make a calculation using ak array operations for a particular task.
I post here the formula I would like to compute.

At the moment I tried with for loops, but it is very inefficient, and also with Numba some operations are not supported. Can you suggest some fast way to implement this please?

[agoose77]

Could you share some example arrays so that we can work out how your problem maps to this formula?

[agoose77]

I'm still not fully sure what all of your quantities mean. Perhaps you could share a loop version?

In general, you can use boolean slicing or array masking to cut out particles that do not meet the relevant criteria, e.g.
Let's assume your array is N * var, where N is the number of events, and var is the multiplicity of each event.

To get a boolean slice that takes only the particles with pdgId i, we can write

is_pdgid_i = event.pdgId == i

[cesarecazzaniga]

So, when I tried to construct the ak.ArrayBuilder(), I did something like this:

def compute_builder_isolation(builder,jet_pf_cands_plus_leptons,leptons_pdgid,pt_min,R):                                                                                                                                                                                                                   
   for event_particles_list in jet_pf_cands_plus_leptons:                                                                                                                                                                          
     builder.begin_list()
     for P in range(len(event_particles_list)):
      isolation = -99999.
      if (event_particles_list[P].pdgId == leptons_pdgid):
       for i in range(len(event_particles_list)):
           sum = 0
           if i != P:
             delta_r_iP = np.sqrt((event_particles_list[i].phi - event_particles_list[P].phi)**2 + (event_particles_list[i].eta - event_particles_list[P].eta)**2)
             if event_particles_list[i].pt > pt_min and delta_r_iP < R:
               sum += event_particles_list[i].pt
       isolation = sum/event_particles_list[P].pt
      builder.append(isolation)
     builder.end_list()

   return builder 

[jpivarski]

The purpose of ak.ArrayBuilder is to convert non-columnar data into columnar data. If jet_pf_cands_plus_leptons is already an Awkward Array, there is no reason to bring ArrayBuilder into it. (For instance, even if you are unconcerned about the time the computation will take and will do it with nested for loops, you could have those loops create Python lists of dicts and pass them to ak.from_iter.) Oh! Maybe you're setting this up to run in Numba, and so ArrayBuilder is the only way to get it out of the Numba-compiled function?

Another thing to know is that you can avoid having to calculate ΔR yourself by installing and using Vector:

import vector
vector.register_awkward()

jet_pf_cands_plus_leptons = ak.with_name(jet_pf_cands_plus_leptons, "Momentum4D")

jet_pf_cands_plus_leptons.deltaR(...)   # put another array of vectors in the "..."

---

Anyway, it looks to me like what you want to do is

pdgid_cut = (jet_pf_cands_plus_leptons.pdgId == leptons_pdgid)
pt_cut = (jet_pf_cands_plus_leptons.pt > pt_min)

good_leptons = jet_pf_cands_plus_leptons[pdgid_cut & pt_cut]

to get good_leptons that have the right pdgid (note: do you want both pdgid and -pdgid? if so, use abs) and high enough pt.

Now you're going to want to consider all pairs, indexed with P and i, of these good_leptons per event. I would normally point you to ak.combinations, which makes exactly the "N choose 2" pairs without repeating that you want, except that you're going to want to sum over all leptons i for a given P, so they're going to need to be grouped by P. I think the clearest way of doing that is to use ak.cartesian with nested=True and we'll exclude the cases in which i == P (they have ΔR == 0, after all).

Since I don't have your dataset, I'm going to make a simplified one with the relevant features:

>>> good_leptons = ak.Array([
...     [{"x": 0, "y": 1, "z": 2}],
...     [{"x": 0, "y": 1, "z": 2}, {"x": 1, "y": 2, "z": 3}],
...     [{"x": 0, "y": 1, "z": 2}, {"x": 1, "y": 2, "z": 3}, {"x": 2, "y": 3, "z": 4}],
...     [{"x": 0, "y": 1, "z": 2}, {"x": 1, "y": 2, "z": 3}, {"x": 2, "y": 3, "z": 4}, {"x": 3, "y": 4, "z": 5}],
... ], with_name="Momentum3D")
>>> good_leptons
<MomentumArray3D [[{x: 0, y: 1, z: 2}, ... y: 4, z: 5}]] type='4 * var * Momentu...'>

ak.cartesian takes a two lists, like "A, B, C" and "1, 2", and returns their Cartesian product, like "A-1, A-2, B-1, B-2, C-1, C-2". Setting nested=True turns the input array of lists of objects into an output array of lists of lists of objects—one level deeper—so that the output can be grouped by one of the input values: "[A-1, A-2], [B-1, B-2], [C-1, C-2]". We'll want that so that we can ask for the sum over each inner list and get the sum over A, followed by the sum over B, followed by the sum over C. ak.unzip splits the left and right halves of each pair into separate arrays so that we can use them both in calculations (such as the ΔR calcuation).

>>> leptons_P, leptons_i = ak.unzip(ak.cartesian([good_leptons, good_leptons], nested=True))

>>> leptons_P.x.tolist()
[
    [[0]],
    [[0, 0], [1, 1]],
    [[0, 0, 0], [1, 1, 1], [2, 2, 2]],
    [[0, 0, 0, 0], [1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]
]
>>> leptons_i.x.tolist()
[
    [[0]],
    [[0, 1], [0, 1]],
    [[0, 1, 2], [0, 1, 2], [0, 1, 2]],
    [[0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3]]
]

See how the innermost numbers are in same-valued groups for leptons_P and in different-valued groups for leptons_i? When we do any calculations involving elements from leptons_P and leptons_i, we'll be doing it on all combinations as though they're in two nested for loops.

Now that we have these two collections, we can calculate the ΔR between them:

>>> deltaR = leptons_P.deltaR(leptons_i)

>>> deltaR.tolist()
[
    [[0.0]],
    [[0.0, 0.5749801549254422], [0.5749801549254422, 0.0]],
    [[0.0, 0.5749801549254422, 0.763476592636774], [0.5749801549254422, 0.0, 0.1924914766026644], [0.763476592636774, 0.1924914766026644, 0.0]],
    [[0.0, 0.5749801549254422, 0.763476592636774, 0.8545361946077011], [0.5749801549254422, 0.0, 0.1924914766026644, 0.28587761696419933], [0.763476592636774, 0.1924914766026644, 0.0, 0.09352784779524788], [0.8545361946077011, 0.28587761696419933, 0.09352784779524788, 0.0]]
]

We don't want the ΔR that are zero, because they're the same particle, and we don't want values larger than R. That's another cut:

>>> pairs_cut = (0 < deltaR) & (deltaR < 0.5)
>>> pairs_cut
<Array [[[False]], ... True, True, False]]] type='4 * var * var * bool'>

and it can be applied to these deeply nested leptons:

>>> isolated_i = leptons_i[pairs_cut]

>>> isolated_i.tolist()
[
    [[]],
    [[], []],
    [[], [{'x': 2, 'y': 3, 'z': 4}], [{'x': 1, 'y': 2, 'z': 3}]],
    [[], [{'x': 2, 'y': 3, 'z': 4}, {'x': 3, 'y': 4, 'z': 5}], [{'x': 1, 'y': 2, 'z': 3}, {'x': 3, 'y': 4, 'z': 5}], [{'x': 1, 'y': 2, 'z': 3}, {'x': 2, 'y': 3, 'z': 4}]]
]

Because we wanted to sum over the pt of those leptons in i with these cuts:

>>> sums_over_i = ak.sum(isolated_i.pt, axis=-1)

>>> sums_over_i.tolist()
[
    [0.0],
    [0.0, 0.0],
    [0.0, 3.605551275463989, 2.23606797749979],
    [0.0, 8.60555127546399, 7.23606797749979, 5.841619252963779]
]

Notice something about this sum, though: it's an array of lists of numbers—no more "lists of lists". That's because summation turns each of the inner lists into a number, and now we're back to the same multiplicity as good_leptons. Your isolation variable is therefore

>>> isolation = sums_over_i / good_leptons.pt

>>> isolation.tolist()
[
    [0.0],
    [0.0, 0.0],
    [0.0, 1.6124515496597098, 0.6201736729460423],
    [0.0, 3.8485195271594996, 2.0069241635091153, 1.1683238505927558]
]

The good_leptons that I constructed was fake, intended to make the step-by-step logic of this more clear. When you try to apply this with your leptons, do it interactively in a terminal, IPython, or Jupyter notebook! It will be absolutely necessary to be able to interact with the data to get this right. If you put it all in a script and run it, you'll very likely run into some minor indexing problem that would be resolvable if you could just look at the data. It would be sufficient to take a sample of 10 events or less—that would be enough to figure out if you're aligning the indexes correctly.

---

I'm just realizing, at the end of all this, that you might have the whole thing working in Numba. Is that true? If so, it isn't necessary to convert it to a columnar form—usually, you do the columnar analysis (with ak.cartesian and all that) because you get to see your intermediate steps more readily than when it's all hidden inside a for loop, especially if it's compiled by Numba.

Anyway, now you can compare the for loop form and the columnar form, to see how the two approaches differ.

[cesarecazzaniga]

thanks a lot for the help !

[cesarecazzaniga]

I think this is a very nice solution, and works.