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

    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:
                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 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:
        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)

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