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
        VHbbNameSpace=config.get('VHbbNameSpace','library')
        ROOT.gSystem.Load(VHbbNameSpace)

    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
        CuttedTree = self.tc.get_tree(job,'1')
        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'])

            #options

            if job.type != 'DATA':
                if CuttedTree.GetEntries():
                    if 'RTight' in treeVar or 'RMed' in treeVar: 
                        drawoption = '(%s)*(%s & %s)'%(weightF,treeCut,BDT_add_cut)
                    else: 
                        drawoption = '(%s)*(%s)'%(weightF,treeCut)
                    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. & %s' %treeCut, "goff,e")
                    else:
                        CuttedTree.Draw('%s>>%s(%s,%s,%s)' %(treeVar,name,nBins,xMin,xMax),treeVar+'<0 & %s'%treeCut, "goff,e")

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