analysis_INTT/plot/plotUtil.py

0001 from ROOT import TH1F, TH2F, TFile, TCanvas, TLegend, TColor, gROOT, gSystem, TPad, gPad, gStyle, kHAlignRight, kVAlignBottom
0002
0003 TickSize = 0.03
0004 AxisTitleSize = 0.05
0005 AxisLabelSize = 0.04
0006 LeftMargin = 0.15
0007 RightMargin = 0.08
0008 TopMargin = 0.08
0009 BottomMargin = 0.13
0010
0011 black_hex = '#1B1A17'
0012 red_hex = '#9A031E'
0013 blue_hex = '#0B60B0'
0014 green_hex = '#186F65'
0015 orange_hex = '#e99960'
0016 purple_hex = '#7F167F'
0017 pink_hex = '#FFC0CB'
0018 yellow_hex = '#ffcc66'
0019 cyan_hex = '#7FE9DE'
0020
0021 def markerset(i):
0022     if i == 1:
0023         return [20]
0024     elif i == 2:
0025         return [20, 21]
0026     elif i == 3:
0027         return [20, 21, 33]
0028     elif i == 4:
0029         return [20, 21, 33, 34]
0030     elif i == 5:
0031         return [20, 21, 33, 34, 47]
0032     elif i == 6:
0033         return [20, 21, 33, 34, 47, 43]
0034     elif i == 7:
0035         return [20, 21, 33, 34, 47, 43, 45]
0036     elif i == 8:
0037         return [20, 21, 33, 34, 47, 43, 45, 49]
0038     else:
0039         print ("Attempt to use more than 8 markers, probably too many histograms. Set all markers to 20")
0040         return [20 for _ in range(i)]
0041
0042 def colorset(i):
0043     if i == 1:
0044         return ['#810000']
0045     elif i == 2:
0046         return ['#810000', '#0F4C75']
0047     elif i == 3:
0048         return ['#810000', '#0F4C75', '#7F167F']
0049     elif i == 4:
0050         return ['#810000', '#0F4C75', '#7F167F', '#5E8B7E']
0051     elif i == 5:
0052         return ['#810000', '#0F4C75', '#7F167F', '#5E8B7E', '#e99960']
0053     elif i == 6:
0054         return ['#810000', '#0F4C75', '#7F167F', '#5E8B7E', '#e99960', '#FFC0CB']
0055     elif i == 7:
0056         return ['#810000', '#0F4C75', '#7F167F', '#5E8B7E', '#e99960', '#FFC0CB', '#ffcc66']
0057     elif i == 8:
0058         return ['#810000', '#0F4C75', '#7F167F', '#5E8B7E', '#e99960', '#FFC0CB', '#ffcc66', '#7FE9DE']
0059     else:
0060         print ("Attempt to use more than 5 colors")
0061         return ['#810000', '#0F4C75', '#7F167F', '#5E8B7E', '#e99960', '#FFC0CB', '#ffcc66', '#7FE9DE']
0062
0063
0064 def colorset2(i):
0065     if i == 1:
0066         return ['#f2777a']
0067     elif i == 2:
0068         return ['#f2777a', '#6699cc']
0069     elif i == 3:
0070         return ['#f2777a', '#6699cc', '#9999cc']
0071     elif i == 4:
0072         return ['#f2777a', '#6699cc', '#9999cc', '#99cc99']
0073     elif i == 5:
0074         return ['#f2777a', '#6699cc', '#9999cc', '#99cc99', '#e99960']
0075     elif i == 6:
0076         return ['#f2777a', '#6699cc', '#9999cc', '#99cc99', '#e99960', '#FFC0CB']
0077     elif i == 7:
0078         return ['#f2777a', '#6699cc', '#9999cc', '#99cc99', '#e99960', '#FFC0CB', '#ffcc66']
0079     elif i == 8:
0080         return ['#f2777a', '#6699cc', '#9999cc', '#99cc99', '#e99960', '#FFC0CB', '#ffcc66', '#7FE9DE']
0081     else:
0082         print ("Attempt to use more than 5 colors")
0083         return ['#f2777a', '#6699cc', '#9999cc', '#99cc99', '#e99960', '#FFC0CB', '#ffcc66', '#7FE9DE']
0084
0085
0086 def Draw_1Dhist(hist, IsData, norm1, logy, ymaxscale, XaxisName, Ytitle_unit, outname):
0087     hist.Sumw2()
0088     binwidth = hist.GetXaxis().GetBinWidth(1)
0089     c = TCanvas('c', 'c', 800, 700)
0090     if norm1:
0091         hist.Scale(1. / hist.Integral(-1, -1))
0092     if logy:
0093         c.SetLogy()
0094     c.cd()
0095     gPad.SetRightMargin(RightMargin)
0096     gPad.SetTopMargin(TopMargin)
0097     gPad.SetLeftMargin(LeftMargin)
0098     gPad.SetBottomMargin(BottomMargin)
0099     if norm1:
0100         if Ytitle_unit == '':
0101             hist.GetYaxis().SetTitle(
0102                 'Normalized entries / ({:g})'.format(binwidth))
0103         else:
0104             hist.GetYaxis().SetTitle(
0105                 'Normalized entries / ({:g} {unit})'.format(binwidth, unit=Ytitle_unit))
0106     else:
0107         if Ytitle_unit == '':
0108             hist.GetYaxis().SetTitle('Entries / ({:g})'.format(binwidth))
0109         else:
0110             hist.GetYaxis().SetTitle(
0111                 'Entries / ({:g} {unit})'.format(binwidth, unit=Ytitle_unit))
0112
0113     # hist.GetXaxis().SetRangeUser(hist.GetBinLowEdge(1)-binwidth, hist.GetBinLowEdge(hist.GetNbinsX())+2*binwidth)
0114     if logy:
0115         hist.GetYaxis().SetRangeUser(hist.GetMinimum(0)*0.5, (hist.GetMaximum()) * ymaxscale)
0116     else:
0117         hist.GetYaxis().SetRangeUser(0., (hist.GetMaximum()) * ymaxscale)
0118     hist.GetXaxis().SetTitle(XaxisName)
0119     hist.GetXaxis().SetTickSize(TickSize)
0120     hist.GetXaxis().SetTitleSize(AxisTitleSize)
0121     hist.GetXaxis().SetLabelSize(AxisLabelSize)
0122     hist.GetYaxis().SetTickSize(TickSize)
0123     hist.GetYaxis().SetTitleSize(AxisTitleSize)
0124     hist.GetYaxis().SetLabelSize(AxisLabelSize)
0125     hist.GetXaxis().SetTitleOffset(1.1)
0126     hist.GetYaxis().SetTitleOffset(1.35)
0127     hist.SetLineColor(1)
0128     hist.SetLineWidth(2)
0129     hist.Draw('hist')
0130     leg = TLegend((1-RightMargin)-0.5, (1-TopMargin)-0.13,
0131                   (1-RightMargin)-0.1, (1-TopMargin)-0.03)
0132     leg.SetTextSize(0.04)
0133     leg.SetFillStyle(0)
0134     if IsData:
0135         leg.AddEntry("", "#it{#bf{sPHENIX}} Internal", "")
0136         leg.AddEntry("", "Au+Au #sqrt{s_{NN}}=200 GeV", "")
0137     else:
0138         leg.AddEntry("", "#it{#bf{sPHENIX}} Simulation", "")
0139         leg.AddEntry("", "Au+Au #sqrt{s_{NN}}=200 GeV", "")
0140     leg.Draw()
0141     c.RedrawAxis()
0142     c.Draw()
0143     c.SaveAs(outname+'.pdf')
0144     c.SaveAs(outname+'.png')
0145     if(c):
0146         c.Close()
0147         gSystem.ProcessEvents()
0148         del c
0149         c = 0
0150
0151
0152 def Draw_1DhistsComp(lhist, norm1, logx, logy, ymaxscale, XaxisName, Ytitle_unit, outname):
0153     color = ['#1B1A17', '#035397', '#9B0000']
0154     legtext = ['1st+2nd layers', '2nd+3rd layers', '1st+3rd layers']
0155     ymax = -1
0156     ymin = 10e10
0157     for h in lhist:
0158         h.Sumw2()
0159         if h.GetMaximum() > ymax:
0160             ymax = h.GetMaximum()
0161         if h.GetMinimum(0) < ymin:
0162             ymin = h.GetMinimum(0)
0163         if norm1:
0164             h.Scale(1. / h.Integral(-1, -1))
0165             ymax = h.GetMaximum()
0166             ymin = h.GetMinimum(0)
0167
0168     binwidth_bin1 = lhist[0].GetXaxis().GetBinWidth(1)
0169     binwidth_bin2 = lhist[0].GetXaxis().GetBinWidth(2)
0170     printbinwidth = True
0171     if binwidth_bin1 != binwidth_bin2:
0172         printbinwidth = False
0173
0174     c = TCanvas('c', 'c', 800, 700)
0175     if logx:
0176         c.SetLogx()
0177     if logy:
0178         c.SetLogy()
0179     c.cd()
0180     gPad.SetRightMargin(RightMargin)
0181     gPad.SetTopMargin(TopMargin)
0182     gPad.SetLeftMargin(LeftMargin)
0183     gPad.SetBottomMargin(BottomMargin)
0184     if printbinwidth:
0185         if norm1:
0186             if Ytitle_unit == '':
0187                 lhist[0].GetYaxis().SetTitle('Normalized entries / ({:g})'.format(binwidth_bin1))
0188             else:
0189                 lhist[0].GetYaxis().SetTitle('Normalized entries / ({:g} {unit})'.format(binwidth_bin1, unit=Ytitle_unit))
0190         else:
0191             if Ytitle_unit == '':
0192                 lhist[0].GetYaxis().SetTitle('Entries / ({:g})'.format(binwidth_bin1))
0193             else:
0194                 lhist[0].GetYaxis().SetTitle('Entries / ({:g} {unit})'.format(binwidth_bin1, unit=Ytitle_unit))
0195     else:
0196         if norm1:
0197             if Ytitle_unit == '':
0198                 lhist[0].GetYaxis().SetTitle('Normalized entries')
0199             else:
0200                 lhist[0].GetYaxis().SetTitle('Normalized entries {unit})'.format(unit=Ytitle_unit))
0201         else:
0202             if Ytitle_unit == '':
0203                 lhist[0].GetYaxis().SetTitle('Entries')
0204             else:
0205                 lhist[0].GetYaxis().SetTitle('Entries {unit}'.format(unit=Ytitle_unit))
0206         # lhist[0].GetXaxis().SetRangeUser(lhist[0].GetBinLowEdge(1)-binwidth, lhist[0].GetBinLowEdge(lhist[0].GetNbinsX())+2*binwidth)
0207     if logy:
0208         lhist[0].GetYaxis().SetRangeUser(ymin * 0.05, ymax * 100)
0209     else:
0210         lhist[0].GetYaxis().SetRangeUser(0., ymax * ymaxscale)
0211     lhist[0].GetXaxis().SetTitle(XaxisName)
0212     lhist[0].GetXaxis().SetTickSize(TickSize)
0213     lhist[0].GetXaxis().SetTitleSize(AxisTitleSize)
0214     lhist[0].GetXaxis().SetLabelSize(AxisLabelSize)
0215     lhist[0].GetYaxis().SetTickSize(TickSize)
0216     lhist[0].GetYaxis().SetTitleSize(AxisTitleSize)
0217     lhist[0].GetYaxis().SetLabelSize(AxisLabelSize)
0218     lhist[0].GetXaxis().SetTitleOffset(1.1)
0219     lhist[0].GetYaxis().SetTitleOffset(1.4)
0220     for i, h in enumerate(lhist):
0221         if i == 0:
0222             h.SetLineColor(TColor.GetColor(color[i]))
0223             h.SetLineWidth(2)
0224             h.Draw('hist')
0225         else:
0226             h.SetLineColor(TColor.GetColor(color[i]))
0227             h.SetLineWidth(2)
0228             h.Draw('histsame')
0229
0230     leg = TLegend((1-RightMargin)-0.45, (1-TopMargin)-0.15,
0231                   (1-RightMargin)-0.1, (1-TopMargin)-0.03)
0232     leg.SetTextSize(0.045)
0233     leg.SetFillStyle(0)
0234     leg.AddEntry("", "#it{#bf{sPHENIX}} Simulation", "")
0235     leg.AddEntry("", "Au+Au #sqrt{s_{NN}}=200 GeV", "")
0236     leg.Draw()
0237
0238     leg1 = TLegend(LeftMargin+0.04, (1-TopMargin)-0.21,
0239                    LeftMargin+0.34, (1-TopMargin)-0.03)
0240     leg1.SetTextSize(0.035)
0241     leg1.SetFillStyle(0)
0242     for i, h in enumerate(lhist):
0243         leg1.AddEntry(h, legtext[i], "l")
0244     leg1.Draw()
0245     c.RedrawAxis()
0246     c.Draw()
0247     c.SaveAs(outname+'.pdf')
0248     c.SaveAs(outname+'.png')
0249     if(c):
0250         c.Close()
0251         gSystem.ProcessEvents()
0252         del c
0253         c = 0
0254
0255
0256 def Draw_2Dhist(hist, IsData, logz, norm1, rmargin, XaxisName, YaxisName, ZaxisName, drawopt, outname):
0257     c = TCanvas('c', 'c', 800, 700)
0258     if logz:
0259         c.SetLogz()
0260     c.cd()
0261     if ZaxisName == '':
0262         gPad.SetRightMargin(rmargin)
0263     else:
0264         gPad.SetRightMargin(rmargin+0.03)
0265
0266     gPad.SetTopMargin(TopMargin)
0267     gPad.SetLeftMargin(LeftMargin)
0268     gPad.SetBottomMargin(BottomMargin)
0269     if norm1:
0270         hist.Scale(1. / hist.Integral(-1, -1, -1, -1))
0271     hist.GetXaxis().SetTitle(XaxisName)
0272     hist.GetYaxis().SetTitle(YaxisName)
0273     hist.GetXaxis().SetTickSize(TickSize)
0274     hist.GetYaxis().SetTickSize(TickSize)
0275     hist.GetXaxis().SetTitleSize(AxisTitleSize)
0276     hist.GetYaxis().SetTitleSize(AxisTitleSize)
0277     hist.GetXaxis().SetLabelSize(AxisLabelSize)
0278     hist.GetYaxis().SetLabelSize(AxisLabelSize)
0279     if ZaxisName != '':
0280         hist.GetZaxis().SetTitle(ZaxisName)
0281         hist.GetZaxis().SetTitleSize(AxisTitleSize)
0282         hist.GetZaxis().SetTitleOffset(1.4)
0283
0284     hist.GetXaxis().SetTitleOffset(1.1)
0285     hist.GetYaxis().SetTitleOffset(1.3)
0286     hist.GetZaxis().SetLabelSize(AxisLabelSize)
0287     hist.SetContour(1000)
0288     hist.Draw(drawopt)
0289
0290     rightshift = 0.1
0291     leg = TLegend((1-RightMargin)-0.5, (1-TopMargin)+0.01, 1-gPad.GetRightMargin(), (1-TopMargin)+0.04)
0292     leg.SetTextAlign(kHAlignRight+kVAlignBottom)
0293     leg.SetTextSize(0.045)
0294     leg.SetFillStyle(0)
0295     if IsData:
0296         leg.AddEntry("", "#it{#bf{sPHENIX}} Internal", "")
0297         # leg.AddEntry("", "Au+Au #sqrt{s_{NN}}=200 GeV", "")
0298     else:
0299         leg.AddEntry("", "#it{#bf{sPHENIX}} Simulation", "")
0300         # leg.AddEntry("", "Au+Au #sqrt{s_{NN}}=200 GeV", "")
0301     leg.Draw()
0302     c.RedrawAxis()
0303     c.Draw()
0304     c.SaveAs(outname+'.pdf')
0305     c.SaveAs(outname+'.png')
0306     if(c):
0307         c.Close()
0308         gSystem.ProcessEvents()
0309         del c
0310         c = 0
0311
0312 def plot_Stack(totalhist, list_hist, color, list_legtext, logy, ymaxscale, XaxisName, Ytitle_unit, plotname):
0313     # color = ['#073b4c','#118ab2','#06d6a0','#ffd166','#ef476f']
0314     binwidth = totalhist.GetXaxis().GetBinWidth(1)
0315     hs = THStack('hs','hs');
0316     for i, hist in enumerate(list_hist):
0317         hist.SetLineColor(1)
0318         hist.SetLineWidth(1)
0319         hist.SetFillColor(TColor.GetColor(color[i]))
0320         hs.Add(hist)
0321
0322     # print(totalhist.GetMaximum(), hs.GetMaximum())
0323
0324     c = TCanvas('c', 'c', 800, 700)
0325     if logy:
0326         c.SetLogy()
0327     c.cd()
0328     hs.Draw()
0329     hs.GetXaxis().SetTitle(XaxisName)
0330     if Ytitle_unit == '':
0331         hist.GetYaxis().SetTitle('Entries / ({:g})'.format(binwidth))
0332     else:
0333         hist.GetYaxis().SetTitle('Entries / ({:g} {unit})'.format(binwidth, unit=Ytitle_unit))
0334     hs.GetXaxis().SetTitleSize(AxisTitleSize)
0335     hs.GetYaxis().SetTitleSize(AxisTitleSize)
0336     hs.GetXaxis().SetTickSize(TickSize)
0337     hs.GetYaxis().SetTickSize(TickSize)
0338     hs.GetXaxis().SetLabelSize(AxisLabelSize)
0339     hs.GetYaxis().SetLabelSize(AxisLabelSize)
0340     hs.GetYaxis().SetTitleOffset(1.3)
0341     hs.SetMaximum(totalhist.GetMaximum() * ymaxscale)
0342     hs.SetMinimum(0.1)
0343     totalhist.SetLineWidth(3)
0344     totalhist.SetLineColor(1)
0345     totalhist.Draw('histsame')
0346     c.Update()
0347     leg = TLegend((1-RightMargin)-0.45, (1-TopMargin)-0.15,
0348                   (1-RightMargin)-0.1, (1-TopMargin)-0.03)
0349     leg.SetTextSize(0.045)
0350     leg.SetFillStyle(0)
0351     leg.AddEntry('', '#it{#bf{sPHENIX}} Simulation', '')
0352     leg.AddEntry('', 'Au+Au #sqrt{s_{NN}}=200 GeV', '')
0353     leg.Draw()
0354     leg1 = TLegend(LeftMargin+0.05, (1-TopMargin)-0.22,
0355                    LeftMargin+0.3, (1-TopMargin)-0.01)
0356     # leg1.SetNColumns(3)
0357     leg1.SetTextSize(0.03)
0358     leg1.SetFillStyle(0)
0359     for i, text in enumerate(list_legtext):
0360         leg1.AddEntry(list_hist[i], text, 'f')
0361     leg1.Draw()
0362     c.Update()
0363     c.SaveAs(plotname+'.png')
0364     c.SaveAs(plotname+'.pdf')
0365     if(c):
0366         c.Close()
0367         gSystem.ProcessEvents()
0368         del c
0369         c = 0
0370
0371
0372 def GetHistogram(filename, histname):
0373     f = TFile(filename, 'r')
0374     hm = f.Get(histname)
0375     hm.SetDirectory(0)
0376     f.Close()
0377     return hm
0378
0379
0380 def str_pttop(s):
0381     return str(s).replace('.', 'p')
0382