%matplotlib inline
from coffea import hist
from coffea.analysis_objects import JaggedCandidateArray
import coffea.processor as processor
# This program will plot the MET for events which have an opposite-sign muon pair that has mass in the range of 60-120 GeV.
class Processor(processor.ProcessorABC):
def __init__(self):
dataset_axis = hist.Cat("dataset", "")
muon_axis = hist.Bin("MET", "MET [GeV]", 50, 1, 100)
self._accumulator = processor.dict_accumulator({
'MET': 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"]
muons = events.Muon
MET = events.MET.pt
output['cutflow']['all events'] += muons.size
output['cutflow']['all muons'] += muons.mass.counts.sum()
# Get all combinations of muon pairs in every event.
dimuons = muons.distincts()
# Check that pairs have opposite charge.
opposites = (dimuons.i0.charge != dimuons.i1.charge)
# Get only muons with mass between 60 and 120. Add the dimuon p4's, take mass.
limits = ((dimuons.i0 + dimuons.i1).mass >= 60) & ((dimuons.i0 + dimuons.i1).mass < 120)
# Mask the dimuons with the opposites and the limits to get dimuons with opposite charge and mass between 60 and 120 GeV.
good_dimuons = dimuons[opposites & limits]
# Mask the MET to get it only if an associated dimuon pair meeting the conditions exists.
good_MET = MET[good_dimuons.counts >= 1]
output['cutflow']['final events'] += good_MET.size
output['MET'].fill(dataset=dataset, MET=good_MET)
return output
def postprocess(self, accumulator):
return accumulator
fileset = {'SingleMu' : ["root://eospublic.cern.ch//eos/root-eos/benchmark/Run2012B_SingleMu.root"]}
from dask.distributed import Client
from coffea_casa import CoffeaCasaCluster
client = Client("tls://localhost:8786")
output = processor.run_uproot_job(fileset=fileset,
treename="Events",
processor_instance=Processor(),
executor=processor.dask_executor,
executor_args={'client': client, 'nano': True},
chunksize=250000)
[## ] | 7% Completed | 34.0s0s
distributed.comm.tcp - WARNING - Closing dangling stream in <TLS local=tls://127.0.0.1:45192 remote=tls://localhost:8786>
[########################################] | 100% Completed | 1min 39.3s
hist.plot1d(output['MET'], overlay='dataset', fill_opts={'edgecolor': (0,0,0,0.3), 'alpha': 0.8})
<matplotlib.axes._subplots.AxesSubplot at 0x7fdc262c7580>
for key, value in output['cutflow'].items():
print(key, value)
all events 53446198 all muons 74439173 final events 2970884