Fork me on GitHub

WorkBook/Tutorials/Student: plotStack.py

File plotStack.py, 8.9 KB (added by mselvaggi, 12 months ago)
Line 
1
2#python plotStack --h histRecoMM_M --fd out_pp_ll_zd.root --ld 'data' --fb1 out_pp_ll_sm.root --lb1 'Drell-Yan' --wb1 1.8  --rebin 2 --xmin 30. --xmax 110.  --tx 'm_{#mu#mu} [GeV]' --ty 'number of events' --out plots/test; display plots/test.pdf
3
4import argparse
5
6#_____________________________________________________________________________
7def options():
8    parser = argparse.ArgumentParser(description="generic 1D data vs Background and signal")
9
10    parser.add_argument('--h', dest='h', type=str, help='name of TH1', default='hpt')
11
12    # files ...
13    parser.add_argument('--fd', dest='fd', type=str, help='file containing pseudo data', default='histoD.root')
14    parser.add_argument('--fs', dest='fs', type=str, help='file containing signal hypothesis', default='')
15    parser.add_argument('--fb1', dest='fb1', type=str, help='file containing background hypothesis 1', default='histoB.root')
16    parser.add_argument('--fb2', dest='fb2', type=str, help='file containing background hypothesis 2', default='')
17
18    # labels
19    parser.add_argument('--ld', dest='ld', type=str, help='label for TH1 in data file to be appeared in the legend', default='')   
20    parser.add_argument('--ls', dest='ls', type=str, help='label for TH1 in signal file to be appeared in the legend', default='')   
21    parser.add_argument('--lb1', dest='lb1', type=str, help='label for TH1 in bck1 file to be appeared in the legend', default='')   
22    parser.add_argument('--lb2', dest='lb2', type=str, help='label for TH1 in bck2 file to be appeared in the legend', default='')   
23
24    # labels
25    parser.add_argument('--wd', dest='wd', type=float, help='weight for TH1 in data file to be appeared in the legend', default=1.)   
26    parser.add_argument('--ws', dest='ws', type=float, help='weight for TH1 in signal file to be appeared in the legend', default=1.)   
27    parser.add_argument('--wb1', dest='wb1', type=float, help='weight for TH1 in bck1 file to be appeared in the legend', default=1.)   
28    parser.add_argument('--wb2', dest='wb2', type=float, help='weight for TH1 in bck2 file to be appeared in the legend', default=1.)   
29
30    # general parameters
31    parser.add_argument('--tx', dest='tx', type=str, help='title of x-axis', default='p_{T}')   
32    parser.add_argument('--ty', dest='ty', type=str, help='title of y-axis', default='normalized event rate')   
33    parser.add_argument('--cap_upr', dest='cap_upr', type=str, help='caption to appear in top right corner', default='')   
34   
35    # if commas are present text is distributed among various lines
36    parser.add_argument('--cap_in', dest='cap_in', type=str, help='caption to appear in the canvas (separation with comma for multiple captions)', default='#sqrt{s} = 14 TeV, L = 10 pb^{-1}')   
37
38    parser.add_argument('--xmin', dest='xmin', type=float, help='minimum x value')   
39    parser.add_argument('--xmax', dest='xmax', type=float, help='maximum x value')   
40    parser.add_argument('--ymin', dest='ymin', type=float, help='minimum y value')   
41    parser.add_argument('--ymax', dest='ymax', type=float, help='maximum y value')   
42
43    parser.add_argument('--log', dest='log', default=False, help='plot y-axis in log scale', action='store_true')
44    parser.add_argument('--rebin', dest='rebin', type=int, help='rebin by amount specified', default=1)
45
46    #parser.add_argument('--draw_opt', dest='draw_opt', type=str, help='specifiy root draw option', default='hist')
47
48    parser.add_argument('--out', dest='out', type=str, help='output filename', default='plots/plot')
49
50    return parser.parse_args()
51
52#______________________________________________________________________________
53def main():
54   
55    ops = options()
56
57    import ROOT, sys, re, os
58    from ROOT import TFile, TH1F, TCanvas, TLegend, THStack
59       
60    # define canvas
61    canvas = ROOT.TCanvas("", "", 600, 600)
62    canvas.SetLogy(ops.log)
63    canvas.SetTicks(1,1)
64    canvas.SetLeftMargin(0.16)
65    canvas.SetRightMargin(0.08)
66    canvas.SetBottomMargin(0.12)
67    ROOT.gROOT.SetBatch(True)
68    ROOT.gStyle.SetOptStat(0000000)   
69
70
71    # define stacked histo
72    hStack = ROOT.THStack("hstack","")
73
74    draw_option = 'same hist'
75
76    legsize = 0.09
77    ylegsize = legsize
78    maxi = 0
79
80    # if option 2 is defined
81    max2 = 0
82    if ops.fb1:
83        hfileb1 = ROOT.TFile(ops.fb1)
84        hb1 = hfileb1.Get(ops.h)
85        hb1.SetTitle(ops.lb1)
86        hb1.SetLineWidth(0)
87        hb1.SetLineColor(29)
88        hb1.SetFillColor(29)
89
90        hb1.Scale(ops.wb1)
91        hb1.Rebin(ops.rebin)
92
93        max_ =  hb1.GetMaximum()
94        maxi = max(maxi, max_)
95        ylegsize += legsize
96       
97        hStack.Add(hb1)
98       
99        #hb1.Draw(draw_option)
100
101    # if option 3 is defined
102    if ops.fb2:
103        hfileb2 = ROOT.TFile(ops.fb2)
104        hb2 = hfileb2.Get(ops.h)
105        hb2.SetTitle(ops.lb2)
106        hb2.SetLineWidth(0)
107        hb2.SetLineColor(38)
108        hb2.SetFillColor(38)
109        hb2.Scale(ops.wb2)
110
111        hb2.Rebin(ops.rebin)
112
113        max_ =  hb2.GetMaximum()
114        maxi = max(maxi, max_)
115        ylegsize += legsize
116
117        #hb2.Draw(draw_option)
118        hStack.Add(hb2)
119       
120        legsize += legsize
121
122    if ops.fs:
123        hfiles = ROOT.TFile(ops.fs)
124        hs = hfiles.Get(ops.h)
125        hs.SetTitle(ops.ls)
126        hs.SetLineWidth(0)
127        hs.SetFillColor(46)
128       
129        hs.Scale(ops.ws)
130
131        hs.Rebin(ops.rebin)
132        max_ =  hs.GetMaximum()
133        maxi = max(maxi, max_)
134        ylegsize += legsize
135
136        #hs.Draw(draw_option)
137
138        hStack.Add(hs)
139        legsize += legsize
140
141    print 'pass', hStack
142
143    hStack.Draw("hist")
144
145    if ops.xmin is not None and ops.xmax is not None:
146       hStack.GetXaxis().SetRangeUser(ops.xmin, ops.xmax)
147
148    # get data files
149    hfiled = ROOT.TFile(ops.fd)
150   
151    # get and initialize histos
152    hd = hfiled.Get(ops.h)
153    hd.SetTitle(ops.ld)
154    hd.SetLineWidth(3)
155    hd.SetLineColor(ROOT.kBlack)
156    hd.SetMarkerSize(0.5)
157    hd.SetMarkerStyle(8)
158    hd.SetTitle('')
159    hd.Scale(ops.wd)
160
161    max_ =  hd.GetMaximum()
162    maxi = max(maxi, max_)
163
164    hd.Rebin(ops.rebin)
165
166    # set histogram boundaries
167    hd.Draw("same E")
168
169    if not ops.ymax:
170       if not ops.log:
171           hStack.SetMaximum(maxi*1.5)
172       else:
173           hStack.SetMaximum(maxi*100.)
174    else:
175       hStack.SetMaximum(ops.ymax)
176
177
178    if not ops.ymin:
179       if not ops.log:
180           hStack.SetMinimum(0.)
181       else:
182           print 'here'
183           hStack.SetMinimum(10.)
184    else:
185       hStack.SetMinimum(ops.ymin)
186
187    hStack.GetXaxis().SetTitle(ops.tx)
188    hStack.GetYaxis().SetTitle(ops.ty)
189    #hStack.GetYaxis().SetTitleOffset(1.75)
190    #hStack.GetXaxis().SetTitleOffset(1.40)
191    hStack.GetYaxis().SetTitleOffset(1.6)
192    hStack.GetXaxis().SetTitleOffset(1.05)
193    hStack.GetXaxis().SetTitleSize(0.045)
194    hStack.GetYaxis().SetTitleSize(0.045)
195
196    print 'set maxima and minima', ops.log
197
198    # build legend
199    leg = TLegend(0.58,0.86-legsize,0.90,0.86)
200    leg.AddEntry(hd,ops.ld,"lep")
201    if ops.fs:
202        leg.AddEntry(hs,ops.ls,"f")
203    if ops.fb1:
204        leg.AddEntry(hb1,ops.lb1,"f")
205    if ops.fb2:
206        leg.AddEntry(hb2,ops.lb2,"f")
207
208    leg.SetFillColor(0)
209    leg.SetFillStyle(0)
210    leg.SetLineColor(0)
211    leg.Draw()
212
213    print 'building legend'
214
215    # add captions
216    Text = ROOT.TLatex()
217    Text.SetNDC()
218    Text.SetTextAlign(31);
219    Text.SetTextSize(0.04)
220
221    # start with upright
222    text = '#it{' + ops.cap_upr +'}'
223    Text.DrawLatex(0.90, 0.92, text)
224
225    # now text inside
226    rt = re.split(",", ops.cap_in)
227    text = '#bf{#it{   ' + rt[0] +'}}'
228    #text = text.replace('#','\\')
229   
230    Text.SetTextAlign(22)
231    Text.SetNDC(ROOT.kTRUE)
232    #Text.SetTextSize(0.04)
233    Text.DrawLatex(0.30, 0.83, text)
234   
235    Text.SetTextAlign(22)
236    Text.SetNDC(ROOT.kTRUE)
237    text = '#bf{#it{' + rt[1] +'}}'
238    if 'ell' in text:
239       text = text.replace('#','\\')
240    Text.SetTextSize(0.03)
241    Text.DrawLatex(0.30, 0.76, text)
242
243    print rt[0]
244    print rt[1]
245
246    print 'writing text'
247
248    # this happens only if additional text
249    if len(rt)>2:
250        text = '#it{#bf{' + rt[2] +'}}'
251        if 'ell' in text:
252           text = text.replace('#','\\')
253        Text.SetTextSize(0.03)
254        Text.DrawLatex(0.30, 0.70, text)
255
256    if len(rt)>3:
257        text = '#it{#bf{' + rt[3] +'}}'
258        if 'ell' in text:
259           text = text.replace('#','\\')
260        Text.SetTextSize(0.03)
261        Text.DrawLatex(0.50, 0.60, text)
262
263    canvas.RedrawAxis()
264    canvas.Update()
265    canvas.GetFrame().SetBorderSize( 12 )
266    canvas.Modified()
267    canvas.Update()
268
269    print 'canvas updating'
270
271    pdir = os.path.dirname(ops.out)
272    name = os.path.basename(ops.out)
273    filename = pdir + '/' + name
274   
275    if not os.path.exists(pdir):
276       os.makedirs(pdir)
277   
278    #canvas.Print('{}.png'.format(filename), 'png')
279    #canvas.Print('{}.eps'.format(filename), 'eps')
280    canvas.Print('{}.pdf'.format(filename), 'pdf')
281
282#______________________________________________________________________________
283if __name__ == '__main__':
284    main()