# 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.Cat("dataset", "")
Jet_axis = hist.Bin("Jet_pt", "Jet [GeV]", 50, 15, 200)
b_tag_axis = hist.Bin("b_tag", "b-tagging discriminant", 50, 0, 1)
self._accumulator = processor.dict_accumulator({
'Jet_pt': hist.Hist("Counts", dataset_axis, Jet_axis),
'b_tag': hist.Hist("Counts", dataset_axis, b_tag_axis),
'cutflow': processor.defaultdict_accumulator(int)
})
@property
def accumulator(self):
return self._accumulator
def process(self, events):
output = self.accumulator.identity()
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 = delta.min()
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 = cut_jets.choose(3)
# 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)
output['Jet_pt'].fill(dataset=dataset, Jet_pt=closest_pt.flatten())
output['b_tag'].fill(dataset=dataset, b_tag=closest_btag.flatten())
return output
def postprocess(self, accumulator):
return accumulator