Fork me on GitHub

WorkBook/Tutorials/Student: plotStack.py

File plotStack.py, 8.9 KB (added by Michele Selvaggi, 5 years 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()