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