fourbody_rad.py (raw text file)# Validation script of HepMC: https://atlaswww.hep.anl.gov/hepsim/
# Run inside Jas4pp
# Look at 4-body invariant masses
# S.Chekanov (ANL)
from java.awt import Color,Font
from java.lang import *
from proto import FileMC
from jhplot import HPlot,P1D,HLabel,H1D,H2D
from jhplot.utils import FileList
from hepsim import HepSim
from java.util import ArrayList
from hephysics.hepsim import PromcUtil
from hephysics.jet import JetN2
from hephysics.particle import LParticle
from hephysics.jet import ParticleD
import math,sys
ktjet=JetN2(0.4,"antikt",20)
TotalEvents=20000000
dataset="tev13pp_pythia8_gkk2radion2gg"
tag="" # only truth level
url=""
if len(sys.argv)>1:
flist=FileList.get(sys.argv[1],"promc")
if (len(sys.argv)>2): NMax=int(sys.argv[2])
else:
sites=HepSim.urlRedirector(dataset)
url=sites[0]+"/"+tag
flist=HepSim.getList(url)
if len(flist)==0: print "Error: No input file!"; sys.exit(0)
n=0
hh1=H1D("M(jj)",50,0,9000)
hh2=H1D("M(jjll)",50,0,9000)
hh3=H1D("M(ll)",50,0,9000)
binsX=30
binsY=30
xminB=0
xmaxB=8000
yminB=0
ymaxB=8000
h0d= H2D("Mkk vs Mradion",binsX,xminB,xmaxB,binsY,yminB,ymaxB)
h1d= H2D("Reco vs Reco",binsX,xminB,xmaxB,binsY,yminB,ymaxB)
h2d= H2D("Reco vs Reco 1 lepton",binsX,xminB,xmaxB,binsY,yminB,ymaxB)
h3d= H2D("Reco vs Reco 2 lepton",binsX,xminB,xmaxB,binsY,yminB,ymaxB)
nev=0
Nfiles=len(flist)
Total=len(flist)
for m in range(Nfiles): # loop over all files in this list
file=FileMC(url+flist[m]) # get input file
header = file.getHeader()
un=float(header.getMomentumUnit()) # conversion units
lunit=float(header.getLengthUnit())
print "Read=",flist[m]
if (nev>TotalEvents): print "Max Nr of events are done"; break # stop after some events
for i in range(file.size()):
nev=nev+1
if (nev%200==0):
print "Event=",nev
eve = file.read(i)
pa = eve.getParticles() # particle information
ve = eve.getEvent() # event information
leptons=[]
allpart=[]
# isolaton
mKK=0
mRAD=0
for j in xrange(pa.getPxCount()):
if (abs(pa.getPdgId(j))==5100039): # graviton mass
mKK=pa.getMass(j)/un
if (abs(pa.getPdgId(j))==32): # radion mass
mRAD=pa.getMass(j)/un
if pa.getStatus(j)==1:
if (abs(pa.getPdgId(j))==12 or abs(pa.getPdgId(j))==14 or abs(pa.getPdgId(j))==15): continue # neutrino
par=LParticle(pa.getPx(j)/un,pa.getPy(j)/un,pa.getPz(j)/un,pa.getEnergy(j)/un,0)
if (par.perp()>0.1): allpart.append(par)
if (abs(pa.getPdgId(j))==11 or abs(pa.getPdgId(j))==13 ):
if (par.perp()>60 and abs(par.pseudoRapidity())<2.47): leptons.append(par)
if (mKK==0 or mRAD==0): continue
# mimic electron identification.
# we collaps several leptons into 1 if dR less than 0.1
# genearlly, correct for electrons, not muons.
leptons_reco=[]
if (len(leptons) ==1): leptons_reco.append(leptons[0])
if (len(leptons)>1):
for l1 in xrange(len(leptons)-1):
phi1=leptons[l1].phi()
eta1=leptons[l1].rapidity()
for l2 in xrange(l1+1,len(leptons)):
phi2=leptons[l2].phi()
eta2=leptons[l2].rapidity()
dR=math.sqrt( (phi1-phi2)*(phi1-phi2) + (eta1-eta1)*(eta2-eta2))
if (dR<0.1):
leptons[l1].add(leptons[l2])
leptons_reco.append(leptons[l1])
else:
leptons_reco.append(leptons[l1])
leptons_reco.append(leptons[l2])
# create isolated leptons
leptons_iso=[]
for lep in leptons_reco:
phiL=lep.phi()
etaL=lep.rapidity()
ppT=0
for pp in allpart:
phi=pp.phi()
eta=pp.rapidity()
dR=math.sqrt( (phi-phiL)*(phi-phiL) + (eta-etaL)*(eta-etaL))
if (dR<0.2): ppT=ppT+pp.perp()
if (ppT/lep.perp() < 1.1): leptons_iso.append(lep) # 90% isolation
# take events with at least 1 lepton
if len(leptons_iso)<1: continue
if (len(leptons_iso)>1):
lep1=ParticleD(leptons_iso[0].px(), leptons_iso[0].py(), leptons_iso[0].pz(), leptons_iso[0].e())
lep2=ParticleD(leptons_iso[1].px(), leptons_iso[1].py(), leptons_iso[1].pz(), leptons_iso[1].e())
lep1.add( lep2 )
mass=lep1.mass()
hh3.fill(mass)
# make jets
particles=PromcUtil.getParticleDList(file.getHeader(), pa, 1, 0, 1000);
ktjet.buildJets(particles);
jets=ktjet.getJetsSorted()
finaljets=[]
for jet in jets:
phi=jet.getPhi()
eta=jet.getRapidity()
isRemove=False
for lep in leptons_iso:
phiL=lep.phi()
etaL=lep.rapidity()
dR=math.sqrt( (phi-phiL)*(phi-phiL) + (eta-etaL)*(eta-etaL))
if (dR<0.2): isRemove=True # double counted
if (isRemove==False):
finaljets.append(jet)
# reconstructed masses
massGKK=0
massRAD=0
# 1 and 2 lepton cases
massGKK1=0
massGKK2=0
# ktjet.printJets();
if (len(finaljets)>1): # mass of 2 leading jets
ptlead1=finaljets[0].perp()
ptlead2=finaljets[1].perp()
#print "jet1=",ptlead1
#print "jet2=",ptlead2
finaljets[0].add(finaljets[1])
massRAD=finaljets[0].mass()
hh1.fill(massRAD)
if (len(leptons_iso)>1):
lep1=ParticleD(leptons_iso[0].px(), leptons_iso[0].py(), leptons_iso[0].pz(), leptons_iso[0].e())
lep2=ParticleD(leptons_iso[1].px(), leptons_iso[1].py(), leptons_iso[1].pz(), leptons_iso[1].e())
finaljets[0].add( lep1 )
finaljets[0].add( lep2 )
massGKK=finaljets[0].mass()
massGKK2=massGKK
hh2.fill(massGKK)
if (len(leptons_iso)==1):
lep1=ParticleD(leptons_iso[0].px(), leptons_iso[0].py(), leptons_iso[0].pz(), leptons_iso[0].e())
finaljets[0].add( lep1 )
massGKK=finaljets[0].mass()
massGKK1=massGKK
hh2.fill(massGKK)
#print massRAD, massGKK, mKK, mRAD, len(finaljets), len(leptons_iso)
h0d.fill(mKK,mRAD) # truth level
if (massRAD>0.5*mRAD and massRAD<1.5*mRAD):
if (massGKK>0.5*mKK and massGKK<1.5*mKK):
h1d.fill(mKK,mRAD) # to calculate efficiency
if (massGKK1>0.5*mKK and massGKK1<1.5*mKK):
h2d.fill(mKK,mRAD) # 1 lepton case
if (massGKK2>0.5*mKK and massGKK2<1.5*mKK):
h3d.fill(mKK,mRAD) # 2 lepton case
#print ""
#print mKK, massGKK
#print mRAD, massRAD
stat = file.getStat()
xcross=stat.getCrossSectionAccumulated()
erro=stat.getCrossSectionErrorAccumulated();
file.close()
crossD=hh1.sumAllBinHeights() # get total cross section
c1 = HPlot("HepSim",600,400) # plot histogram
c1.setNameX("M [GeV]")
# c1.setNameY("d σ / d M_{jj} [pb / GeV]")
c1.setNameY("Entries / GeV")
c1.visible(True)
c1.setAutoRange()
c1.setRangeX(0,10000)
c1.setRangeY(0,20000)
c1.setMarginLeft(100)
# c1.setLogScale(1,True)
# plot for mass=4 TeV
c1.draw(hh1)
hh3.setColor(Color.blue)
c1.draw(hh3)
hh2.setColor(Color.red)
c1.draw(hh2)
# h1.toTable()
#l1=HLabel("Pythia8 QCD σ=%.3e pb"%crossD, 0.5, 0.65, "NDC")
#l1.setFont(Font("Helvetica", Font.PLAIN, 14))
#c1.add(l1)
l2=HLabel("=HepSim=",0.75,0.84, "NDC")
l2.setFont(Font("Helvetica", Font.PLAIN, 14))
l2.setColor(Color.gray)
c1.add(l2)
# create file/image using name of the file
name="output"
if len(sys.argv[0])>0: name=sys.argv[0].replace(".py","")
from jhplot.io import HBook
file=HBook(name+".jdat","w"); print name+".jdat created"
file.write("mjj",hh1)
file.write("mjjll",hh2)
file.write("mll",hh3)
file.write("mkkmrad_tru",h0d)
file.write("mkkmrad_rec",h1d)
file.write("mkkmrad_rec_1lep",h2d)
file.write("mkkmrad_rec_2lep",h3d)
file.close()
c1.export(name+".svg"); print name+".svg created"
# h1.toFile(name+".txt"); print name+".txt created"
# sys.exit(0)