python/myutils/HistoMaker.py

import sys,os
import pickle
import ROOT 
from array import array
from printcolor import printc
from BetterConfigParser import BetterConfigParser
from TreeCache import TreeCache
from math import sqrt
from copy import copy

class HistoMaker:
    def __init__(self, samples, path, config, optionsList):
        self.path = path
        self.config = config
        self.optionsList = optionsList
        self.nBins = optionsList[0]['nBins']
        self.lumi=0.
        self.cuts = []
        for options in optionsList:
            self.cuts.append(options['cut'])
        #self.tc = TreeCache(self.cuts,samples,path) 
        self.tc = TreeCache(self.cuts,samples,path,config)
        self._rebin = False
        self.mybinning = None

    def get_histos_from_tree(self,job):
        if self.lumi == 0: 
            raise Exception("You're trying to plot with no lumi")
         
        hTreeList=[]

        #get the conversion rate in case of BDT plots
        TrainFlag = eval(self.config.get('Analysis','TrainFlag'))
        BDT_add_cut='EventForTraining == 0'


        plot_path = self.config.get('Directories','plotpath')
        addOverFlow=eval(self.config.get('Plot_general','addOverFlow'))

        # get all Histos at once
        for options in self.optionsList:
            name=job.name
            group=job.group
            treeVar=options['var']
            name=options['name']
            nBins=self.nBins
            #int(options['nBins'])
            xMin=float(options['xMin'])
            xMax=float(options['xMax'])
            weightF=options['weight']
            treeCut='%s'%(options['cut'])
            CuttedTree = self.tc.get_tree(job,treeCut)

            #options

            if job.type != 'DATA':
                if CuttedTree.GetEntries():
                    
                    if 'RTight' in treeVar or 'RMed' in treeVar: 
                        drawoption = '(%s)*(%s)'%(weightF,BDT_add_cut)
                    else: 
                        drawoption = '%s'%(weightF)
                    CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax), drawoption, "goff,e")
                    full=True
                else:
                    full=False
            elif job.type == 'DATA':
                if options['blind']:
                    if treeVar == 'H.mass':
                        CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax),treeVar+'<90. || '+treeVar + '>150.' , "goff,e")
                    else:
                        CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax),treeVar+'<0', "goff,e")

                else:
                    CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax),'1', "goff,e")
                full = True
            if full:
                hTree = ROOT.gDirectory.Get(name)
            else:
                hTree = ROOT.TH1F('%s'%name,'%s'%name,nBins,xMin,xMax)
                hTree.Sumw2()
            if job.type != 'DATA':
                if 'RTight' in treeVar or 'RMed' in treeVar:
                    if TrainFlag:
                        MC_rescale_factor=2.
                        print 'I RESCALE BY 2.0'
                    else: 
                        MC_rescale_factor = 1.
                    ScaleFactor = self.tc.get_scale(job,self.config,self.lumi)*MC_rescale_factor
                else: 
                    ScaleFactor = self.tc.get_scale(job,self.config,self.lumi)
                if ScaleFactor != 0:
                    hTree.Scale(ScaleFactor)
            #print '\t-->import %s\t Integral: %s'%(job.name,hTree.Integral())
            if addOverFlow:
                uFlow = hTree.GetBinContent(0)+hTree.GetBinContent(1)
                oFlow = hTree.GetBinContent(hTree.GetNbinsX()+1)+hTree.GetBinContent(hTree.GetNbinsX())
                uFlowErr = ROOT.TMath.Sqrt(ROOT.TMath.Power(hTree.GetBinError(0),2)+ROOT.TMath.Power(hTree.GetBinError(1),2))
                oFlowErr = ROOT.TMath.Sqrt(ROOT.TMath.Power(hTree.GetBinError(hTree.GetNbinsX()),2)+ROOT.TMath.Power(hTree.GetBinError(hTree.GetNbinsX()+1),2))
                hTree.SetBinContent(1,uFlow)
                hTree.SetBinContent(hTree.GetNbinsX(),oFlow)
                hTree.SetBinError(1,uFlowErr)
                hTree.SetBinError(hTree.GetNbinsX(),oFlowErr)
            hTree.SetDirectory(0)
            gDict = {}
            if self._rebin:
                gDict[group] = self.mybinning.rebin(hTree)
            else: 
                gDict[group] = hTree
            hTreeList.append(gDict)
        return hTreeList
       
    @property
    def rebin(self):
        return self._rebin

    @property
    def rebin(self, value):
        if self._rebin and value:
            return True
        elif self._rebin and not value:
            self.nBins = self.norebin_nBins
            self._rebin = False
        elif not self._rebin and value:
            if self.mybinning is None:
                raise Exception('define rebinning first')
            else:
                self.nBins = self.rebin_nBins
                self._rebin = True
                return True
        elif not self._rebin and not self.value:
            return False

    def calc_rebin(self, bg_list, nBins_start=1000, tolerance=0.35):
        self.norebin_nBins = self.nBins
        self.rebin_nBins = nBins_start
        self.nBins = nBins_start
        i=0
        #add all together:
        for job in bg_list:
            if not i:
                totalBG = self.get_histos_from_tree(job)[0].values()[0]
            else:
                totalBG.Add(self.get_histos_from_tree(job)[0].values()[0],1)
            i+=1
        ErrorR=0
        ErrorL=0
        TotR=0
        TotL=0
        binR=self.rebin_nBins
        binL=1
        rel=1.0
        #---- from right
        while rel > tolerance:
            TotR+=totalBG.GetBinContent(binR)
            ErrorR=sqrt(ErrorR**2+totalBG.GetBinError(binR)**2)
            binR-=1
            if not TotR == 0 and not ErrorR == 0:
                rel=ErrorR/TotR
                #print rel
        #print 'upper bin is %s'%binR

        #---- from left
        rel=1.0
        while rel > tolerance:
            TotL+=totalBG.GetBinContent(binL)
            ErrorL=sqrt(ErrorL**2+totalBG.GetBinError(binL)**2)
            binL+=1
            if not TotL == 0 and not ErrorL == 0:
                rel=ErrorL/TotL
                #print rel
        #it's the lower edge
        binL+=1
        #print 'lower bin is %s'%binL

        inbetween=binR-binL
        stepsize=int(inbetween)/(int(self.norebin_nBins)-2)
        modulo = int(inbetween)%(int(self.norebin_nBins)-2)

        #print'stepsize %s'% stepsize
        #print 'modulo %s'%modulo
        binlist=[binL]
        for i in range(0,int(self.norebin_nBins)-3):
            binlist.append(binlist[-1]+stepsize)
        binlist[-1]+=modulo
        binlist.append(binR)
        binlist.append(self.rebin_nBins+1)

        self.mybinning = Rebinner(int(self.norebin_nBins),array('d',[-1.0]+[totalBG.GetBinLowEdge(i) for i in binlist]),True)
        self._rebin = True


    @staticmethod
    def orderandadd(histo_dicts,setup):
        print histo_dicts
        ordered_histo_dict = {}
        for sample in setup:
            nSample = 0
            for histo_dict in histo_dicts:
                if histo_dict.has_key(sample):
                    if nSample == 0:
                        ordered_histo_dict[sample] = histo_dict[sample]
                    else:
                        printc('magenta','','\t--> added %s to %s'%(sample,sample))
                        ordered_histo_dict[sample].Add(histo_dict[sample])
                    nSample += 1
        return ordered_histo_dict 

class Rebinner:
    def __init__(self,nBins,lowedgearray,active=True):
        self.lowedgearray=lowedgearray
        self.nBins=nBins
        self.active=active
    def rebin(self, histo):
        if not self.active: return histo
        #print 'rebinning'
        #print histo.Integral()
        ROOT.gDirectory.Delete('hnew')
        histo.Rebin(self.nBins,'hnew',self.lowedgearray)
        binhisto=ROOT.gDirectory.Get('hnew')
        #print binhisto.Integral()
        newhisto=ROOT.TH1F('new','new',self.nBins,self.lowedgearray[0],self.lowedgearray[-1])
        newhisto.Sumw2()
        for bin in range(1,self.nBins+1):
            newhisto.SetBinContent(bin,binhisto.GetBinContent(bin))
            newhisto.SetBinError(bin,binhisto.GetBinError(bin))
        newhisto.SetName(binhisto.GetName())
        newhisto.SetTitle(binhisto.GetTitle())
        #print newhisto.Integral()
        return copy(newhisto)
Revision:	1.4
Committed:	Tue Feb 5 10:02:03 2013 UTC (12 years, 3 months ago) by peller
Content type:	text/x-python
Branch:	MAIN
Changes since 1.3:	+114 -4 lines
Log Message:	added rebinning
#	User	Rev	Content
1	nmohr	1.2	import sys,os
2	peller	1.1	import pickle
3			import ROOT
4			from array import array
5	nmohr	1.2	from printcolor import printc
6	peller	1.1	from BetterConfigParser import BetterConfigParser
7	nmohr	1.2	from TreeCache import TreeCache
8	peller	1.4	from math import sqrt
9			from copy import copy
10	peller	1.1
11			class HistoMaker:
12	nmohr	1.2	def __init__(self, samples, path, config, optionsList):
13	peller	1.1	self.path = path
14			self.config = config
15			self.optionsList = optionsList
16	peller	1.4	self.nBins = optionsList[0]['nBins']
17	peller	1.1	self.lumi=0.
18	nmohr	1.2	self.cuts = []
19			for options in optionsList:
20			self.cuts.append(options['cut'])
21	nmohr	1.3	#self.tc = TreeCache(self.cuts,samples,path)
22			self.tc = TreeCache(self.cuts,samples,path,config)
23	peller	1.4	self._rebin = False
24			self.mybinning = None
25	peller	1.1
26	nmohr	1.2	def get_histos_from_tree(self,job):
27			if self.lumi == 0:
28			raise Exception("You're trying to plot with no lumi")
29	peller	1.1
30			hTreeList=[]
31
32			#get the conversion rate in case of BDT plots
33			TrainFlag = eval(self.config.get('Analysis','TrainFlag'))
34			BDT_add_cut='EventForTraining == 0'
35
36
37			plot_path = self.config.get('Directories','plotpath')
38			addOverFlow=eval(self.config.get('Plot_general','addOverFlow'))
39
40			# get all Histos at once
41			for options in self.optionsList:
42	nmohr	1.2	name=job.name
43			group=job.group
44			treeVar=options['var']
45			name=options['name']
46	peller	1.4	nBins=self.nBins
47			#int(options['nBins'])
48	nmohr	1.2	xMin=float(options['xMin'])
49			xMax=float(options['xMax'])
50			weightF=options['weight']
51			treeCut='%s'%(options['cut'])
52			CuttedTree = self.tc.get_tree(job,treeCut)
53	peller	1.1
54			#options
55
56			if job.type != 'DATA':
57			if CuttedTree.GetEntries():
58
59	nmohr	1.2	if 'RTight' in treeVar or 'RMed' in treeVar:
60			drawoption = '(%s)*(%s)'%(weightF,BDT_add_cut)
61			else:
62			drawoption = '%s'%(weightF)
63	peller	1.1	CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax), drawoption, "goff,e")
64			full=True
65			else:
66			full=False
67			elif job.type == 'DATA':
68	nmohr	1.3	if options['blind']:
69	peller	1.1	if treeVar == 'H.mass':
70			CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax),treeVar+'<90. \|\| '+treeVar + '>150.' , "goff,e")
71			else:
72			CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax),treeVar+'<0', "goff,e")
73
74			else:
75			CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax),'1', "goff,e")
76			full = True
77			if full:
78			hTree = ROOT.gDirectory.Get(name)
79			else:
80			hTree = ROOT.TH1F('%s'%name,'%s'%name,nBins,xMin,xMax)
81			hTree.Sumw2()
82			if job.type != 'DATA':
83			if 'RTight' in treeVar or 'RMed' in treeVar:
84			if TrainFlag:
85			MC_rescale_factor=2.
86			print 'I RESCALE BY 2.0'
87	nmohr	1.2	else:
88			MC_rescale_factor = 1.
89			ScaleFactor = self.tc.get_scale(job,self.config,self.lumi)*MC_rescale_factor
90			else:
91			ScaleFactor = self.tc.get_scale(job,self.config,self.lumi)
92	peller	1.1	if ScaleFactor != 0:
93			hTree.Scale(ScaleFactor)
94			#print '\t-->import %s\t Integral: %s'%(job.name,hTree.Integral())
95			if addOverFlow:
96			uFlow = hTree.GetBinContent(0)+hTree.GetBinContent(1)
97			oFlow = hTree.GetBinContent(hTree.GetNbinsX()+1)+hTree.GetBinContent(hTree.GetNbinsX())
98			uFlowErr = ROOT.TMath.Sqrt(ROOT.TMath.Power(hTree.GetBinError(0),2)+ROOT.TMath.Power(hTree.GetBinError(1),2))
99			oFlowErr = ROOT.TMath.Sqrt(ROOT.TMath.Power(hTree.GetBinError(hTree.GetNbinsX()),2)+ROOT.TMath.Power(hTree.GetBinError(hTree.GetNbinsX()+1),2))
100			hTree.SetBinContent(1,uFlow)
101			hTree.SetBinContent(hTree.GetNbinsX(),oFlow)
102			hTree.SetBinError(1,uFlowErr)
103			hTree.SetBinError(hTree.GetNbinsX(),oFlowErr)
104			hTree.SetDirectory(0)
105	nmohr	1.3	gDict = {}
106	peller	1.4	if self._rebin:
107			gDict[group] = self.mybinning.rebin(hTree)
108			else:
109			gDict[group] = hTree
110	nmohr	1.3	hTreeList.append(gDict)
111			return hTreeList
112	peller	1.4
113			@property
114			def rebin(self):
115			return self._rebin
116
117			@property
118			def rebin(self, value):
119			if self._rebin and value:
120			return True
121			elif self._rebin and not value:
122			self.nBins = self.norebin_nBins
123			self._rebin = False
124			elif not self._rebin and value:
125			if self.mybinning is None:
126			raise Exception('define rebinning first')
127			else:
128			self.nBins = self.rebin_nBins
129			self._rebin = True
130			return True
131			elif not self._rebin and not self.value:
132			return False
133
134			def calc_rebin(self, bg_list, nBins_start=1000, tolerance=0.35):
135			self.norebin_nBins = self.nBins
136			self.rebin_nBins = nBins_start
137			self.nBins = nBins_start
138			i=0
139			#add all together:
140			for job in bg_list:
141			if not i:
142			totalBG = self.get_histos_from_tree(job)[0].values()[0]
143			else:
144			totalBG.Add(self.get_histos_from_tree(job)[0].values()[0],1)
145			i+=1
146			ErrorR=0
147			ErrorL=0
148			TotR=0
149			TotL=0
150			binR=self.rebin_nBins
151			binL=1
152			rel=1.0
153			#---- from right
154			while rel > tolerance:
155			TotR+=totalBG.GetBinContent(binR)
156			ErrorR=sqrt(ErrorR2+totalBG.GetBinError(binR)2)
157			binR-=1
158			if not TotR == 0 and not ErrorR == 0:
159			rel=ErrorR/TotR
160			#print rel
161			#print 'upper bin is %s'%binR
162
163			#---- from left
164			rel=1.0
165			while rel > tolerance:
166			TotL+=totalBG.GetBinContent(binL)
167			ErrorL=sqrt(ErrorL2+totalBG.GetBinError(binL)2)
168			binL+=1
169			if not TotL == 0 and not ErrorL == 0:
170			rel=ErrorL/TotL
171			#print rel
172			#it's the lower edge
173			binL+=1
174			#print 'lower bin is %s'%binL
175
176			inbetween=binR-binL
177			stepsize=int(inbetween)/(int(self.norebin_nBins)-2)
178			modulo = int(inbetween)%(int(self.norebin_nBins)-2)
179
180			#print'stepsize %s'% stepsize
181			#print 'modulo %s'%modulo
182			binlist=[binL]
183			for i in range(0,int(self.norebin_nBins)-3):
184			binlist.append(binlist[-1]+stepsize)
185			binlist[-1]+=modulo
186			binlist.append(binR)
187			binlist.append(self.rebin_nBins+1)
188
189			self.mybinning = Rebinner(int(self.norebin_nBins),array('d',[-1.0]+[totalBG.GetBinLowEdge(i) for i in binlist]),True)
190			self._rebin = True
191
192	peller	1.1
193	nmohr	1.3	@staticmethod
194			def orderandadd(histo_dicts,setup):
195			print histo_dicts
196			ordered_histo_dict = {}
197			for sample in setup:
198			nSample = 0
199			for histo_dict in histo_dicts:
200			if histo_dict.has_key(sample):
201			if nSample == 0:
202			ordered_histo_dict[sample] = histo_dict[sample]
203			else:
204			printc('magenta','','\t--> added %s to %s'%(sample,sample))
205			ordered_histo_dict[sample].Add(histo_dict[sample])
206			nSample += 1
207			return ordered_histo_dict
208	peller	1.4
209			class Rebinner:
210			def __init__(self,nBins,lowedgearray,active=True):
211			self.lowedgearray=lowedgearray
212			self.nBins=nBins
213			self.active=active
214			def rebin(self, histo):
215			if not self.active: return histo
216			#print 'rebinning'
217			#print histo.Integral()
218			ROOT.gDirectory.Delete('hnew')
219			histo.Rebin(self.nBins,'hnew',self.lowedgearray)
220			binhisto=ROOT.gDirectory.Get('hnew')
221			#print binhisto.Integral()
222			newhisto=ROOT.TH1F('new','new',self.nBins,self.lowedgearray[0],self.lowedgearray[-1])
223			newhisto.Sumw2()
224			for bin in range(1,self.nBins+1):
225			newhisto.SetBinContent(bin,binhisto.GetBinContent(bin))
226			newhisto.SetBinError(bin,binhisto.GetBinError(bin))
227			newhisto.SetName(binhisto.GetName())
228			newhisto.SetTitle(binhisto.GetTitle())
229			#print newhisto.Integral()
230			return copy(newhisto)