import numpy as np
%matplotlib inline
import hist
import coffea.processor as processor
import awkward as ak
from coffea.nanoevents import schemas
# 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.axis.StrCategory(name="dataset", label="", categories=[], growth=True)
muon_axis = hist.axis.Regular(name="massT", label="Transverse Mass [GeV]", bins=50, start=15, stop=250)
self.output = processor.dict_accumulator({
'massT': hist.Hist(dataset_axis, muon_axis),
'cutflow': processor.defaultdict_accumulator(int)
})
def process(self, events):
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
self.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)]
self.output['cutflow']['all muons'] += ak.sum(ak.count(muons, axis=1))
self.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]
self.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, [], axis=0)
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)]
self.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)))
self.output['massT'].fill(dataset=dataset, massT=ak.flatten(mt_lep))
return self.output
def postprocess(self, accumulator):
pass
from dask.distributed import Client
client = Client("tls://localhost:8786")
fileset = {'SingleMu' : ["root://eospublic.cern.ch//eos/root-eos/benchmark/Run2012B_SingleMu.root"]}
executor = processor.DaskExecutor(client=client)
run = processor.Runner(executor=executor,
schema=schemas.NanoAODSchema,
savemetrics=True
)
output, metrics = run(fileset, "Events", processor_instance=Processor())
metrics
[########################################] | 100% Completed | 3min 35.1s
{'bytesread': 9982563224, 'columns': ['Muon_pt', 'Electron_charge', 'MET_CovXY', 'Muon_phi', 'Muon_pfRelIso04_all', 'nMuon', 'Muon_charge', 'Muon_mass', 'nJet', 'Muon_softId', 'Muon_dzErr', 'Electron_dz', 'Muon_pfRelIso03_all', 'Electron_pfRelIso03_all', 'Muon_dz', 'Electron_dzErr', 'MET_pt', 'Electron_dxy', 'MET_CovXX', 'MET_phi', 'Muon_tightId', 'Electron_mass', 'Electron_phi', 'Electron_jetIdx', 'Electron_dxyErr', 'nElectron', 'Muon_dxy', 'Electron_eta', 'Electron_genPartIdx', 'MET_sumet', 'Electron_cutBasedId', 'Muon_eta', 'Muon_dxyErr', 'MET_significance', 'MET_CovYY', 'Electron_pt', 'Electron_pfId', 'Muon_jetIdx', 'Muon_genPartIdx'], 'entries': 53446198, 'processtime': 8462.17432808876, 'chunks': 534}
output['massT'].plot1d()
[StairsArtists(stairs=<matplotlib.patches.StepPatch object at 0x7f3be3fb4b50>, errorbar=None, legend_artist=None)]
for key, value in output['cutflow'].items():
print(key, value)
all events 53446198 all muons 991421725 all electrons 87391040 trileptons 5043516 number of final leading leptons 1227755