# This program plots the transverse mass of MET and a third lepton, where the third lepton is associated with a lepton pair
# that has the same flavor, opposite charge, and closest mass to 91.2.
import math
class Processor(processor.ProcessorABC):
def __init__(self):
dataset_axis = hist.Cat("dataset", "MET and Third Lepton")
muon_axis = hist.Bin("massT", "Transverse Mass", 50, 15, 250)
self._accumulator = processor.dict_accumulator({
'massT': hist.Hist("Counts", dataset_axis, muon_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"]
# Keep track of muons and electrons by tagging them 0/1.
muons = ak.with_field(events.Muon, 0, 'flavor')
electrons = ak.with_field(events.Electron, 1, 'flavor')
MET = events.MET
output['cutflow']['all events'] += ak.size(events.MET, axis=0)
# A few reasonable muon and electron selection cuts
muons = muons[(muons.pt > 10) & (np.abs(muons.eta) < 2.4)]
electrons = electrons[(electrons.pt > 10) & (np.abs(electrons.eta) < 2.5)]
output['cutflow']['all muons'] += ak.sum(ak.count(muons, axis=1))
output['cutflow']['all electrons'] += ak.sum(ak.count(electrons, axis=1))
# Stack muons and electrons into a single array.
leptons = ak.with_name(ak.concatenate([muons, electrons], axis=1), 'PtEtaPhiMCandidate')
# Filter out events with less than 3 leptons.
trileptons = leptons[ak.num(leptons, axis=1) >= 3]
output['cutflow']['trileptons'] += ak.sum(ak.num(trileptons, axis=1))
# Generate the indices of every pair; indices because we'll be removing these elements later.
lepton_pairs = ak.argcombinations(trileptons, 2, fields=['i0', 'i1'])
# Select pairs that are SFOS.
SFOS_pairs = lepton_pairs[(trileptons[lepton_pairs['i0']].flavor == trileptons[lepton_pairs['i1']].flavor) & (trileptons[lepton_pairs['i0']].charge != trileptons[lepton_pairs['i1']].charge)]
# Find the pair with mass closest to Z.
closest_pairs = SFOS_pairs[ak.local_index(SFOS_pairs) == ak.argmin(np.abs((trileptons[SFOS_pairs['i0']] + trileptons[SFOS_pairs['i1']]).mass - 91.2), axis=1)]
# Make trileptons and closest_pairs have same shape. First, fill nones with empty arrays. Then filter out events that don't meet the pair requirement.
closest_pairs = ak.fill_none(closest_pairs, [])
closest_pairs = closest_pairs[ak.num(closest_pairs) > 0]
trileptons = trileptons[ak.num(closest_pairs) > 0]
MET = MET[ak.num(closest_pairs) > 0]
# Remove elements of the closest pairs from leptons, because we want the pt of the third lepton.
trileptons_no_pair = trileptons[(ak.local_index(trileptons) != ak.flatten(closest_pairs.i0)) & (ak.local_index(trileptons) != ak.flatten(closest_pairs.i1))]
# Find the highest-pt lepton out of the ones that remain.
leading_lepton = trileptons_no_pair[ak.argmax(trileptons_no_pair.pt, axis=1)]
output['cutflow']['number of final leading leptons'] += ak.sum(ak.num(trileptons_no_pair, axis=1))
# Cross MET with the leading lepton.
met_plus_lep = ak.cartesian({'i0': MET, 'i1': leading_lepton})
# Do some math to get what we want.
dphi_met_lep = (met_plus_lep.i0.phi - met_plus_lep.i1.phi + math.pi) % (2*math.pi) - math.pi
mt_lep = np.sqrt(2.0*met_plus_lep.i0.pt*met_plus_lep.i1.pt*(1.0-np.cos(dphi_met_lep)))
output['massT'].fill(dataset=dataset, massT=ak.flatten(mt_lep))
return output
def postprocess(self, accumulator):
return accumulator