Распространение выделения на подмножестве неудобного массива

Question

Есть ли лучший способ сделать эту логику c? Я хочу распространить выборку из выборки более низкого уровня, доступную только на подмножество внутренних элементов, вверх

В частности, я хочу сократить уровень событий для противоположно заряженной пары мюон-электрон.

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]