def test_kernel_set_get_in_offsets(self):
print("kernel_set_get_in_offsets")
dataset = self.dataset
muons = dataset.structs["Muon"][0]
arr = muons.pt.copy()
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)
inds = self.NUMPY_LIB.zeros(muons.numevents(),
dtype=self.NUMPY_LIB.int8)
#set the pt of the first muon in each event to 1
inds[:] = 0
target = self.NUMPY_LIB.ones(muons.numevents(), dtype=muons.pt.dtype)
kernels.set_in_offsets(self.ha, muons.offsets, arr, inds, target,
sel_ev, sel_mu)
print("checking set_in_offsets")
asnp = self.NUMPY_LIB.asnumpy
self.assertTrue(
verify_set_in_offsets(asnp(muons.offsets), asnp(inds), asnp(arr),
asnp(target)))
print("checking get_in_offsets")
z = kernels.get_in_offsets(self.ha, muons.offsets, arr, inds, sel_ev,
sel_mu)
self.assertTrue(
verify_get_in_offsets(asnp(muons.offsets), asnp(inds), asnp(arr),
asnp(target), asnp(z)))
return muons.numevents()
def analyze_data_function(data, parameters):
ret = Results()
num_events = data["num_events"]
muons = data["Muon"]
mu_pt = nplib.sqrt(muons.Px**2 + muons.Py**2)
muons.attrs_data["pt"] = mu_pt
mask_events = nplib.ones(muons.numevents(), dtype=nplib.bool)
mask_muons_passing_pt = muons.pt > parameters["muons_ptcut"]
num_muons_event = kernels.sum_in_offsets(backend, muons.offsets,
mask_muons_passing_pt,
mask_events, muons.masks["all"],
nplib.int8)
mask_events_dimuon = num_muons_event == 2
#get the leading muon pt in events that have exactly two muons
inds = nplib.zeros(num_events, dtype=nplib.int32)
leading_muon_pt = kernels.get_in_offsets(backend, muons.offsets, muons.pt,
inds, mask_events_dimuon,
mask_muons_passing_pt)
#compute a weighted histogram
weights = nplib.ones(num_events, dtype=nplib.float32)
bins = nplib.linspace(0, 300, 101, dtype=nplib.float32)
hist_muons_pt = Histogram(*kernels.histogram_from_vector(
backend, leading_muon_pt[mask_events_dimuon],
weights[mask_events_dimuon], bins))
#save it to the output
ret["hist_leading_muon_pt"] = hist_muons_pt
return ret
def test_kernel_get_in_offsets(dataset):
muons = dataset.structs["Muon"][0]
sel_ev = nplib.ones(muons.numevents(), dtype=nplib.bool)
sel_mu = nplib.ones(muons.numobjects(), dtype=nplib.bool)
inds = nplib.zeros(muons.numevents(), dtype=nplib.int8)
inds[:] = 0
z = kernels.get_in_offsets(backend, muons.offsets, muons.pt, inds, sel_ev,
sel_mu)
def test_kernel_get_in_offsets(self):
dataset = self.dataset
muons = dataset.structs["Muon"][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)
inds = self.NUMPY_LIB.zeros(muons.numevents(),
dtype=self.NUMPY_LIB.int8)
inds[:] = 0
z = kernels.get_in_offsets(self.ha, muons.offsets, muons.pt, inds,
sel_ev, sel_mu)
return muons.numevents()