fourbody_mass.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 
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=20000

dataset="tev13pp_pythia8_X1000toAB"
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)",70,0,1500)
hh2=H1D("M(jjll)",70,0,1500)
hh3=H1D("M(ll)",70,0,1500)

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
      for j in range(pa.getPxCount()):
         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) 

      # create isolated leptons
      leptons_iso=[]
      for lep in leptons:
         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 

      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) 



      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)         

      # 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])
                 mass=finaljets[0].mass()
                 hh1.fill(mass)
                 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  )
                              mass=finaljets[0].mass()
                              hh2.fill(mass)

    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,1500)
c1.setRangeY(0,5000)
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(hh1)
file.close()
c1.export(name+".svg");    print name+".svg created"
# h1.toFile(name+".txt");  print name+".txt created"
# sys.exit(0)