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