1

Is there a better way to do this logic? I want to propagate a selection from a lower-level selection available only on a subset of inner elements upwards

Specifically, I am looking to have an event level cut for oppositely charged muon-electron pair.

req_mu = (events.Muon.counts >= 1)
req_ele = (events.Electron.counts >= 1)
req = req_ele & req_mu

def propagate_up(subset, selection):
    '''
    subset: bool array slice on upper level
    '''
    dummy = np.zeros_like(subset)
    dummy[subset] = selection
    return dummy

req_opposite_charge = propagate_up(req, events[req].Muon[:, 0].charge * events[req].Electron[:, 0].charge == -1)

Andrzej Novák
  • 136
  • 1
  • 8

1 Answers1

1

The easiest way to select a pair from separate collections is with cross, e.g.

good_el = electrons[electrons.pt > 10]
good_mu = muons[muons.pt > 10]
pairs = good_el.cross(good_mu)
# filter our pairs to have opposite charge (jagged mask)
pairs = pairs[pairs.i0.charge == -1 * pairs.i1.charge]
# filter events to have exactly one good pair
pairs = pairs[pairs.counts == 1]
# get the two leptons as a now flat array
el, mu = pairs.i0[:, 0], pairs.i1[:, 0]