examples of doing heterogenous combinations in awkward0

column_combinatorics.py

from uproot_methods import TLorentzVectorArray
import awkward as ak
from coffea.nanoaod import NanoEvents

def make_labeled_p4(x, indices, itype):
    p4 = TLorentzVectorArray.from_ptetaphim(x.pt, x.eta, x.phi, x.mass)
    return ak.JaggedArray.zip(p4=p4,
                              ptype=itype*x.pt.ones_like().astype(np.int),
                              pidx=indices
                              charge=x.charge)

def awkward0_pstack(one, two):
    one_indices = ak.JaggedArray.fromoffsets(one.pt.offsets, 
                                             np.arange(0, one.pt.content.size)) - one.pt.offsets[:-1]
    two_indices = ak.JaggedArray.fromoffsets(two.pt.offsets, 
                                             np.arange(0, two.pt.content.size)) - two.pt.offsets[:-1]
    one_p4 = make_labeled_p4(one, one_indices, 0)
    two_p4 = make_labeled_p4(two, two_indices, 1)
    
    stacked_p4 = ak.concatenate((one_p4, two_p4), axis=1)
    
    return stacked_p4

# this is from within a cloned coffea repository
e = NanoEvents.from_file('tests/samples/nano_dimuon.root')
  
#loose dilepton event veto
loose_mu = e.Muon[e.Muon.looseId]
loose_e = e.Electron[e.Electron.cutBased > 0]
enms = awkward0_pstack(loose_mu, loose_e)
pairs = enms.argchoose(2)
mass = (enms[pairs.i0].p4 + enms[pairs.i1].p4).mass
loose_preselect = ((mass > 12).all() & (mass.counts > 0))
print(loose_preselect)

#dilepton (I've left out the IDs on the leptons since I only had a small sample to test)
muon_dilep = e.Muon # tight muons
ele_dilep = e.Electron # tight electrons
#muons have ptype = 0, electrons ptype = 1
leptons = awkward0_pstack(muon_dilep, ele_dilep)
pairs = leptons.argchoose(2)
masses = (leptons[pairs.i0].p4 + leptons[pairs.i1].p4).mass
types = (leptons[pairs.i0].ptype + leptons[pairs.i1].ptype)
di_electron_z_veto = ((types < 2) | (np.abs(91.18 - masses) > 10.)).any()
dilep_select = loose_preselect & di_electron_z_veto & (pairs.counts == 1)
print(dilep_select)

def has_ossf_dileptons(dilep):
    return ((dilep.i0.charge != dilep.i1.charge) & 
            (np.abs(91.18 - (dilep.i0 + dilep.i1).mass) < 15.)).any()

#trilepton (ditto, we're assuming we're operating on the tight leptons here)
muon_trilep = e.Muon # tight muons
ele_trilep = e.Electron # tight electrons
dimuon_fortri = muon_trilep.choose(2)
mu_z_cand = has_ossf_dileptons(dimuon_fortri)
dielectron_fortri = ele_trilep.choose(2)
ele_z_cand = has_ossf_dileptons(dielectron_fortri)
leptons = awkward0_pstack(muon_trilep, ele_trilep)
triples = leptons.argchoose(3)
trilep_select = loose_preselect & (~mu_z_cand) & (~ele_z_cand) & (triples.counts == 1)