Skip to content
Snippets Groups Projects
Commit 60faf480 authored by Jingjing Pan's avatar Jingjing Pan
Browse files

MC truth plots for mZd ranges 55~100 GeV

parent 997a33e3
No related branches found
No related tags found
No related merge requests found
# do the following first:
# asetup AthAnalysis,21.2.94,here
#
# https://twiki.cern.ch/twiki/bin/viewauth/AtlasProtected/AthAnalysis#How_to_inspect_files_in_pyROOT
# https://twiki.cern.ch/twiki/bin/viewauth/AtlasComputing/SoftwareTutorialxAODAnalysisInCMake#Using_PyROOT
#
# python truthPlots_pyROOT.py /direct/usatlas+u/chweber/usatlasdata/signalDAODs/mc16_13TeV.343234.MadGraphPythia8EvtGen_A14NNPDF23LO_HAHMggfZZd4l_mZd15.deriv.DAOD_HIGG2D1.e4551_e5984_a875_r9364_r9315_p3654/DAOD_HIGG2D1.16330910._000004.pool.root.1 --nEventsToProcess 100
import math
import ROOT
from array import array # to fill the TTree eventally
import argparse # to parse command line options
import collections # so we can use collections.defaultdict to more easily construct nested dicts on the fly
#import matplotlib.pyplot as plt
#import numpy as np
def calculateInvariantMass(lep1, lep2):
E = lep1.e() + lep2.e()
px = lep1.px() + lep2.px()
py = lep1.py() + lep2.py()
pz = lep1.pz() + lep2.pz()
invariantMassSq = (E**2 - px**2 - py**2 - pz**2)
invariantMass = math.copysign( abs(invariantMassSq)**0.5, invariantMassSq)
return invariantMass
def calculatePT(lep1, lep2):
px = lep1.px() + lep2.px()
py = lep1.py() + lep2.py()
pT =((px**2 + py**2)**0.5) / 1000
return pT
def calculateEta(lep1, lep2):
px = lep1.px() + lep2.px()
py = lep1.py() + lep2.py()
pz = lep1.pz() + lep2.pz()
pAbs = (px**2 + py**2 + pz**2)**0.5
eta = 0.5 * math.log( (pAbs + pz)/(pAbs - pz))
#math.atanh( pz / pAbs)
return eta
def getEtaPtPdgId(truthParticle): return { "pT" : truthParticle.pt() / 1000 , "eta" : truthParticle.eta(), "pdgId" : truthParticle.pdgId()}
# divide by 1000 to convert MeV to GeV
def propertiesOfVectorBosonAndDoughterLeptons( vectorBoson):
outputDict = collections.defaultdict(dict)
lepton1 = vectorBoson.child(0)
lepton2 = vectorBoson.child(1)
if lepton1 == None or lepton2 == None: return None
outputDict["vBoson"] = getEtaPtPdgId(vectorBoson)
outputDict["vBoson"]["m"] = vectorBoson.m() / 1000 # convert to GeV
outputDict["vBoson"]["mInv"] = calculateInvariantMass(lepton1, lepton2) / 1000 # convert to GeV
outputDict["lepton1"] = getEtaPtPdgId( lepton1 )
outputDict["lepton2"] = getEtaPtPdgId( lepton2 )
return outputDict
def getOrderedZAndZdFromHiggs(higgsBoson):
# first output is regular Z, second output is Zd
# pdgId 23 = Z boson
# pdgId 32 = Zd boson
vectorBoson1 = higgsBoson.child(0)
vectorBoson2 = higgsBoson.child(1)
if vectorBoson1.pdgId() == 23:
ZBoson = vectorBoson1 ; ZdBoson = vectorBoson2
else:
ZBoson = vectorBoson2 ; ZdBoson = vectorBoson1
return ZBoson , ZdBoson
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("input", type=str, help="name or path to the input files")
parser.add_argument( "--outputName", type=str, default="tree.root" , help = "Pick the name of the output TFile." )
parser.add_argument( "--tTreeName", type=str, default="Zd_TruthTree" , help = "Pick the name of the TTree" )
parser.add_argument( "--nEventsToProcess", type=int, default=-1 , help = "Pick the name of the TTree" )
parser.add_argument( "--mZd", type=int, default=55 , help = "mass of Zd" )
args = parser.parse_args()
print( args.input ) # print the event for debug purposes when running on condor
print( args.outputName )
print( args.tTreeName )
print( args.mZd )
mZd = args.mZd
evt = ROOT.POOL.TEvent(ROOT.POOL.TEvent.kClassAccess) #argument is optional, but ClassAccess is faster than the default POOLAccess
evt.readFrom(args.input)
tFile = ROOT.TFile(args.outputName, 'recreate') # https://wiki.physik.uzh.ch/cms/root:pyroot_ttree
TTree = ROOT.TTree(args.tTreeName, args.outputName)
# create 1 dimensional float arrays as fill variables, in this way the float
# array serves as a pointer which can be passed to the branch
# create the branches and assign the fill-variables to them as doubles (D)
Z_pT = array('f',[0]) ; TTree.Branch("Z_pT" , Z_pT , 'Z_pT/F')
Z_eta = array('f',[0]) ; TTree.Branch("Z_eta" , Z_eta , 'Z_eta/F')
Z_pdgId = array('i',[0]) ; TTree.Branch("Z_pdgId", Z_pdgId , 'Z_pdgId/I')
Z_m = array('f',[0]) ; TTree.Branch("Z_m" , Z_m , 'Z_m/F')
hist1 = ROOT.TH1F("llij_mInv","llij_mInv", 30,0,80)
hist2 = ROOT.TH1F("ll12_mInv","ll12_mInv", 30,0,80)
llij_mInv = array('f',[0]) ; TTree.Branch("llij_mInv", llij_mInv , 'llij_mInv/F')
ll12_mInv = array('f',[0]) ; TTree.Branch("ll12_mInv", ll12_mInv , 'll12_mInv/F')
candidate = array('f',[0]) ; TTree.Branch("candidate", candidate , 'candidate/F')
l1_pT = array('f',[0]) ; TTree.Branch("l1_pT" , l1_pT , 'l1_pT/F')
l1_eta = array('f',[0]) ; TTree.Branch("l1_eta" , l1_eta , 'l1_eta/F')
l1_pdgId = array('i',[0]) ; TTree.Branch("l1_pdgId", l1_pdgId, 'l1_pdgId/I')
l2_pT = array('f',[0]) ; TTree.Branch("l2_pT" , l2_pT , 'l2_pT/F')
l2_eta = array('f',[0]) ; TTree.Branch("l2_eta" , l2_eta , 'l2_eta/F')
l2_pdgId = array('i',[0]) ; TTree.Branch("l2_pdgId", l2_pdgId, 'l2_pdgId/I')
Zd_pT = array('f',[0]) ; TTree.Branch("Zd_pT" , Zd_pT , 'Zd_pT/F')
Zd_eta = array('f',[0]) ; TTree.Branch("Zd_eta" , Zd_eta , 'Zd_eta/F')
Zd_pdgId = array('i',[0]) ; TTree.Branch("Zd_pdgId", Zd_pdgId , 'Zd_pdgId/I')
Zd_m = array('f',[0]) ; TTree.Branch("Zd_m" , Zd_m , 'Zd_m/F')
ll34_mInv = array('f',[0]) ; TTree.Branch("ll34_mInv", ll34_mInv , 'll34_mInv/F')
l3_pT = array('f',[0]) ; TTree.Branch("l3_pT" , l3_pT , 'l3_pT/F')
l3_eta = array('f',[0]) ; TTree.Branch("l3_eta" , l3_eta , 'l3_eta/F')
l3_pdgId = array('i',[0]) ; TTree.Branch("l3_pdgId", l3_pdgId, 'l3_pdgId/I')
l4_pT = array('f',[0]) ; TTree.Branch("l4_pT" , l4_pT , 'l4_pT/F')
l4_eta = array('f',[0]) ; TTree.Branch("l4_eta" , l4_eta , 'l4_eta/F')
l4_pdgId = array('i',[0]) ; TTree.Branch("l4_pdgId", l4_pdgId, 'l4_pdgId/I')
# PdgIds: http://pdg.lbl.gov/2010/download/rpp-2010-JPhys-G-37-075021.pdf
# 25 = Higgs
# 23 = Z-boson
# 32 = Zd boson
# 11 = electron
# 13 = muon
if args.nEventsToProcess < 0 : nEvents = evt.getEntries()
else: nEvents = args.nEventsToProcess
#for n in range(0, nEvents):
for n in range(0, nEvents):
evt.getEntry(n) #would call this method inside a loop if you want to loop over events .. argument is the entry number
#print("the number of this event: % d" %n)
truthParticles = evt.retrieve("xAOD::TruthParticleContainer","TruthParticles")
# print out all (?) the information of the truth particles in the given event
# for p in truthParticles: ROOT.AAH.printAuxElement(p)
leptons = [ particle for particle in truthParticles if abs(particle.pdgId()) == 11 or abs(particle.pdgId()) == 13]
for i in range(0, len(leptons)-2):
# search for same-parentId, same-flavor, opposite-sign pairs
# no need to loop for the last lep in the list
# not considering the possibility that there may be more than one lepton[j] meeting the condition (the histogram shape may be off...)
lepiParent = leptons[i].parent()
if lepiParent.pdgId() == 25:
# keep this lepton if its parent is Z* while MISLABELED as higgs
for j in range(i+1, len(leptons)-1):
# no need to go over the leps before (including) lep[i]
# go until the last lep
lepjParent = leptons[j].parent()
if lepiParent.pdgId() == lepjParent.pdgId() and leptons[i].pdgId() == -(leptons[j].pdgId()):
llij_mInv[0] = calculateInvariantMass(leptons[i], leptons[j])/ 1000 #convert to GeV
hist1.Fill(llij_mInv[0])
Z_pT[0] = calculatePT(leptons[i], leptons[j])
Z_eta[0] = calculateEta(leptons[i], leptons[j])
break
# found a same-parentId, same-flavor, opposite-sign pair
# recorded for event n
#print("number of leptons: % d" %len(leptons))
#print("index i: % d, index j: % d" %(i, j))
#import pdb; pdb.set_trace() # import the debugger and instruct it to stop here
for particle in truthParticles:
if particle.pdgId() != 32: continue
# we found the Zd only if particle.pdgId() == 32
higgs = particle.parent()
if higgs == None: break # sometimes the truth associations between the truthHiggs and its vector children seem to be broken. Just skip the event in this case.
ZBoson, darkZBoson = getOrderedZAndZdFromHiggs(higgs)
ZProperties = propertiesOfVectorBosonAndDoughterLeptons( ZBoson)
darkZProperties = propertiesOfVectorBosonAndDoughterLeptons( darkZBoson)
if ZProperties == None or darkZProperties == None: break
candidate[0] = ZProperties["vBoson"]["mInv"]
if candidate[0] != mZd and candidate[0] > 54.5:
ll12_mInv[0] = candidate[0] # sometimes the truth associations break, then we discard the event
hist2.Fill(ll12_mInv[0])
#import pdb; pdb.set_trace() # import the debugger and instruct it to stop here
# Z_pT[0] = ZProperties["vBoson"]["pT"]
# Z_eta[0] = ZProperties["vBoson"]["eta"]
# Z_pdgId[0] = ZProperties["vBoson"]["pdgId"]
# Z_m[0] = ZProperties["vBoson"]["m"]
#ll12_mInv[0] = ZProperties["vBoson"]["mInv"]
# l1_pT[0] = ZProperties["lepton1"]["pT"]
# l1_eta[0] = ZProperties["lepton1"]["eta"]
# l1_pdgId[0] = ZProperties["lepton1"]["pdgId"]
# l2_pT[0] = ZProperties["lepton2"]["pT"]
# l2_eta[0] = ZProperties["lepton2"]["eta"]
# l2_pdgId[0] = ZProperties["lepton2"]["pdgId"]
Zd_pT[0] = darkZProperties["vBoson"]["pT"]
Zd_eta[0] = darkZProperties["vBoson"]["eta"]
Zd_pdgId[0] = darkZProperties["vBoson"]["pdgId"]
Zd_m[0] = darkZProperties["vBoson"]["m"]
ll34_mInv[0] = darkZProperties["vBoson"]["mInv"]
l3_pT[0] = darkZProperties["lepton1"]["pT"]
l3_eta[0] = darkZProperties["lepton1"]["eta"]
l3_pdgId[0] = darkZProperties["lepton1"]["pdgId"]
l4_pT[0] = darkZProperties["lepton2"]["pT"]
l4_eta[0] = darkZProperties["lepton2"]["eta"]
l4_pdgId[0] = darkZProperties["lepton2"]["pdgId"]
#import pdb; pdb.set_trace() # import the debugger and instruct it to stop here
TTree.Fill()
break # once we found the Z and Zd no need to go over the other truthParticles, better go on with the next event
#import pdb; pdb.set_trace() # import the debugger and instruct it to stop here
hist1.SetBinContent(1,100)
hist2.SetBinContent(1,100)
# change line color and style, and width
# see here: https://root.cern.ch/doc/master/classTAttLine.html
# scroll down here for color wheel: https://root.cern.ch/doc/master/classTColor.html
hist1.SetLineColor(ROOT.kOrange )
hist2.SetLineColor(ROOT.kSpring )
#hist1.SetFillColor(ROOT.kOrange)
#hist2.SetFillColor(ROOT.kSpring)
hist1.SetLineStyle(ROOT.kSolid )
hist1.SetLineWidth( 3 )
hist2.SetLineStyle(ROOT.kDashed )
hist2.SetLineWidth( 3 )
# scale the y-axis range of the first hist to the maximum of any of the hists
histMaximumValues = []
for hist in [hist1, hist2]: histMaximumValues.append( hist.GetMaximum())
hist1.GetYaxis().SetRangeUser(0, max(histMaximumValues) * 1.1)
# setup canvas to 'draw' the histograms on.
# When doing hist.Draw() one is usually implicity created, but we need to create one explicitly here, so that we can use it's 'update' method later
tCanvas = ROOT.TCanvas( "testCanvas","testCanvas", 1280,720)
hist1.Draw()
hist2.Draw("same") # need to draw with 'same' so that the second 'draw' does not overwrite the first one
# see https://root.cern/doc/master/classTHistPainter.html for additional draw options
#setup legend
xOffset = 0.5; yOffset = 0.4
xWidth = 0.4; ywidth = 0.3
TLegend = ROOT.TLegend(xOffset, yOffset ,xOffset + xWidth, yOffset+ ywidth)
TLegend.SetFillColor(ROOT.kWhite)
TLegend.SetLineColor(ROOT.kWhite)
TLegend.SetNColumns(1);
TLegend.SetFillStyle(0); # make legend background transparent
TLegend.SetBorderSize(0); # and remove its border without a border
# Add legend entries
TLegend.AddEntry(hist1 , "mInv of Z", "l");
TLegend.AddEntry(hist2 , "mInv of Z*", "l");
# need to draw the legend
TLegend.Draw()
# update the canvas so that it shows the things we added since we created the tCanvas
tCanvas.Update()
# print the TCanvas as PDF, PNG, and ROOT file, though don't forget that we can weite the tCanvas to a root file too with tCanvas.Write()
tCanvas.Print("mInva_mZd.pdf")
tCanvas.Print("mInva_mZd.png")
tCanvas.Print("mInva_mZd.root")
#import pdb; pdb.set_trace() # import the debugger and instruct it to stop here
tFile.Write()
tFile.Close()
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment