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