awkward-array - 在笨拙数组备份的子集上传播选择

Question

有没有更好的方法来做这个逻辑？我想从仅在内部元素子集上可用的较低级别选择中传播选择

具体来说，我希望降低带相反电荷的 μ 电子对的事件水平。

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)

Answer 1

从单独的集合中选择一对的最简单方法是使用cross，例如

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]