TWTrigger.py~

#! /usr/bin/env python

###################################################################
##								 ##
## Name: TWrate.py						 ##
## Author: Kevin Nash 						 ##
## Date: 5/28/2015						 ##
## Purpose: This program creates the numerator and denominator 	 ##
##          used by TWTrigger_Maker.py to create trigger  	 ##
##          Efficiency curves.					 ##
##								 ##
###################################################################

import os
import glob
import math
from math import sqrt,exp
import ROOT
from ROOT import std,ROOT,TFile,TLorentzVector,TMath,gROOT, TF1,TH1F,TH1D,TH2F,TH2D
from ROOT import TVector
from ROOT import TFormula

import sys
from DataFormats.FWLite import Events, Handle
from optparse import OptionParser
from array import *

parser = OptionParser()

parser.add_option('-s', '--set', metavar='F', type='string', action='store',
                  default	=	'data',
                  dest		=	'set',
                  help		=	'dataset (ie data,ttbar etc)')
parser.add_option('-u', '--ptreweight', metavar='F', type='string', action='store',
                  default	=	'none',
                  dest		=	'ptreweight',
                  help		=	'on or off')

parser.add_option('-n', '--num', metavar='F', type='string', action='store',
                  default	=	'all',
                  dest		=	'num',
                  help		=	'job number')
parser.add_option('-j', '--jobs', metavar='F', type='string', action='store',
                  default	=	'1',
                  dest		=	'jobs',
                  help		=	'number of jobs')
parser.add_option('-g', '--grid', metavar='F', type='string', action='store',
                  default	=	'off',
                  dest		=	'grid',
                  help		=	'running on grid off or on')
parser.add_option('-S', '--split', metavar='F', type='string', action='store',
                  default	=	'file',
                  dest		=	'split',
                  help		=	'split by event of file')

parser.add_option('-b', '--bx', metavar='F', type='string', action='store',
                  default	=	'25ns',
                  dest		=	'bx',
                  help		=	'bunch crossing 50ns or 25ns')

parser.add_option('-t', '--tname', metavar='F', type='string', action='store',
                  default	=	'HLT_PFHT800_v2',
                  dest		=	'tname',
                  help		=	'trigger name')

parser.add_option('-p', '--pretname', metavar='F', type='string', action='store',
                  default	=	'NONE',
                  dest		=	'pretname',
                  help		=	'prescaled trigger name')


(options, args) = parser.parse_args()

tname = options.tname.split(',')
tnamestr = ''
for iname in range(0,len(tname)):
	tnamestr+=tname[iname]
	if iname!=len(tname)-1:
		tnamestr+='OR'
	
gROOT.Macro("rootlogon.C")


import Bstar_Functions	
from Bstar_Functions import *
gROOT.Macro("rootlogon.C")
#Load up cut values based on what selection we want to run 
Cuts = LoadCuts("default")
wpt = Cuts['wpt']
tpt = Cuts['tpt']
dy = Cuts['dy']
tmass = Cuts['tmass']
nsubjets = Cuts['nsubjets']
tau32 = Cuts['tau32']
minmass = Cuts['minmass']
sjbtag = Cuts['sjbtag']
wmass = Cuts['wmass']
eta1 = Cuts['eta1']
eta2 = Cuts['eta2']




print "Options summary"
print "=================="
for  opt,value in options.__dict__.items():
	#print str(option)+ ": " + str(options[option]) 
	print str(opt) +': '+ str(value)
print "=================="
print ""

#If running on the grid we access the script within a tarred directory
di = ""
if options.grid == 'on':
	di = "./tardir/"
	sys.path.insert(0, 'tardir/')


#For large datasets we need to parallelize the processing
jobs=int(options.jobs)
if jobs != 1:
	num=int(options.num)
	jobs=int(options.jobs)
	print "Running over " +str(jobs)+ " jobs"
	print "This will process job " +str(num)
else:
	print "Running over all events"


#Based on what set we want to analyze, we find all Ntuple root files 
files = Load_Ntuples(options.set,di)
jobiter = 0
# We select all the events: 
splitfiles = []

if jobs != 1 and options.split=="file":
    for ifile in range(1,len(files)+1):
    	if (ifile-1) % jobs == 0:
		jobiter+=1
	count_index = ifile  - (jobiter-1)*jobs
	if count_index==num:
		splitfiles.append(files[ifile-1])

    events = Events (splitfiles)
if options.split=="event" or jobs == 1:	  
	events = Events (files)
#events = ChainEvent(files)
#Here we load up handles and labels.
#These are used to grab entries from the Ntuples.
#To see all the current types in an Ntuple use edmDumpEventContent /PathtoNtuple/Ntuple.root

AK8HL = Initlv("jetsAK8")
	
BDiscHandle 	= 	Handle (  "vector<float>"  )
BDiscLabel  	= 	( "jetsAK8" , "jetAK8CSV")

TbitHandle 	= 	Handle (  "vector<float>"  )
TbitLabel  	= 	( "TriggerUserData" , "triggerBitTree")


TstrHandle 	= 	Handle (  "vector<string>"  )
TstrLabel  	= 	( "TriggerUserData" , "triggerNameTree")

TpsHandle 	= 	Handle (  "vector<int>"  )
TpsLabel  	= 	( "TriggerUserData" , "triggerPrescaleTree")

puHandle    	= 	Handle("int")
puLabel     	= 	( "eventUserData", "puNtrueInt" )

minmassHandle 	= 	Handle (  "vector<float> "  )
minmassLabel  	= 	( "jetsAK8" , "jetAK8minmass")

nSubjetsHandle 	= 	Handle (  "vector<float> "  )
nSubjetsLabel  	= 	( "jetsAK8" , "jetAK8nSubJets")

# not used
softDropMassHandle 	= 	Handle (  "vector<float> "  )
softDropMassLabel  	= 	( "jetsAK8" , "jetAK8softDropMass")

# for top mass
softDropMassuncorrHandle 	= 	Handle (  "vector<float> "  )
softDropMassuncorrLabel  	= 	( "jetsAK8" , "jetAK8softDropMassuncorr")

# for W mass
PrunedMassHandle 	= 	Handle (  "vector<float> "  )
PrunedMassLabel  	= 	( "jetsAK8" , "jetAK8prunedMass")

vsubjets0indexHandle 	= 	Handle (  "vector<float> "  )
vsubjets0indexLabel  	= 	( "jetsAK8" , "jetAK8vSubjetIndex0")

vsubjets1indexHandle 	= 	Handle (  "vector<float> "  )
vsubjets1indexLabel  	= 	( "jetsAK8" , "jetAK8vSubjetIndex1")

subjetsAK8CSVHandle 	= 	Handle (  "vector<float> "  )
subjetsAK8CSVLabel  	= 	( "subjetsAK8" , "subjetAK8CSV")

tau1Handle 	= 	Handle (  "vector<float> "  )
tau1Label  	= 	( "jetsAK8" , "jetAK8tau1")

tau2Handle 	= 	Handle (  "vector<float> "  )
tau2Label  	= 	( "jetsAK8" , "jetAK8tau2")

tau3Handle 	= 	Handle (  "vector<float> "  )
tau3Label  	= 	( "jetsAK8" , "jetAK8tau3")

topMassHandle 	= 	Handle (  "vector<float> "  )
topMassLabel  	= 	( "jetsAK8" , "jetAK8topMass")

subjetsCSVHandle 	= 	Handle (  "vector<float> "  )
subjetsCSVLabel  	= 	( "subjetsCmsTopTag" , "subjetCmsTopTagCSV")

subjets0indexHandle 	= 	Handle (  "vector<float> "  )
subjets0indexLabel  	= 	( "jetsAK8" , "jetAK8topSubjetIndex0")

subjets1indexHandle 	= 	Handle (  "vector<float> "  )
subjets1indexLabel  	= 	( "jetsAK8" , "jetAK8topSubjetIndex1")

subjets2indexHandle 	= 	Handle (  "vector<float> "  )
subjets2indexLabel  	= 	( "jetsAK8" , "jetAK8topSubjetIndex2")

subjets3indexHandle 	= 	Handle (  "vector<float> "  )
subjets3indexLabel  	= 	( "jetsAK8" , "jetAK8topSubjetIndex3")

#jetKeysAK8Handle 	= 	Handle (  "vector<vector<int> >"  )
#jetKeysAK8Label  	= 	( "jetKeysAK8" , "")

#subjetKeysAK8Handle 	= 	Handle (  "vector<vector<int> >"  )
#subjetKeysAK8Label  	= 	( "subjetKeysAK8" , "")

#---------------------------------------------------------------------------------------------------------------------#

#Create the output file
if jobs != 1:
	f = TFile( "TWTrigger"+options.set+tnamestr+"_pre_"+options.pretname+"_job"+options.num+"of"+options.jobs+".root", "recreate" )
else:
	f = TFile( "triggerstudies/TWTrigger"+options.set+tnamestr+"_pre_"+options.pretname+".root", "recreate" )




print "Creating histograms"

#Define Histograms
f.cd()
#---------------------------------------------------------------------------------------------------------------------#
Htpreuntrig          = TH1D("Htpreuntrig",           "",             400,  0,  4000 )
Htuntrig          = TH1D("Htuntrig",           "",             400,  0,  4000 )

Htpreuntrig.Sumw2()
Htuntrig.Sumw2()


Htpre          = TH1D("Htpre",           "",             400,  0,  4000 )
Ht          = TH1D("Ht",           "",             400,  0,  4000 )

Htpre.Sumw2()
Ht.Sumw2()

Htprebtags          = TH1D("Htprebtags",           "",             400,  0,  4000 )
Htbtags         = TH1D("Htbtags",           "",             400,  0,  4000 )

Htprebtags.Sumw2()
Htbtags.Sumw2()


Ptpre          = TH1D("Ptpre",           "",             200,  0,  2000 )
Pt          = TH1D("Pt",           "",             200,  0,  2000 )


Ptpre.Sumw2()
Pt.Sumw2()

Mpre          = TH1D("Mpre",           "",             100,  0,  200 )
M          = TH1D("M",           "",             100,  0,  200 )


Mpre.Sumw2()
M.Sumw2()


#---------------------------------------------------------------------------------------------------------------------#

# loop over events
#---------------------------------------------------------------------------------------------------------------------#

count = 0

print "Start looping"
#initialize the ttree variables
#totevents = events.size()
#print str(totevents)  +  ' Events total'
PFIRST = True
numpass = 0
trigdict = {}
prestrig = options.pretname 
for event in events:

    count	= 	count + 1

    if count % 100000 == 0 :
      print  '--------- Processing Event ' + str(count)# +'   -- percent complete ' + str(100*count/totevents) + '% -- '

   # if count > 100000 :
	#break


    #Here we split up event processing based on number of jobs 
    #This is set up to have jobs range from 1 to the total number of jobs (ie dont start at job 0)


    event.getByLabel (TstrLabel, TstrHandle)
    Tstr 		= 	TstrHandle.product() 

    event.getByLabel (TbitLabel, TbitHandle)
    Tbit 		= 	TbitHandle.product() 

    for t in range(0,len(Tstr)):
		
	
		try:
			trigdict[Tstr[t]]+=Tbit[t]
		except:
			trigdict[Tstr[t]]=0.


    AK8LV = Makelv(AK8HL,event)

	
    #print "ak8"
    if len(AK8LV)==0:
	continue

    tindex,Windex = Hemispherize(AK8LV,AK8LV)

    if jobs != 1 and options.split=="event":
    	if (count - 1) % jobs == 0:
		jobiter+=1
	count_index = count - (jobiter-1)*jobs
	if count_index!=num:
		continue 
	
    wJetsh1 = []
    wJetsh0  =  []
    topJetsh1 = []
    topJetsh0  = []

    for i in range(0,len(windex[1])):
   	wJetsh1.append(AK8LV[windex[1][i]])
    for i in range(0,len(windex[0])):
    	wJetsh0.append(AK8LV[windex[0][i]])
    for i in range(0,len(tindex[1])):
    	topJetsh1.append(AK8LV[tindex[1][i]])
    for i in range(0,len(tindex[0])):
    	topJetsh0.append(AK8LV[tindex[0][i]])
    
    wjh0 = 0
    wjh1 = 0

    #Require 1 pt>150 jet in each hemisphere (top jets already have the 150GeV requirement) 

    for wjet in wJetsh0:
	if wjet.Perp() > 200.0:
		wjh0+=1
    for wjet in wJetsh1:
	if wjet.Perp() > 200.0:
		wjh1+=1
    njets11w0 	= 	((len(topJetsh1) == 1) and (wjh0 == 1))
    njets11w1 	= 	((len(topJetsh0) == 1) and (wjh1 == 1))
    #We consider both the case that the b is the leading (highest pt) jet (hemis0) and the case where the top is the leading jet (hemis1)
    for hemis in ['hemis0']:#,'hemis1']:
    	if hemis == 'hemis0'   :
		if not njets11w0:
			continue 
		#The Ntuple entries are ordered in pt, so [0] is the highest pt entry
		#We are calling a candidate b jet (highest pt jet in hemisphere0)  

		tindexval = tindex[1][0]
		windexval = windex[0][0]

		wjet = wJetsh0[0]
		tjet = topJetsh1[0]

    	if hemis == 'hemis1'  :
		if not njets11w1:
			continue 

		tindexval = tindex[0][0]
		windexval = windex[1][0]

		wjet = wJetsh1[0]
		tjet = topJetsh0[0]
	HT = tjet.Perp() + wjet.Perp()
	PT = tjet.Perp()

        event.getByLabel (softDropMassLabel, softDropMassHandle)
        topJetMass 	= 	softDropMassHandle.product()

	# For W mass
	event.getByLabel (PrunedMassLabel, PrunedMassHandle)
	prunedJetMass 	= 	PrunedMassHandle.product()

	# For top mass
	event.getByLabel (softDropMassuncorrLabel, softDropMassuncorrHandle)
	topJetMassuncorr 	= 	softDropMassuncorrHandle.product()

	MA = topJetMassuncorr[tindexval]

	tmass_cut = tmass[0]<topJetMassuncorr[tindexval]<tmass[1]
	wmass_cut = wmass[0]<=prunedJetMass[windexval]<wmass[1]

    	event.getByLabel (TpsLabel, TpsHandle)
    	Tps 		= 	TpsHandle.product() 


	TRIGBOOL = []
	presTRIGBOOL = False
	#print "starting loops"
	for t in range(0,len(Tstr)):
		if PFIRST:
			#trigdict[Tstr[t]]=0.
			print Tstr[t]
			print "prescale " + str(Tps[t])
		#trigdict[Tstr[t]]+=Tbit[t]
		if Tstr[t]==prestrig:
			presTRIGBOOL = Tbit[t]
		for X in tname:	
			
			if Tstr[t]==X:

				if PFIRST:
					print "-----FOUND-----"
					print X
					print "prescale " + str(Tps[t])
					print "-----FOUND-----"
				
				TRIGBOOL.append(Tbit[t])
	if options.pretname == 'NONE':
		presTRIGBOOL = True
	PFIRST=False
	TPASS = False


	for TB in TRIGBOOL:
		
		if TB:
			TPASS = True
			#numpass+=1
			#print numpass
			break 	


	Htpreuntrig.Fill(HT)
	if TPASS:
		Htuntrig.Fill(HT)


	if not presTRIGBOOL:
		continue 

	Htpre.Fill(HT)
	Ptpre.Fill(PT)
	Mpre.Fill(MA)



	if TPASS:
		Ht.Fill(HT)
		Pt.Fill(PT)
		M.Fill(MA)


        event.getByLabel ( nSubjetsLabel , nSubjetsHandle )
    	nSubjets 		= 	nSubjetsHandle.product()

    	event.getByLabel (subjets0indexLabel, subjets0indexHandle)
    	subjets0index 		= 	subjets0indexHandle.product() 

    	event.getByLabel (subjets1indexLabel, subjets1indexHandle)
    	subjets1index 		= 	subjets1indexHandle.product() 

    	event.getByLabel (subjets2indexLabel, subjets2indexHandle)
    	subjets2index 		= 	subjets2indexHandle.product() 

    	event.getByLabel (subjets3indexLabel, subjets3indexHandle)
    	subjets3index 		= 	subjets3indexHandle.product()

    	event.getByLabel (subjetsCSVLabel, subjetsCSVHandle)
    	subjetsCSV 		= 	subjetsCSVHandle.product()  

	if nSubjets[tindexval]==0:
		continue 

	#SJ_csvs = [subjets0index,subjets1index,subjets2index,subjets3index]
			
	#SJ_csvvals = []

	#for icsv in range(0,int(nSubjets[tindexval])):
	#	if int(SJ_csvs[icsv][tindexval])!=-1:
	#		SJ_csvvals.append(subjetsCSV[int(SJ_csvs[icsv][tindexval])])
	#	else:
	#		SJ_csvvals.append(0.)
	#SJ_csvmax = max(SJ_csvvals)

	event.getByLabel (vsubjets0indexLabel,vsubjets0indexHandle )
	vsubjets0index 		= 	vsubjets0indexHandle.product() 

	event.getByLabel (vsubjets1indexLabel,vsubjets1indexHandle )
	vsubjets1index 		= 	vsubjets1indexHandle.product() 

	event.getByLabel (subjetsAK8CSVLabel,subjetsAK8CSVHandle )
	subjetsAK8CSV		= 	subjetsAK8CSVHandle.product() 


	if len(subjetsAK8CSV)==0:
		continue
	if len(subjetsAK8CSV)<2:
		subjetsAK8CSV[int(vsubjets0index[tindexval])]
	else:
		SJ_csvvals = [subjetsAK8CSV[int(vsubjets0index[tindexval])],subjetsAK8CSV[int(vsubjets1index[tindexval])]]


	sjbtag_cut = sjbtag[0]<SJ_csvmax<=sjbtag[1]

#   	event.getByLabel (BDiscLabel, BDiscHandle)
#    	bJetBDisc 	= 	BDiscHandle.product()
		
#       btag_cut = btag[0]<bJetBDisc[bindexval]<=btag[1]
	#For statistics
#	btag_cut=False	
#	for b in bJetBDisc:
#		if btag[0]<bJetBDisc[bindexval]<=btag[1]:
#			btag_cut=True

	if sjbtag_cut and tmass_cut: #and btag_cut?
		Htprebtags.Fill(HT)
		if TPASS:
			Htbtags.Fill(HT)

f.cd()
f.Write()
f.Close()
for key in trigdict.keys():
	print key
	print trigdict[key]
	print str(trigdict[key]/count)+'%'
	print 


print "number of events: " + str(count)