dijetmass_dm_a_mjj.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",20)
NMax=2 # max number of files per pT. Set to -1 to run over all files

dataset="tev13pp_mg5_dm_a_boson"
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)
else: print "Reading "+str(len(flist))+" files. Nr of files= ",NMax

print flist
xmin=[]
for j in flist:
   if (j.endswith("promc")):
      d1=j.split("mz_m");
      d2=d1[1].split(".")
      d3=d2[0].split("_")
      print d1, d3[0] 
      xmin.append(int(d3[0]))

xmin = list(set(xmin))

print "Mass ranges: ",xmin
histo=[]        # m_jj mass    
crosssection=[] # keeps cross sections
xmin.sort()

# debug
# xmin=[250]
n=0
for i in xmin:
     hh=H1D("M="+str(i),100,0,7000)
     hh.setColor(Color(40+n*20,20+n*5,215-2*n)) 
     histo.append( hh )  
     n=n+1

print "Total number of mass ranges=",len(xmin)
Total=0
for ipt in range(len(xmin)):
  xmass=xmin[ipt]
  xfile=""
  for jjj in flist:
       if (jjj.find("a1_mz_m"+str(xmass)+"_1.promc")>-1): xfile=jjj
  if (len(xfile)<1): continue
  
  Total=Total+1 
  print ipt, "Mass = "+str(xmass) 
  cross=0; xcross=0; nev=0; 
  print "Read=",url+xfile
  file=FileMC(url+xfile)         # get input file
  header = file.getHeader()
  un=float(header.getMomentumUnit()) # conversion units
  lunit=float(header.getLengthUnit())
  nev=0
  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

      # 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()
  cross=stat.getCrossSectionAccumulated()
  erro=stat.getCrossSectionErrorAccumulated()
  crosssection.append(cross)
  file.close()
  lumi=nev/cross
  #histo[ipt].scale(1.0/lumi)   # for cross section 
  #histo[ipt]=histo[ipt].getDividedByBinWidth() # get differential cross section 
  print "Cross section = %.3e pb"%cross, " Lumi=%.3e pb"%lumi


c1 = HPlot("HepSim",600,400) # plot histogram
c1.setNameX("M_{jj} [TeV]")
# c1.setNameY("d σ / d M_{jj} [pb / GeV]")
c1.setNameY("Events")
c1.visible(True)
c1.setAutoRange()
c1.setRangeX(0,7000)
c1.setRangeY(0,3000)
c1.setMarginLeft(100)
#c1.setLogScale(1,True)
#plot for mass=4 TeV

print "Cross sections for masses [pb]"
print crosssection

for ipt in range(Total):
       xc=float(crosssection[ipt]) 
       m=float(xmin[ipt])*0.001
       print "M=",m," TeV x-section [pb]=",xc
       lab='%.1e TeV' % m 
       lab=lab+' %.1e pb' % xc
       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.8,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)