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)
                del hTree
            else: 
                gDict[group] = hTree
            hTreeList.append(gDict)
            CuttedTree.IsA().Destructor(CuttedTree)
            del CuttedTree
        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:
            print job
            htree = self.get_histos_from_tree(job)[0].values()[0]
            if not i:
                totalBG = copy(htree)
            else:
                totalBG.Add(htree,1)
            del htree
            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):
        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].Clone()
                    else:
                        printc('magenta','','\t--> added %s to %s'%(sample,sample))
                        ordered_histo_dict[sample].Add(histo_dict[sample])
                    nSample += 1
        del histo_dicts
        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 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()
        del histo
        del binhisto
        return copy(newhisto)
Revision:	1.7
Committed:	Wed Feb 6 12:55:04 2013 UTC (12 years, 3 months ago) by peller
Content type:	text/x-python
Branch:	MAIN
Changes since 1.6:	+10 -4 lines
Log Message:	bugfix
#	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	peller	1.7	del hTree
109	peller	1.4	else:
110			gDict[group] = hTree
111	nmohr	1.3	hTreeList.append(gDict)
112	nmohr	1.6	CuttedTree.IsA().Destructor(CuttedTree)
113			del CuttedTree
114	nmohr	1.3	return hTreeList
115	peller	1.4
116			@property
117			def rebin(self):
118			return self._rebin
119
120			@property
121			def rebin(self, value):
122			if self._rebin and value:
123			return True
124			elif self._rebin and not value:
125			self.nBins = self.norebin_nBins
126			self._rebin = False
127			elif not self._rebin and value:
128			if self.mybinning is None:
129			raise Exception('define rebinning first')
130			else:
131			self.nBins = self.rebin_nBins
132			self._rebin = True
133			return True
134			elif not self._rebin and not self.value:
135			return False
136
137			def calc_rebin(self, bg_list, nBins_start=1000, tolerance=0.35):
138			self.norebin_nBins = self.nBins
139			self.rebin_nBins = nBins_start
140			self.nBins = nBins_start
141			i=0
142			#add all together:
143			for job in bg_list:
144	peller	1.7	print job
145			htree = self.get_histos_from_tree(job)[0].values()[0]
146	peller	1.4	if not i:
147	peller	1.7	totalBG = copy(htree)
148	peller	1.4	else:
149	peller	1.7	totalBG.Add(htree,1)
150			del htree
151	peller	1.4	i+=1
152			ErrorR=0
153			ErrorL=0
154			TotR=0
155			TotL=0
156			binR=self.rebin_nBins
157			binL=1
158			rel=1.0
159			#---- from right
160			while rel > tolerance:
161			TotR+=totalBG.GetBinContent(binR)
162			ErrorR=sqrt(ErrorR2+totalBG.GetBinError(binR)2)
163			binR-=1
164			if not TotR == 0 and not ErrorR == 0:
165			rel=ErrorR/TotR
166			#print rel
167			#print 'upper bin is %s'%binR
168
169			#---- from left
170			rel=1.0
171			while rel > tolerance:
172			TotL+=totalBG.GetBinContent(binL)
173			ErrorL=sqrt(ErrorL2+totalBG.GetBinError(binL)2)
174			binL+=1
175			if not TotL == 0 and not ErrorL == 0:
176			rel=ErrorL/TotL
177			#print rel
178			#it's the lower edge
179			binL+=1
180			#print 'lower bin is %s'%binL
181
182			inbetween=binR-binL
183			stepsize=int(inbetween)/(int(self.norebin_nBins)-2)
184			modulo = int(inbetween)%(int(self.norebin_nBins)-2)
185
186			#print'stepsize %s'% stepsize
187			#print 'modulo %s'%modulo
188			binlist=[binL]
189			for i in range(0,int(self.norebin_nBins)-3):
190			binlist.append(binlist[-1]+stepsize)
191			binlist[-1]+=modulo
192			binlist.append(binR)
193			binlist.append(self.rebin_nBins+1)
194
195			self.mybinning = Rebinner(int(self.norebin_nBins),array('d',[-1.0]+[totalBG.GetBinLowEdge(i) for i in binlist]),True)
196			self._rebin = True
197
198	peller	1.1
199	nmohr	1.3	@staticmethod
200			def orderandadd(histo_dicts,setup):
201			ordered_histo_dict = {}
202			for sample in setup:
203			nSample = 0
204			for histo_dict in histo_dicts:
205			if histo_dict.has_key(sample):
206			if nSample == 0:
207	peller	1.7	ordered_histo_dict[sample] = histo_dict[sample].Clone()
208	nmohr	1.3	else:
209			printc('magenta','','\t--> added %s to %s'%(sample,sample))
210			ordered_histo_dict[sample].Add(histo_dict[sample])
211			nSample += 1
212	peller	1.7	del histo_dicts
213	nmohr	1.3	return ordered_histo_dict
214	peller	1.4
215			class Rebinner:
216			def __init__(self,nBins,lowedgearray,active=True):
217			self.lowedgearray=lowedgearray
218			self.nBins=nBins
219			self.active=active
220			def rebin(self, histo):
221			if not self.active: return histo
222			#print histo.Integral()
223			ROOT.gDirectory.Delete('hnew')
224			histo.Rebin(self.nBins,'hnew',self.lowedgearray)
225			binhisto=ROOT.gDirectory.Get('hnew')
226			#print binhisto.Integral()
227			newhisto=ROOT.TH1F('new','new',self.nBins,self.lowedgearray[0],self.lowedgearray[-1])
228			newhisto.Sumw2()
229			for bin in range(1,self.nBins+1):
230			newhisto.SetBinContent(bin,binhisto.GetBinContent(bin))
231			newhisto.SetBinError(bin,binhisto.GetBinError(bin))
232			newhisto.SetName(binhisto.GetName())
233			newhisto.SetTitle(binhisto.GetTitle())
234			#print newhisto.Integral()
235	peller	1.7	del histo
236			del binhisto
237	peller	1.4	return copy(newhisto)