dijetmass_charhw_tb.py (raw text file)

# Validation script of HepMC: https://atlaswww.hep.anl.gov/hepsim/
# Task: reconstruct invariant mass of 2 anti-kT jets
# Comment: Run inside Jas4pp 
# author: 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
import sys

ktjet=JetN2(0.4,"antikt",100)
NMax=2 # max number of files per pT. Set to -1 to run over all files

dataset="tev13pp_mg5_chaHW_tb"
tag="" # only truth level
url=""
if len(sys.argv)>1:
   flist=FileList.get(sys.argv[1],"slcio")
   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)
else: print "Reading "+str(len(flist))+" files. Nr of files= ",NMax

# get list of files, splitting in masses
xmin=HepSim.getRanges(flist,"tb_m","_");

# reduce range
# del xmin[-1]; del xmin[-1]; 
print "Mass ranges: ",xmin

# create file list using mass ranges 
filelist=[]
histo=[]        # m_jj mass    
crosssection=[] # keeps cross sections

n=0
for i in xmin:
     ma=[]
     for j in flist:
        if (j.find("m"+str(i)+"_")>-1): ma.append(j)
     filelist.append(ma)
     hh=H1D("M="+str(i),100,0,7000)
     hh.setColor(Color(40+n*10,20+n*5,215-2*n)) 
     histo.append( hh )  
     n=n+1

print "Total number of mass ranges=",len(filelist)
Total=len(filelist)
for ipt in range(Total):
  ptlist=filelist[ipt]
  print ipt, "Mass = "+str(xmin[ipt]) 
  cross=0; xcross=0; nev=0;  Nfiles=len(ptlist)
  if (NMax>-1):  Nfiles=NMax           # limit to NMax files 
  for m in range(Nfiles):              # loop over all files in this list    
    file=FileMC(url+ptlist[m])         # get input file
    header = file.getHeader()
    un=float(header.getMomentumUnit()) # conversion units
    lunit=float(header.getLengthUnit())
    print "Read=",url+ptlist[m]
    for i in range(file.size()):
      nev=nev+1
      if (nev%200==0):
           if (m==0): print "Event=",nev
           else: print "Event=",nev," done=",int((100.0*m)/Nfiles),"%", 'x-cross=',xcross 
      eve = file.read(i)
      pa = eve.getParticles()    # particle information

      # are any leptons? If not, skip the event 
      NN=0
      for j in range(pa.getPxCount()):
         if (Math.abs(pa.getPdgId(j))==11 or Math.abs(pa.getPdgId(j))==13):
           p=LParticle(pa.getPx(j)/un,pa.getPy(j)/un,pa.getPz(j)/un,pa.getEnergy(j)/un,0)
           if (p.perp()>60): NN+=1
      if (NN==0): continue;  

      ve = eve.getEvent()        # event information
      particles=PromcUtil.getParticleDList(file.getHeader(), pa, 1, -1, 1000);
      ktjet.buildJets(particles);
      # ktjet.printJets();
      jets=ktjet.getJetsSorted()
      if (len(jets)>1): # mass of 2 leading jets 
                 ptlead1=jets[0].perp()
                 ptlead2=jets[1].perp()
                 #print "jet1=",jets[0]
                 #print "jet2=",jets[1]
                 jets[0].add(jets[1])
                 mass=jets[0].mass()
                 histo[ipt].fill(mass)

    stat = file.getStat()
    xcross=stat.getCrossSectionAccumulated()
    cross=cross+xcross;
    erro=stat.getCrossSectionErrorAccumulated();
    file.close()
  cross=cross/Nfiles
  lumi=nev/cross
  crosssection.append( histo[ipt].entries()/lumi )
  # histo[ipt].scale(1.0/lumi)   # for cross section 
  histo[ipt]=histo[ipt].getDividedByBinWidth() # get differential cross section 
  # print ipt, ") Cross section = %.3e pb"%cross, " Lumi=%.3e pb"%lumi

crossD=histo[0].sumAllBinHeights() # get total cross section 
c1 = HPlot("HepSim",600,400) # plot histogram
c1.setNameX("M_{jj} [GeV]")
# c1.setNameY("d σ / d M_{jj} [pb / GeV]")
c1.setNameY("Entries / GeV")
c1.visible(True)
c1.setAutoRange()
c1.setRangeX(0,8000)
c1.setRangeY(0,200)
c1.setMarginLeft(100)
# c1.setLogScale(1,True)
# plot for mass=4 TeV

print "Cross sections for masses [pb]"
for ipt in range(Total):
       lab="{:.1e}".format(xmin[ipt]*0.001)+" TeV "+"{:.1e} pb".format(crosssection[ipt])
       print ipt, lab
       histo[ipt].setTitle(lab)

c1.draw(histo)
# 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"
for i in histo: file.write(i)
file.close()
c1.export(name+".svg");    print name+".svg created"
# h1.toFile(name+".txt");  print name+".txt created"
# sys.exit(0)