Applying cuts on branches that are sometimes empty

[delaere]

Hi,

I have a pretty standard tree (showing only part of it):

name                 | typename                 | interpretation                
---------------------+--------------------------+-------------------------------
nMuons               | int32_t                  | AsDtype('>i4')
MuonsPt              | double[]                 | AsJagged(AsDtype('>f8'))
invMass              | double                   | AsDtype('>f8')

nMuons can be 0, in which cas MuonsPt is empty. I want to filter on the MuonsPt, so I naively do:

tree.arrays(["invMass"],"(nMuons==2) & (MuonsPt[:,0]>20)")

but it fails because the cut on MuonsPt is evaluated even for events with nMuons==0, and fails with an index out of range error.

I managed to do what I want in two steps (first filtering on nMuons, later on MuonsPt) with the array filter synthax, but I would prefer the cut string approach for pedagogical reasons.

What would you suggest in such a case?

Thanks!

[jpivarski]

This is not so much about making it a cut string in Uproot as finding a one-liner in Awkward Array. Since operations are columnar, (nMuons == 2) is calculated independently of (MuonsPt[:, 0] > 20) before computing their intersection (&), and the latter is invalid for some array entries. For this application, you could use the nMuons == 2 as a slice on the first dimension, rather than selecting everything in the first dimension with :.

For instance, if you start with

>>> import awkward as ak
>>> array = ak.Array([[1, 2, 3], [], [4, 5]])
>>> array
<Array [[1, 2, 3], [], [4, 5]] type='3 * var * int64'>

you can't do

>>> (ak.num(array) > 0) & (array[:, 0])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jpivarski/irishep/awkward-1.0/awkward/highlevel.py", line 996, in __getitem__
    tmp = ak._util.wrap(self.layout[where], self._behavior)
ValueError: in ListOffsetArray64 attempting to get 0, index out of range

but you can do

>>> array[ak.num(array) > 0, 0]
<Array [1, 4] type='2 * int64'>

or even

>>> array[ak.num(array) > 0][:, 0]
<Array [1, 4] type='2 * int64'>

The thing is that the selection of events (first dimension) that makes the selection of particles (second dimension) valid has to be applied to the same array. When ak.num(array) > 0 is applied to the first dimension, there aren't any events left that don't have muons.

Following this logic, I was about to say that your cut string could be "MuonsPt[nMuons == 2, 0] > 20". But the > 20 part is a cut on "events with two muons," and you can't apply that cut to "events," which is what the cut argument to tree.arrays does. In my example, array[ak.num(array) > 0][:, 0] returns a boolean array of length 2, but array has length 3: it won't be an applicable cut. Here's something you can do instead:

MuonsPt.mask[nMuons == 2][:, 0] > 20

By analogy (going back to my example),

>>> array.mask[ak.num(array) > 0]
<Array [[1, 2, 3], None, [4, 5]] type='3 * option[var * int64]'>

puts a "None" in the event that fails the multiplicity cut, and then

>>> array.mask[ak.num(array) > 0][:, 0]
<Array [1, None, 4] type='3 * ?int64'>

passes that "None" through, and

>>> array.mask[ak.num(array) > 0][:, 0] > 1
<Array [False, None, True] type='3 * ?bool'>

creates a boolean-or-None array, which at least has the right length; it applies to "all events," not just "events with two muons" (going back to your example, by analogy). This would drop values of "invMass" that are in the same position as a "False" and it would return "None" for values that are in the same position as a "None". Maybe you don't want that; you probably want both of those cases to drop values. You want the equivalent of

>>> ak.fill_none(array.mask[ak.num(array) > 0][:, 0] > 1, False)
<Array [False, False, True] type='3 * bool'>

which, in your case, is

ak.fill_none(MuonsPt.mask[nMuons == 2][:, 0] > 20, False)

This replaces all "None" values with "False" values, giving you a purely boolean array with "False" for both the case of not having two muons and the case of the first muon having a pT less than 20.

As usual, turning things into one-liners doesn't always make them more clear. If you're doing this for pedagogy, will your students/audience be able to read that one-liner? They'd need to know about ak.mask and ak.fill_none.

(I'm not sure it was a great idea to introduce the expressions and cut arguments in TTree.arrays because it puts an artificial boundary between "calculations that can be done in the expressions/cut strings" and "all possible calculations." If someone is following the example of a script that manages to express everything in expressions/cut but they need to extend it to do something that can't, it won't be clear from the source script how to do that. It's like the ROOT TBrowser: double-clicking is an easier way of getting to the plots, but the fact that double-clicking is possible prevents some new students from finding out that they can also TTree::Draw, and those students don't know how to apply a cut (just as TTree:Draw prevents students from finding out that they can compute more complex functions by scripting). I think this is why Python has an "only one way to do something" ideal, because the existence of a shortcut makes it harder to find out that there's a non-shortcut method, and sometimes the non-shortcut method is needed. So I'm just a little nervous that I made a bad design decision by including the cut argument at all...)

[delaere]

Thanks a lot for this detailed and interesting explanation!
I agree with your conclusion that this is probably too advanced for students in my context. I wanted to reduce the exposition to specific array syntax, and the result would be just the opposite.

Thanks again!