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,GroupDict=None):
        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'])
        #print self.cuts
        #self.tc = TreeCache(self.cuts,samples,path) 
        self.tc = TreeCache(self.cuts,samples,path,config)
        self._rebin = False
        self.mybinning = None
        self.GroupDict=GroupDict
        self.calc_rebin_flag = False

    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
            if self.GroupDict is None:
                group=job.group
            else:
                group=self.GroupDict[job.name]
            treeVar=options['var']
            name=options['name']
            if self._rebin or self.calc_rebin_flag:
                nBins = self.nBins
            else:
                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 eval(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: 
                #print 'not rebinning %s'%job.name 
                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.calc_rebin_flag = True
        self.norebin_nBins = copy(self.nBins)
        self.rebin_nBins = nBins_start
        self.nBins = nBins_start
        i=0
        #add all together:
        print '\n\t...calculating rebinning...'
        for job in bg_list:
            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)
        #print 'binning set to %s'%binlist
        self.mybinning = Rebinner(int(self.norebin_nBins),array('d',[-1.0]+[totalBG.GetBinLowEdge(i) for i in binlist]),True)
        self._rebin = True
        print '\t > rebinning is set <\n'

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