# This program plots the pT of the trijet system in each event with mass closest to 172.5, as well as the maximum b-tag among the three plotted jets.
class Processor(processor.ProcessorABC):
def __init__(self):
dataset_axis = hist.axis.StrCategory(name="dataset", label="", categories=[], growth=True)
Jet_axis = hist.axis.Regular(name="Jet_pt", label="Jet_pt [GeV]", bins=50, start=15, stop=200)
b_tag_axis = hist.axis.Regular(name="b_tag", label="b-tagging discriminant", bins=50, start=0, stop=1)
self.output = processor.dict_accumulator({
'Jet_pt': hist.Hist(dataset_axis, Jet_axis),
'b_tag': hist.Hist(dataset_axis, b_tag_axis),
'cutflow': processor.defaultdict_accumulator(int)
})
def process(self, events):
dataset = events.metadata["dataset"]
jets = events.Jet
# Closest calculates the distance from 172.5 of a group of masses, finds the minimum distance, then returns a Boolean array of the original input array shape with True where the minimum-distance mass is located.
def closest(masses):
delta = abs(172.5 - masses)
closest_masses = ak.min(delta, axis=1)
is_closest = (delta == closest_masses)
return is_closest
# We're going to be generating combinations of three jets - that's a lot, and cutting pt off at 30 reduces jets by half.
cut_jets = jets[jets.pt > 30]
# Get all combinations of three jets.
trijets = ak.combinations(cut_jets, 3, fields=['i0', 'i1', 'i2'])
# Get combined masses of those combinations, by adding all p4's and then taking .mass.
trijet_masses = (trijets['i0'] + trijets['i1'] + trijets['i2']).mass
# Get the masses closest to specified value (see function above)
is_closest = closest(trijet_masses)
closest_trijets = trijets[is_closest]
# Get pt of the closest trijets.
closest_pt = (closest_trijets['i0'] + closest_trijets['i1'] + closest_trijets['i2']).pt
# Get btag of the closest trijets. np.maximum(x,y) compares two arrays and gets element-wise maximums. We make two comparisons - once between the first and second jet, then between the first comparison and the third jet.
closest_btag = np.maximum(np.maximum(closest_trijets['i0'].btag, closest_trijets['i1'].btag), closest_trijets['i2'].btag)
self.output['Jet_pt'].fill(dataset=dataset, Jet_pt=ak.flatten(closest_pt))
self.output['b_tag'].fill(dataset=dataset, b_tag=ak.flatten(closest_btag))
return self.output
def postprocess(self, accumulator):
pass