def select_jets(jets, mu, el, selected_muons, selected_electrons, pt_cut,
aeta_cut, jet_lepton_dr_cut, btag_cut):
this_worker = get_worker_wrapper()
NUMPY_LIB = this_worker.NUMPY_LIB
backend = this_worker.backend
passes_id = jets.puId == True
passes_aeta = NUMPY_LIB.abs(jets.eta) < aeta_cut
passes_pt = jets.pt > pt_cut
passes_btag = jets.btag > btag_cut
selected_jets = passes_id & passes_aeta & passes_pt
jets_d = {"eta": jets.eta, "phi": jets.phi, "offsets": jets.offsets}
jets_pass_dr_mu = ha_kernels.mask_deltar_first(backend, jets_d,
selected_jets, {
"eta": mu.eta,
"phi": mu.phi,
"offsets": mu.offsets
}, selected_muons,
jet_lepton_dr_cut)
jets_pass_dr_el = ha_kernels.mask_deltar_first(backend, jets_d,
selected_jets, {
"eta": el.eta,
"phi": el.phi,
"offsets": el.offsets
}, selected_electrons,
jet_lepton_dr_cut)
selected_jets_no_lepton = selected_jets & jets_pass_dr_mu & jets_pass_dr_el
selected_jets_btag = selected_jets_no_lepton & passes_btag
return selected_jets_no_lepton, selected_jets_btag
def test_kernel_mask_deltar_first(self):
print("kernel_mask_deltar_first")
dataset = self.dataset
muons = dataset.structs["Muon"][0]
jet = dataset.structs["Jet"][0]
sel_ev = self.NUMPY_LIB.ones(muons.numevents(),
dtype=self.NUMPY_LIB.bool)
sel_mu = self.NUMPY_LIB.ones(muons.numobjects(),
dtype=self.NUMPY_LIB.bool)
sel_jet = jet.pt > 10
muons_matched_to_jet = kernels.mask_deltar_first(
self.ha,
{
"offsets": muons.offsets,
"eta": muons.eta,
"phi": muons.phi
},
sel_mu,
{
"offsets": jet.offsets,
"eta": jet.eta,
"phi": jet.phi
},
sel_jet,
0.3,
)
self.assertEqual(len(muons_matched_to_jet), muons.numobjects())
self.assertEqual(muons_matched_to_jet.sum(), 1698542)
return muons.numevents()
def test_kernel_mask_deltar_first(dataset):
muons = dataset.structs["Muon"][0]
jet = dataset.structs["Jet"][0]
sel_ev = nplib.ones(muons.numevents(), dtype=nplib.bool)
sel_mu = nplib.ones(muons.numobjects(), dtype=nplib.bool)
sel_jet = (jet.pt > 10)
muons_matched_to_jet = kernels.mask_deltar_first(
backend,
{"offsets": muons.offsets, "eta": muons.eta, "phi": muons.phi},
sel_mu,
{"offsets": jet.offsets, "eta": jet.eta, "phi": jet.phi},
sel_jet, 0.3
)