def efficiency(stype, Zlep=True): genPoints = [800, 1000, 1200, 1400, 1600, 1800, 2000, 2500, 3000, 3500, 4000, 4500] eff = {} channels = [x for x in channelList if len(x)<5] for channel in channels: treeSign = {} ngenSign = {} nevtSign = {} eff[channel] = TGraphErrors() for i, m in enumerate(genPoints): signName = "%s_M%d" % (stype, m) #"%s_M%d" % (channel[:3], m) ngenSign[m] = 0. nevtSign[m] = 0. for j, ss in enumerate(sample[signMass]['files']): if 'nn' in channel and not 'Zinv' in ss: continue if ('en' in channel or 'mn' in channel) and not 'Wlep' in ss: continue if ('ee' in channel or 'mm' in channel) and not 'Zlep' in ss: continue if Zlep and 'Zinv' in ss: continue if not Zlep and 'Zlep' in ss: continue sfile = TFile(NTUPLEDIR + ss + ".root", "READ") if not sfile.Get("Events")==None: ngenSign[m] += sfile.Get("Events").GetEntries() # From trees treeSign[m] = sfile.Get("tree") nevtSign[m] += treeSign[m].GetEntries(selection[channel] + selection['SR']) else: ngenSign[m] = -1 print "Failed reading file", NTUPLEDIR + ss + ".root" sfile.Close() if nevtSign[m] == 0 or ngenSign[m] < 0: continue # Gen Br n = eff[channel].GetN() eff[channel].SetPoint(n, m, nevtSign[m]/ngenSign[m]) eff[channel].SetPointError(n, 0, math.sqrt(nevtSign[m])/ngenSign[m]) eff[channel].SetMarkerColor(color[channel]) eff[channel].SetMarkerStyle(20) eff[channel].SetLineColor(color[channel]) eff[channel].SetLineWidth(2) if channel.count('b')==1: eff[channel].SetLineStyle(3) n = max([eff[x].GetN() for x in channels]) maxEff = 0. # Total efficiency eff["sum"] = TGraphErrors(n) eff["sum"].SetMarkerStyle(24) eff["sum"].SetMarkerColor(1) eff["sum"].SetLineWidth(2) for i in range(n): tot, mass = 0., 0. for channel in channels: if eff[channel].GetN() > i: tot += eff[channel].GetY()[i] mass = eff[channel].GetX()[i] if tot > maxEff: maxEff = tot eff["sum"].SetPoint(i, mass, tot) leg = TLegend(0.15, 0.60, 0.95, 0.8) leg.SetBorderSize(0) leg.SetFillStyle(0) #1001 leg.SetFillColor(0) leg.SetNColumns(len(channels)/4) for i, channel in enumerate(channels): if eff[channel].GetN() > 0: leg.AddEntry(eff[channel], getChannel(channel), "pl") leg.SetY1(leg.GetY2()-len([x for x in channels if eff[x].GetN() > 0])/2.*0.045) legS = TLegend(0.55, 0.85-0.045, 0.95, 0.85) legS.SetBorderSize(0) legS.SetFillStyle(0) #1001 legS.SetFillColor(0) legS.AddEntry(eff['sum'], "Total efficiency", "pl") c1 = TCanvas("c1", "Signal Efficiency", 1200, 800) c1.cd(1) eff['sum'].Draw("APL") for i, channel in enumerate(channels): eff[channel].Draw("SAME, PL") leg.Draw() legS.Draw() setHistStyle(eff["sum"], 1.1) eff["sum"].SetTitle(";m_{"+stype[1]+"'} (GeV);Acceptance #times efficiency") eff["sum"].SetMinimum(0.) eff["sum"].SetMaximum(max(1., maxEff*1.5)) #0.65 eff["sum"].GetXaxis().SetTitleSize(0.045) eff["sum"].GetYaxis().SetTitleSize(0.045) eff["sum"].GetYaxis().SetTitleOffset(1.1) eff["sum"].GetXaxis().SetTitleOffset(1.05) eff["sum"].GetXaxis().SetRangeUser(750, 5500) if stype=='XWH' or (stype=='XZH' and Zlep): line = drawLine(750, 2./3., 4500, 2./3.) drawCMS(-1,YEAR, "Simulation") #Preliminary drawAnalysis("ZH") suffix = "" if stype=='XZH' and Zlep: suffix = "ll" elif stype=='XZH' and not Zlep: suffix = "nn" elif stype=='XWH': suffix = "ln" c1.Print("plotsSignal/Efficiency/"+stype+suffix+".pdf") c1.Print("plotsSignal/Efficiency/"+stype+suffix+".png") # print print "category", for m in range(0, eff["sum"].GetN()): print " & %d" % int(eff["sum"].GetX()[m]), print "\\\\", "\n\\hline" for i, channel in enumerate(channels+["sum"]): if channel=='sum': print "\\hline" print getChannel(channel).replace("high ", "H").replace("low ", "L").replace("purity", "P").replace("b-tag", ""), for m in range(0, eff[channel].GetN()): print "& %.1f" % (100.*eff[channel].GetY()[m]), print "\\\\"
def efficiencyAll(): #signals = {'XZHeebb':['eebb'],'XZHmmbb':['mmbb'],'XZHnnbb':['nnbb'],'XZHee0b':['ee0b'],'XZHmm0b':['mm0b'],'XZHnn0b':['nn0b'],'XZHVBFeebbVBF':['eebbVBF'],'XZHVBFmmbbVBF':['mmbbVBF'],'XZHVBFnnbbVBF':['nnbbVBF'],'XZHVBFee0bVBF':['ee0bVBF'],'XZHVBFmm0bVBF':['mm0bVBF'],'XZHVBFnn0bVBF':['nn0bVBF']} labels = {'XZHeebb' : "eeb#bar{b}",'XZHmmbb' : "#mu#mub#bar{b}",'XZHnnbb' : "#nu#nub#bar{b}",'XZHee0b' : "ee0b",'XZHmm0b' : "#mu#mu0b",'XZHnn0b' : "#nu#nu0b",'XZHVBFeebbVBF' : "eeb#bar{b}VBF",'XZHVBFmmbbVBF' : "#mu#mub#bar{b}VBF",'XZHVBFnnbbVBF' : "#nu#nub#bar{b}VBF",'XZHVBFee0bVBF' : "ee0bVBF",'XZHVBFmm0bVBF' : "#mu#mu0bVBF",'XZHVBFnn0bVBF' : "#nu#nu0bVBF"} colors = {'XZHeebb' : 2, 'XZHmmbb' : 4, 'XZHnnbb' : 2,'XZHee0b' : 3, 'XZHmm0b' : 6, 'XZHnn0b' : 4,'XZHVBFeebbVBF' : 2, 'XZHVBFmmbbVBF' : 4, 'XZHVBFnnbbVBF' : 2,'XZHVBFee0bVBF' : 3, 'XZHVBFmm0bVBF' : 6, 'XZHVBFnn0bVBF' : 4} styles = {'XZHeebb' : 1, 'XZHmmbb' : 1, 'XZHnnbb' : 1,'XZHee0b' : 1, 'XZHmm0b' : 1, 'XZHnn0b' : 1,'XZHVBFeebbVBF' : 1, 'XZHVBFmmbbVBF' : 1, 'XZHVBFnnbbVBF' : 1,'XZHVBFee0bVBF' : 1, 'XZHVBFmm0bVBF' : 1, 'XZHVBFnn0bVBF' : 1} marker = {'XZHeebb' : 22, 'XZHmmbb' : 20, 'XZHnnbb' : 22,'XZHee0b' : 22, 'XZHmm0b' : 20, 'XZHnn0b' : 20, 'XZHVBFeebbVBF' : 22, 'XZHVBFmmbbVBF' : 20, 'XZHVBFnnbbVBF' : 22,'XZHVBFee0bVBF' : 22, 'XZHVBFmm0bVBF' : 20, 'XZHVBFnn0bVBF' : 20} genPoints = [800, 1000, 1200, 1400, 1600, 1800, 2000, 2500, 3000, 3500, 4000, 4500, 5000] eff = {} for signal_samples in ['ZlepHinc','ZinvHinc','ZinvHincVBF','ZlepHincVBF']: if signal_samples == 'ZinvHinc': signals = {'XZHnnbb':['nnbb'],'XZHnn0b':['nn0b']} sign_list = ['XZHnnbb','XZHnn0b'] elif signal_samples == 'ZlepHinc': signals = {'XZHeebb':['eebb'],'XZHmmbb':['mmbb'],'XZHee0b':['ee0b'],'XZHmm0b':['mm0b']} sign_list = ['XZHeebb','XZHmmbb','XZHee0b', 'XZHmm0b'] elif signal_samples == 'ZinvHincVBF': signals = {'XZHVBFnnbbVBF':['nnbbVBF'],'XZHVBFnn0bVBF':['nn0bVBF']} sign_list = ['XZHVBFnnbbVBF', 'XZHVBFnn0bVBF'] elif signal_samples == 'ZlepHincVBF': signals = {'XZHVBFeebbVBF':['eebbVBF'],'XZHVBFmmbbVBF':['mmbbVBF'],'XZHVBFee0bVBF':['ee0bVBF'],'XZHVBFmm0bVBF':['mm0bVBF']} sign_list = ['XZHVBFeebbVBF', 'XZHVBFmmbbVBF','XZHVBFee0bVBF', 'XZHVBFmm0bVBF'] for sign, channels in signals.iteritems(): treeSign = {} ngenSign = {} nevtSign = {} eff[sign] = TGraphErrors() eff[sign].SetTitle(sign) eff[sign].SetMarkerColor(colors[sign]) eff[sign].SetMarkerSize(1.25) eff[sign].SetLineColor(colors[sign]) eff[sign].SetLineWidth(2) eff[sign].SetLineStyle(styles[sign]) eff[sign].SetMarkerStyle(marker[sign]) for i, m in enumerate(genPoints): neff = 0. for channel in channels: if signal_samples == 'ZinvHinc': file_list = ['Ntuples2016/XZH/ZprimeToZHToZinvHall_narrow_M%s'%m,'Ntuples2017/XZH/ZprimeToZHToZinvHall_narrow_M%s'%m,'Ntuples2018/XZH/ZprimeToZHToZinvHall_narrow_M%s'%m] elif signal_samples == 'ZlepHinc': file_list = ['Ntuples2016/XZH/ZprimeToZHToZlepHinc_narrow_M%s'%m,'Ntuples2017/XZH/ZprimeToZHToZlepHinc_narrow_M%s'%m,'Ntuples2018/XZH/ZprimeToZHToZlepHinc_narrow_M%s'%m] elif signal_samples == 'ZinvHincVBF': file_list = ['Ntuples2016/XZHVBF/Zprime_VBF_Zh_Zinvhinc_narrow_M-%s'%m,'Ntuples2017/XZHVBF/Zprime_VBF_Zh_Zinvhinc_narrow_M-%s'%m,'Ntuples2018/XZHVBF/Zprime_VBF_Zh_Zinvhinc_narrow_M-%s'%m] elif signal_samples == 'ZlepHincVBF': file_list = ['Ntuples2016/XZHVBF/Zprime_VBF_Zh_Zlephinc_narrow_M-%s'%m,'Ntuples2017/XZHVBF/Zprime_VBF_Zh_Zlephinc_narrow_M-%s'%m,'Ntuples2018/XZHVBF/Zprime_VBF_Zh_Zlephinc_narrow_M-%s'%m] #if not 'VBF' in channel: # signMass = "XZH_M%d" % m #else: # signMass = "XZHVBF_M%d" % m ngenSign[m] = 0. nevtSign[m] = 0. #for j, ss in enumerate(sample[signMass]['files']): for j, ss in enumerate(file_list): sfile = TFile(NTUPLEDIR + ss + ".root", "READ") if not sfile.Get("Events")==None: ngenSign[m] += sfile.Get("Events").GetEntries() # From trees treeSign[m] = sfile.Get("tree") nevtSign[m] += treeSign[m].GetEntries(selection[channel] + selection['SR']) else: ngenSign[m] = -1 print "Failed reading file", NTUPLEDIR + ss + ".root" sfile.Close() if nevtSign[m] == 0 or ngenSign[m] < 0: continue # Gen Br #print "m:",m #print "nevtSign:",nevtSign[m] #print "ngenSign:",ngenSign[m] neff += nevtSign[m]/ngenSign[m] if 'ln' in sign or 'll' in sign: neff *= 1.5 n = eff[sign].GetN() eff[sign].SetPoint(n, m, neff) eff[sign].SetPointError(n, 0, 0) n = 0. #max([eff[x].GetN() for x in channels]) maxEff = 0. #sign_list = ['XZHeebb','XZHmmbb','XZHnnbb', 'XZHee0b', 'XZHmm0b', 'XZHnn0b','XZHVBFeebbVBF', 'XZHVBFmmbbVBF', 'XZHVBFnnbbVBF','XZHVBFee0bVBF', 'XZHVBFmm0bVBF', 'XZHVBFnn0bVBF'] leg = TLegend(0.15, 0.15, 0.95, 0.35) #leg = TLegend(0.15, 0.7, 0.95, 0.8) leg.SetBorderSize(0) leg.SetFillStyle(0) #1001 leg.SetFillColor(0) for sign in sign_list: #if eff[sign].GetN() > 0: leg.AddEntry(eff[sign], labels[sign], "pl") n += 1 leg.SetNColumns(int(n/3)) leg.SetY1(leg.GetY2()-n*0.045/leg.GetNColumns()) n_error = max([eff[x].GetN() for x in sign_list]) # Total efficiency eff["sum"] = TGraphErrors(n_error) eff["sum"].SetMarkerStyle(24) eff["sum"].SetMarkerColor(1) eff["sum"].SetLineWidth(2) for i in range(n_error): tot, mass = 0., 0. for sign in sign_list: if eff[sign].GetN() > i: tot += eff[sign].GetY()[i] mass = eff[sign].GetX()[i] if tot > maxEff: maxEff = tot eff["sum"].SetPoint(i, mass, tot) #legS = TLegend(0.55, 0.85-0.045, 0.95, 0.85) legS = TLegend(0.55, 0.35-0.045, 0.95, 0.35) legS.SetBorderSize(0) legS.SetFillStyle(0) #1001 legS.SetFillColor(0) legS.AddEntry(eff['sum'], "%s Total efficiency"%signal_samples, "pl") c1 = TCanvas("c1", "Signal Efficiency", 1200, 800) c1.cd(1) c1.GetPad(0).SetTicks(1, 1) c1.SetLogy() first = sign_list[0] #if eff['XZHeebb'].GetN()!=0: # first = 'XZHeebb' #else: # first = 'XZHnnbb' eff[first].Draw("APL") for sign, channels in signals.iteritems(): eff[sign].Draw("APL" if i==0 else "SAME, PL") eff["sum"].Draw("SAME, PL") leg.Draw() legS.Draw() setHistStyle(eff[first], 1.1) eff[first].SetTitle(";m_{X} (GeV);Acceptance #times efficiency") eff[first].SetMinimum(0.) eff[first].SetMaximum(max(1., maxEff*1.5)) #0.65 eff[first].GetXaxis().SetTitleSize(0.045) eff[first].GetYaxis().SetTitleSize(0.045) eff[first].GetXaxis().SetLabelSize(0.045) eff[first].GetYaxis().SetLabelSize(0.045) eff[first].GetYaxis().SetTitleOffset(1.1) eff[first].GetXaxis().SetTitleOffset(1.05) eff[first].GetXaxis().SetRangeUser(750,5500) eff[first].GetYaxis().SetRangeUser(0., 0.4) drawCMS(-1,YEAR, "Simulation") """ latex = TLatex() latex.SetNDC() latex.SetTextSize(0.05) latex.SetTextColor(1) latex.SetTextFont(42) latex.SetTextAlign(13) latex.DrawLatex(0.83, 0.99, "(13 TeV)") latex.SetTextFont(62) latex.SetTextSize(0.06) latex.DrawLatex(0.15, 0.90, "CMS") latex.SetTextSize(0.05) latex.SetTextFont(52) """ #c1.Print("plotsSignal/Efficiency/Efficiency.pdf") #c1.Print("plotsSignal/Efficiency/Efficiency.png") c1.Print("plotsSignal/Efficiency/%s_Efficiency.pdf"%signal_samples) c1.Print("plotsSignal/Efficiency/%s_Efficiency.png"%signal_samples)
def limit2HDM(): global signals signals = range(800, 2000 + 1, 50) multF = HTOBB THEORY = ['T1', 'T2'] mass, val = fillValues("./combine/AZh/AZh_M%d.txt") Obs0s = TGraph() Exp0s = TGraph() Exp1s = TGraphAsymmErrors() Exp2s = TGraphAsymmErrors() massB, valB = fillValues("./combine/BBAZh/BBAZh_M%d.txt") Obs0sB = TGraph() Exp0sB = TGraph() Exp1sB = TGraphAsymmErrors() Exp2sB = TGraphAsymmErrors() for i, m in enumerate(mass): if not m in val: print "Key Error:", m, "not in value map" continue n = Exp0s.GetN() Obs0s.SetPoint(n, m, val[m][0] * multF) Exp0s.SetPoint(n, m, val[m][3] * multF) Exp1s.SetPoint(n, m, val[m][3] * multF) Exp1s.SetPointError(n, 0., 0., val[m][3] * multF - val[m][2] * multF, val[m][4] * multF - val[m][3] * multF) Exp2s.SetPoint(n, m, val[m][3] * multF) Exp2s.SetPointError(n, 0., 0., val[m][3] * multF - val[m][1] * multF, val[m][5] * multF - val[m][3] * multF) Obs0sB.SetPoint(n, m, valB[m][0] * multF) Exp0sB.SetPoint(n, m, valB[m][3] * multF) Exp1sB.SetPoint(n, m, valB[m][3] * multF) Exp1sB.SetPointError(n, 0., 0., valB[m][3] * multF - valB[m][2] * multF, valB[m][4] * multF - valB[m][3] * multF) Exp2sB.SetPoint(n, m, valB[m][3] * multF) Exp2sB.SetPointError(n, 0., 0., valB[m][3] * multF - valB[m][1] * multF, valB[m][5] * multF - valB[m][3] * multF) col = 629 Exp2s.SetLineWidth(2) Exp2s.SetLineStyle(1) Obs0s.SetLineWidth(3) Obs0s.SetMarkerStyle(0) Obs0s.SetLineColor(1) Exp0s.SetLineStyle(2) Exp0s.SetLineWidth(3) Exp0s.SetLineColor(1) # Exp1s.SetFillColorAlpha(col, 0.4) #kGreen+1 # Exp1s.SetLineColorAlpha(col, 0.4) # Exp2s.SetFillColorAlpha(col, 0.2) #kOrange # Exp2s.SetLineColorAlpha(col, 0.2) Exp1s.SetFillColor(417) Exp1s.SetLineColor(417) Exp2s.SetFillColor(800) Exp2s.SetLineColor(800) colB = 922 Exp2sB.SetLineWidth(2) Obs0sB.SetLineStyle(9) Obs0sB.SetLineWidth(3) Obs0sB.SetMarkerStyle(0) Obs0sB.SetLineColor(colB) Exp0sB.SetLineStyle(8) Exp0sB.SetLineWidth(3) Exp0sB.SetLineColor(colB) Exp1sB.SetFillColorAlpha(colB, 0.4) #kGreen+1 Exp1sB.SetLineColorAlpha(colB, 0.4) Exp2sB.SetFillColorAlpha(colB, 0.2) #kOrange Exp2sB.SetLineColorAlpha(colB, 0.2) Exp2s.GetXaxis().SetTitle("m_{A} (GeV)") Exp2s.GetXaxis().SetTitleSize(Exp2s.GetXaxis().GetTitleSize() * 1.25) Exp2s.GetXaxis().SetNoExponent(True) Exp2s.GetXaxis().SetMoreLogLabels(True) Exp2s.GetYaxis().SetTitle( "#sigma(A) #bf{#it{#Beta}}(A #rightarrow Zh) #bf{#it{#Beta}}(h #rightarrow bb) (fb)" ) Exp2s.GetYaxis().SetTitleOffset(1.5) Exp2s.GetYaxis().SetNoExponent(True) Exp2s.GetYaxis().SetMoreLogLabels() Theory = {} #for t in THEORY: # Theory[t] = TGraphAsymmErrors() # for m in sorted(THDM[t]['ggA'].keys()): # if m < mass[0] or m > mass[-1]: continue # Xs, Xs_Up, Xs_Down = 0., 0., 0. # Xs = THDM[t]['ggA'][m] # Xs_Up = Xs*(1.+math.sqrt((THDM['PDF']['ggA'][m][0]-1.)**2 + (THDM['QCD']['ggA'][m][0]-1.)**2)) # Xs_Down = Xs*(1.-math.sqrt((1.-THDM['PDF']['ggA'][m][1])**2 + (1.-THDM['QCD']['ggA'][m][1])**2)) # n = Theory[t].GetN() # Theory[t].SetPoint(n, m, Xs) # Theory[t].SetPointError(n, 0., 0., (Xs-Xs_Down), (Xs_Up-Xs)) # Theory[t].SetLineColor(theoryLineColor[t]) # Theory[t].SetFillColor(theoryFillColor[t]) # Theory[t].SetFillStyle(theoryFillStyle[t]) # Theory[t].SetLineWidth(2) # #Theory[t].SetLineStyle(7) c1 = TCanvas("c1", "Exclusion Limits", 800, 600) c1.cd() #SetPad(c1.GetPad(0)) c1.GetPad(0).SetTopMargin(0.06) c1.GetPad(0).SetRightMargin(0.05) c1.GetPad(0).SetLeftMargin(0.12) c1.GetPad(0).SetTicks(1, 1) c1.GetPad(0).SetLogy() Exp2s.Draw("A3") Exp1s.Draw("SAME, 3") Exp0s.Draw("SAME, L") # Exp2sB.Draw("SAME, 3") # Exp1sB.Draw("SAME, 3") Exp0sB.Draw("SAME, L") if not options.blind: Obs0s.Draw("SAME, L") Obs0sB.Draw("SAME, L") for t in THEORY: Theory[t].Draw("SAME, L3") Theory[t].Draw("SAME, L3X0Y0") #setHistStyle(Exp2s) # Exp2s.GetXaxis().SetTitleSize(0.045) # Exp2s.GetYaxis().SetTitleSize(0.04) # Exp2s.GetXaxis().SetLabelSize(0.04) # Exp2s.GetYaxis().SetLabelSize(0.04) # Exp2s.GetXaxis().SetTitleOffset(1) # Exp2s.GetYaxis().SetTitleOffset(1.25) Exp2s.GetXaxis().SetTitleSize(0.050) Exp2s.GetYaxis().SetTitleSize(0.050) Exp2s.GetXaxis().SetLabelSize(0.045) Exp2s.GetYaxis().SetLabelSize(0.045) Exp2s.GetXaxis().SetTitleOffset(0.90) Exp2s.GetYaxis().SetTitleOffset(1.25) Exp2s.GetYaxis().SetMoreLogLabels(True) Exp2s.GetYaxis().SetNoExponent(True) Exp2s.GetYaxis().SetRangeUser(0.5, 1.e3) Exp2s.GetXaxis().SetRangeUser(mass[0], mass[-1]) drawAnalysis('AZh') drawRegion('AZHsl', True) drawCMS(LUMI, "") #Preliminary #drawCMS(LUMI, "Work in Progress", suppressCMS=True) # legend leg = TLegend(0.6, 0.90, 0.99, 0.90) leg.SetBorderSize(0) leg.SetFillStyle(0) #1001 leg.SetFillColor(0) leg.SetHeader("95% CL upper limits") leg.AddEntry(None, "gg #rightarrow A #rightarrow Zh", "") #"95% CL upper limits" leg.AddEntry(Obs0s, "Observed", "l") leg.AddEntry(Exp0s, "Expected", "l") leg.AddEntry(Exp1s, "#pm 1 std. deviation", "f") leg.AddEntry(Exp2s, "#pm 2 std. deviation", "f") leg.AddEntry(None, "", "") leg.AddEntry(None, "bbA #rightarrow Zh", "") leg.AddEntry(Obs0sB, "Observed", "l") leg.AddEntry(Exp0sB, "Expected", "l") leg.SetY1(leg.GetY2() - leg.GetNRows() * 0.045) leg.Draw() # latex = TLatex() # latex.SetNDC() # latex.SetTextSize(0.040) # latex.SetTextFont(42) # latex.DrawLatex(0.65, leg.GetY1()-0.045, "cos(#beta-#alpha)=0.25, tan(#beta)=1") # legB = TLegend(0.12, 0.4-4*0.3/5., 0.65, 0.4) legB = TLegend(0.15, 0.27, 0.68, 0.27) legB.SetBorderSize(0) legB.SetFillStyle(0) #1001 legB.SetFillColor(0) for t in THEORY: legB.AddEntry(Theory[t], theoryLabel[t], "fl") legB.AddEntry(None, "cos(#beta-#alpha)=0.25, tan(#beta)=1", "") legB.SetY1(legB.GetY2() - legB.GetNRows() * 0.045) legB.Draw() c1.GetPad(0).RedrawAxis() leg.Draw() c1.Update() if not gROOT.IsBatch(): raw_input("Press Enter to continue...") c1.Print("plotsLimit/Exclusion/THDM.png") c1.Print("plotsLimit/Exclusion/THDM.pdf")
def make_plot_all_rods(error_dict, rod_dict, name): leg = TLegend(0,0,0,0) if 'lock' in name: leg = TLegend(0.18,0.85,0.50,0.55) else: leg = TLegend(0.18,0.85,0.40,0.55) leg.SetLineColor(0) leg.SetFillStyle(0) leg.SetShadowColor(0) leg.SetBorderSize(0) leg.SetNColumns(3) gStyle.SetLegendTextSize(0.045) v_hists = [] #for e,c in zip(error_bits, error_colors): for e in error_bits: h = TH1F('h'+e,'h'+e, len(rod_dict), 0, len(rod_dict)) h.SetFillStyle(1001) h.SetLineWidth(0) v_hists.append(h) v_hists[-1].SetDirectory(0) if bin(int('0x'+e, 16))[2:].zfill(4) == '0111': leg.AddEntry(v_hists[-1],'GOL 3',"f"); elif bin(int('0x'+e, 16))[2:].zfill(4) == '1011': leg.AddEntry(v_hists[-1],'GOL 2',"f"); elif bin(int('0x'+e, 16))[2:].zfill(4) == '1101': leg.AddEntry(v_hists[-1],'GOL 1',"f"); elif bin(int('0x'+e, 16))[2:].zfill(4) == '1110': leg.AddEntry(v_hists[-1],'GOL 0',"f"); else: leg.AddEntry(v_hists[-1],bin(int('0x'+e, 16))[2:].zfill(4),"f"); h = leg.GetY2()-leg.GetY1(); w = leg.GetX2()-leg.GetX1()*.6; leg.SetMargin(leg.GetNColumns()*h/(leg.GetNRows()*w)) for key,val in error_dict.items(): idx_rod = 0 for i, key2 in enumerate(rod_dict): if key2 == key[:8]: idx_rod = i v_hists[int(key[11:12], 16)].Fill(idx_rod, val) stack = THStack("stack","stack") for hist in v_hists: stack.Add(hist) if 'buff' in name: c1 = TCanvas( 'c1', 'c1', 2000, 500) else: c1 = TCanvas( 'c1', 'c1', 1000, 500) h1 = TH1F('h_1','h_1', len(rod_dict), 0, len(rod_dict)) for i, key in enumerate(rod_dict): h1.GetXaxis().SetBinLabel(i+1,key) h1.SetBinContent(i+1,rod_dict[key]) h1.GetXaxis().LabelsOption("v") h1.GetXaxis().SetTitle("ROD") h1.GetXaxis().SetTitleOffset(2.2) h1.GetYaxis().SetTitle("# of rocketio errors") h1.SetLineColor(kRed) h1.SetLineWidth(1) leg.AddEntry(h1,'total',"l"); c1.SetBottomMargin(0.23) h1.GetXaxis().SetTitle("ROD") h1.Draw("HIST") stack.Draw("PFC PLC SAME HIST") h1.Draw("SAMEHIST") AtlasStyle.ATLAS_LABEL(0.19,.88, 1, "Internal") leg.Draw() c1.Update() c1.Print("plots/"+name +".pdf") c1.Clear()
def MakeLegend(can, x1=.8, y1=.8, x2=.9, y2=.9, textsize=18, ncolumns=1, totalentries=0, option='f', skip=[]): from ROOT import TLegend, TH1, gStyle, TGraph if can.GetPrimitive('pad_top'): MakeLegend(can.GetPrimitive('pad_top'), x1, y1, x2, y2, textsize, ncolumns, totalentries, skip=skip) return if CanvasEmpty(can): print 'Error: trying to make legend from canvas with 0 plots. Will do nothing.' return # # if a previous version exists from this function, delete it # if can.GetPrimitive('legend'): can.GetPrimitive('legend').Delete() leg = TLegend(x1, y1, x2, y2) leg.SetName('legend') tobject_collector.append(leg) leg.SetTextFont(43) leg.SetTextSize(textsize) leg.SetTextFont(43) leg.SetBorderSize(0) leg.SetFillStyle(0) leg.SetNColumns(ncolumns) # # Add by TH1 GetTitle() # the_primitives = can.GetListOfPrimitives() if can.GetPrimitive('pad_top'): the_primitives = can.GetPrimitive('pad_top').GetListOfPrimitives() if can.GetPrimitive('stack'): the_stack = list(reversed(list(can.GetPrimitive('stack').GetHists()))) the_primitives = the_stack + list(the_primitives) if type(option) == type(''): option = [option] * 100 total = 0 for i in the_primitives: if i.GetName() == 'stack': continue drawopt = i.GetDrawOption() if issubclass(type(i), TH1) or issubclass(type(i), TGraph): if i.GetTitle() in skip: continue leg.AddEntry(i, i.GetTitle(), option[total]) # plef total += 1 # # Add empty entries to ensure a standard layout # for i in range(100): if totalentries == 0: break if total >= totalentries: break leg.AddEntry(None, '', '') total += 1 # recipe for making roughly square boxes h = leg.GetY2() - leg.GetY1() w = leg.GetX2() - leg.GetX1() leg.SetMargin(leg.GetNColumns() * h / float(leg.GetNRows() * w)) can.cd() if can.GetPrimitive('pad_top'): can.GetPrimitive('pad_top').cd() leg.Draw() can.Modified() can.Update() return
def efficiency(year): import numpy as np from root_numpy import tree2array, fill_hist from aliases import AK8veto, electronVeto, muonVeto genPoints = [ 1800, 2000, 2500, 3000, 3500, 4000, 4500, 5000, 5500, 6000, 7000, 8000 ] eff = {} vetoes = {"AK8": AK8veto, "electron": electronVeto, "muon": muonVeto} VETO = "AK8" ##could change the veto to investigate here if SEPARATE: eff_add = {} #channels = ['none', 'qq', 'bq', 'bb', 'mumu'] channels = ['qq', 'bq', 'bb', 'mumu'] for channel in channels: treeSign = {} ngenSign = {} nevtSign = {} eff[channel] = TGraphErrors() if SEPARATE: nevtSign_add = {} eff_add[channel] = TGraphErrors() for i, m in enumerate(genPoints): signName = "ZpBB_M" + str(m) ngenSign[m] = 0. nevtSign[m] = 0. if SEPARATE: nevtSign_add[m] = 0. for j, ss in enumerate(sample[signName]['files']): if year == "run2" or year in ss: sfile = TFile(NTUPLEDIR + ss + ".root", "READ") ngenSign[m] += sfile.Get("Events").GetBinContent(1) treeSign[m] = sfile.Get("tree") if BTAGGING == 'semimedium': #if SEPARATE: # temp_array = tree2array(treeSign[m], branches='BTagAK4Weight_deepJet', selection=aliasSM[channel].replace(vetoes[VETO], "")) #else: temp_array = tree2array( treeSign[m], branches='BTagAK4Weight_deepJet', selection=aliasSM[channel]) temp_hist = TH1F('pass', 'pass', 1, 0, 1) fill_hist(temp_hist, np.zeros(len(temp_array)), weights=temp_array) nevtSign[m] += temp_hist.GetBinContent(1) temp_array = None temp_hist.Reset() if SEPARATE: temp_array = tree2array( treeSign[m], branches='BTagAK4Weight_deepJet', selection=aliasSM[channel].replace( vetoes[VETO], "")) temp_hist = TH1F('pass', 'pass', 1, 0, 1) fill_hist(temp_hist, np.zeros(len(temp_array)), weights=temp_array) nevtSign[m] += temp_hist.GetBinContent(1) temp_array = None temp_hist.Reset() else: #if SEPARATE: # temp_array = tree2array(treeSign[m], branches='BTagAK4Weight_deepJet', selection=alias[channel].format(WP=working_points[BTAGGING]).replace(vetoes[VETO], "")) #else: temp_array = tree2array( treeSign[m], branches='BTagAK4Weight_deepJet', selection=alias[channel].format( WP=working_points[BTAGGING])) temp_hist = TH1F('pass', 'pass', 1, 0, 1) fill_hist(temp_hist, np.zeros(len(temp_array)), weights=temp_array) nevtSign[m] += temp_hist.GetBinContent(1) temp_array = None temp_hist.Reset() if SEPARATE: temp_array = tree2array( treeSign[m], branches='BTagAK4Weight_deepJet', selection=alias[channel].format( WP=working_points[BTAGGING]).replace( vetoes[VETO], "")) temp_hist = TH1F('pass', 'pass', 1, 0, 1) fill_hist(temp_hist, np.zeros(len(temp_array)), weights=temp_array) nevtSign_add[m] += temp_hist.GetBinContent(1) temp_array = None temp_hist.Reset() sfile.Close() print channel, ss, ":", nevtSign[m], "/", ngenSign[ m], "=", nevtSign[m] / ngenSign[m] if nevtSign[m] == 0 or ngenSign[m] < 0: continue n = eff[channel].GetN() eff[channel].SetPoint(n, m, nevtSign[m] / ngenSign[m]) eff[channel].SetPointError(n, 0, math.sqrt(nevtSign[m]) / ngenSign[m]) if SEPARATE: eff_add[channel].SetPoint(n, m, nevtSign_add[m] / ngenSign[m]) eff_add[channel].SetPointError( n, 0, math.sqrt(nevtSign_add[m]) / ngenSign[m]) eff[channel].SetMarkerColor(color[channel]) eff[channel].SetMarkerStyle(20) eff[channel].SetLineColor(color[channel]) eff[channel].SetLineWidth(2) if SEPARATE: eff_add[channel].SetMarkerColor(color[channel] + color_shift[channel]) eff_add[channel].SetMarkerStyle(21) eff_add[channel].SetLineColor(color[channel] + color_shift[channel]) eff_add[channel].SetLineWidth(2) eff_add[channel].SetLineStyle(7) if channel == 'qq' or channel == 'none': eff[channel].SetLineStyle(3) n = max([eff[x].GetN() for x in channels]) maxEff = 0. # Total efficiency eff["sum"] = TGraphErrors(n) eff["sum"].SetMarkerStyle(24) eff["sum"].SetMarkerColor(1) eff["sum"].SetLineWidth(2) if SEPARATE: eff_add["sum"] = TGraphErrors(n) eff_add["sum"].SetMarkerStyle(25) eff_add["sum"].SetMarkerColor(1) eff_add["sum"].SetLineWidth(2) eff_add["sum"].SetLineStyle(7) for i in range(n): tot, mass = 0., 0. if SEPARATE: tot_add = 0. for channel in channels: if channel == 'qq' or channel == 'none': continue #not sure if I should include 2mu category in sum if eff[channel].GetN() > i: tot += eff[channel].GetY()[i] if SEPARATE: tot_add += eff_add[channel].GetY()[i] mass = eff[channel].GetX()[i] if tot > maxEff: maxEff = tot eff["sum"].SetPoint(i, mass, tot) if SEPARATE: eff_add["sum"].SetPoint(i, mass, tot_add) if SEPARATE: leg = TLegend(0.15, 0.50, 0.95, 0.8) else: leg = TLegend(0.15, 0.60, 0.95, 0.8) leg.SetBorderSize(0) leg.SetFillStyle(0) #1001 leg.SetFillColor(0) leg.SetNColumns(len(channels) / 4) for i, channel in enumerate(channels): if eff[channel].GetN() > 0: leg.AddEntry(eff[channel], getChannel(channel), "pl") if SEPARATE: leg.AddEntry(eff_add[channel], getChannel(channel) + " no " + VETO + "-veto", "pl") if SEPARATE: leg.SetY1(leg.GetY2() - len([x for x in channels if eff[x].GetN() > 0]) * 0.045) else: leg.SetY1(leg.GetY2() - len([x for x in channels if eff[x].GetN() > 0]) / 2. * 0.045) if SEPARATE: legS = TLegend(0.5, 0.8 - 0.045, 0.9, 0.85) else: legS = TLegend(0.5, 0.85 - 0.045, 0.9, 0.85) legS.SetBorderSize(0) legS.SetFillStyle(0) #1001 legS.SetFillColor(0) legS.AddEntry(eff['sum'], "Total b tag efficiency (1 b tag + 2 b tag + 2 #mu)", "pl") if SEPARATE: legS.AddEntry(eff_add['sum'], "Total b tag efficiency, no " + VETO + "-veto", "pl") c1 = TCanvas("c1", "Signal Efficiency", 1200, 800) c1.cd(1) eff['sum'].Draw("APL") if SEPARATE: eff_add['sum'].Draw("SAME, PL") for i, channel in enumerate(channels): eff[channel].Draw("SAME, PL") if SEPARATE: eff_add[channel].Draw("SAME, PL") leg.Draw() legS.Draw() setHistStyle(eff["sum"], 1.1) eff["sum"].SetTitle(";m_{Z'} (GeV);Acceptance #times efficiency") eff["sum"].SetMinimum(0.) eff["sum"].SetMaximum(max(1., maxEff * 1.5)) #0.65 if SEPARATE: eff_add["sum"].SetTitle(";m_{Z'} (GeV);Acceptance #times efficiency") eff_add["sum"].SetMinimum(0.) eff_add["sum"].SetMaximum(1.) eff["sum"].GetXaxis().SetTitleSize(0.045) eff["sum"].GetYaxis().SetTitleSize(0.045) eff["sum"].GetYaxis().SetTitleOffset(1.1) eff["sum"].GetXaxis().SetTitleOffset(1.05) eff["sum"].GetXaxis().SetRangeUser(1500, 8000) c1.SetTopMargin(0.05) #drawCMS(-1, "Simulation Preliminary", year=year) #Preliminary #drawCMS(-1, "Work in Progress", year=year, suppressCMS=True) drawCMS(-1, "", year=year, suppressCMS=True) drawAnalysis("") if SEPARATE: c1.Print("plots/Efficiency/" + year + "_" + BTAGGING + "_no" + VETO + "veto.pdf") c1.Print("plots/Efficiency/" + year + "_" + BTAGGING + "_no" + VETO + "veto.png") else: c1.Print("plots/Efficiency/" + year + "_" + BTAGGING + ".pdf") c1.Print("plots/Efficiency/" + year + "_" + BTAGGING + ".png") # print print "category", for m in range(0, eff["sum"].GetN()): print " & %d" % int(eff["sum"].GetX()[m]), print "\\\\", "\n\\hline" for i, channel in enumerate(channels + ["sum"]): if channel == 'sum': print "\\hline" print getChannel(channel).replace("high ", "H").replace( "low ", "L").replace("purity", "P").replace("b-tag", ""), for m in range(0, eff[channel].GetN()): print "& %.1f" % (100. * eff[channel].GetY()[m]), print "\\\\"
def hvt(benchmark=['B3', 'A1']): hxs = {} hw = {} gxs = {} gw = {} mg = TMultiGraph() for m in massPoints: hxs[m] = TH2F("hxs_M%d" % m, ";;", 50, -0.04, 3.96, 100, 0., 2.) hw[m] = TH2F("hw_M%d" % m, ";;", 50, -0.04, 3.96, 50, 0., 2.) for m in massPoints: file = TFile.Open("HVT/scanHVT_M%s.root" % m, "READ") tree = file.Get("tree") for entry in range( tree.GetEntries()): # Fill mass points only if NOT excluded tree.GetEntry(entry) gH, gF = tree.gv * tree.ch, tree.g * tree.g * tree.cq / tree.gv XsBr = tree.CX0 * tree.BRbb * 1000. # in fb if XsBr < observed[m]: hxs[m].Fill(gH, gF) if tree.total_widthV0 / float(m) < width: hw[m].Fill(gH, gF) for b in range(hxs[m].GetNbinsX() * hxs[m].GetNbinsY()): hxs[m].SetBinContent(b, 1. if hxs[m].GetBinContent(b) > 0. else 0.) hw[m].SetBinContent(b, 1. if hw[m].GetBinContent(b) > 0. else 0.) #hxs[m].Smooth(20) #hw[m].Smooth(20) gxs[m] = getCurve(hxs[m]) for i, g in enumerate(gxs[m]): g.SetLineColor(massColors[m]) g.SetFillColor(massColors[m]) g.SetFillStyle(massFill[m]) #(3345 if i>1 else 3354) g.SetLineWidth(503 * (1 if i < 2 else -1)) mg.Add(g) #if m==3000: if m == massPoints[-1]: gw[m] = getCurve(hw[m]) for i, g in enumerate(gw[m]): g.SetPoint(0, 0., g.GetY()[0]) g.SetLineWidth(501 * (1 if i < 2 else -1)) g.SetLineColor(920 + 2) g.SetFillColor(920 + 1) g.SetFillStyle(3003) mg.Add(g) if options.root: outFile = TFile("plotsLimit/Model.root", "RECREATE") outFile.cd() for m in massPoints: mg[m].Write("X_M%d" % m) mgW.Write("width") outFile.Close() print "Saved histogram in file plotsLimit/Model.root, exiting..." exit() ### plot ### c1 = TCanvas("c1", "HVT Exclusion Limits", 800, 600) c1.cd() c1.GetPad(0).SetTopMargin(0.06) c1.GetPad(0).SetRightMargin(0.05) c1.GetPad(0).SetTicks(1, 1) mg.Draw("AC") #mg.GetXaxis().SetTitle("g_{V} c_{H}") mg.GetXaxis().SetTitle("Higgs and vector boson coupling g_{H}") mg.GetXaxis().SetRangeUser(-3., 3.) mg.GetXaxis().SetLabelSize(0.045) mg.GetXaxis().SetTitleSize(0.045) mg.GetXaxis().SetTitleOffset(1.) #mg.GetYaxis().SetTitle("g^{2} c_{F} / g_{V}") mg.GetYaxis().SetTitle("Fermion coupling g_{F}") mg.GetYaxis().SetLabelSize(0.045) mg.GetYaxis().SetTitleSize(0.045) mg.GetYaxis().SetTitleOffset(1.) mg.GetYaxis().SetRangeUser(-1.2, 1.2) mg.GetYaxis().SetNdivisions(505) # hxs[3500].Draw("CONTZ") drawCMS(LUMI, "Preliminary", False) # drawAnalysis("XVH"+category, False) # latex = TLatex() # latex.SetNDC() # latex.SetTextFont(62) # latex.SetTextSize(0.06) # latex.DrawLatex(0.10, 0.925, "CMS") # model B g_model = {} for i, b in enumerate(benchmark): g_model[i] = TGraph(1) g_model[i].SetTitle(models_name[b]) g_model[i].SetPoint(0, models_point[b][0], models_point[b][1]) g_model[i].SetMarkerStyle(models_style[b]) g_model[i].SetMarkerColor(models_color[b]) g_model[i].SetMarkerSize(1.5) g_model[i].Draw("PSAME") # text latex = TLatex() latex.SetTextSize(0.045) latex.SetTextFont(42) latex.SetTextColor(630) # for b in benchmark: latex.DrawLatex(models_point[b][0]+0.02, models_point[b][1]+0.02, models_name[b]) latex.SetTextColor(920 + 2) latex.DrawLatex(-2.8, -0.875, "#frac{#Gamma_{Z'}}{m_{Z'}} > %.0f%%" % (width * 100, )) #leg = TLegend(0.68, 0.60, 0.95, 0.94) leg = TLegend(0.68, 0.34, 0.95, 0.66) leg.SetBorderSize(1) leg.SetFillStyle(1001) leg.SetFillColor(0) for m in massPoints: leg.AddEntry(gxs[m][0], "m_{Z'} = %.1f TeV" % (m / 1000.), "fl") for i, b in enumerate(benchmark): leg.AddEntry(g_model[i], g_model[i].GetTitle(), "P") leg.SetY1(leg.GetY2() - leg.GetNRows() * 0.050) leg.SetMargin(0.35) leg.Draw() gxs_ = gxs[massPoints[0]][0].Clone("gxs_") gxs_.SetLineColor(1) # gxs_.SetFillColor(1) latex.SetNDC() latex.SetTextColor(1) latex.SetTextSize(0.04) latex.SetTextFont(52) latex.DrawLatex(0.15, 0.95, "q#bar{q} #rightarrow Z' #rightarrow b#bar{b}") c1.Print("plots/model/HVT.png") c1.Print("plots/model/HVT.pdf") c1.Print("plots/model/HVT.root") c1.Print("plots/model/HVT.C") #g = 0.646879, cH = 0.976246, cF = 1.02433 print "model B = [", 3 * 0.976246, ",", 0.646879 * 0.646879 * 1.02433 / 3, "]" if not gROOT.IsBatch(): raw_input("Press Enter to continue...")
class ROOTPlot(PlotSpec): # pylint: disable=too-many-instance-attributes """ROOT specific implementation of PlotSpec """ def __init__(self, **kwargs): super().__init__() # copy member variables from orig PlotSpec self.copy(kwargs.pop("orig", None)) # It's a ROOT object and we control its name self.name = None # frame constructed to provide axes self.frame = None # ROOT objects self.objects = None # Styles to be used for the styles self.styles = None # text to be added self.pave_boxes = None # legend self.root_legend = None # labels of objects added to this ROOTPlot self.labels = None # lines added self._root_lines = [] # TPad used for the plot self.pad = None # TODO: use PlotSpec's relative coordinates self.size = (300, 300) def __draw_objects(self): """Draw all objects """ for obj, sty in zip(self.objects, self.styles): # adjust draw_option to always contain "same" in addition to # user specified options draw_option = "same" if not sty else f"same {sty.draw_options}" obj.Draw(draw_option) def __draw_lines(self): """Draw all lines """ for l in self._lines: x_low = l.x_low x_up = l.x_up if l.x_orientation == "relative": x_range = self._axes[0].limits[1] - self._axes[0].limits[0] x_low = self._axes[0].limits[0] + l.x_low * x_range x_up = self._axes[0].limits[0] + l.x_up * x_range y_low = l.y_low y_up = l.y_up if l.y_orientation == "relative": y_range = self._axes[1].limits[1] - self._axes[1].limits[0] y_low = self._axes[1].limits[0] + l.y_low * y_range y_up = self._axes[1].limits[0] + l.y_up * y_range line = TLine(x_low, y_low, x_up, y_up) apply_line_style(line, l.style) line.Draw() self._root_lines.append(line) def __draw_legends(self): """Encapsulate legend drawing """ if self.root_legend: self.root_legend.Draw() def __style_objects(self, **kwargs): """Style all objects to be plotted Args: kwargs: dict "keep_stats" to keep potential stats boxes associated with e.g. histograms """ # TODO Make that static again keep_stats = kwargs.pop("keep_stats", False) #scale = sqrt(self.size[0] * self.size[1] / SCALE_BASE) for obj, style in zip(self.objects, self.styles): style_object(obj, style) # , scale) if not keep_stats: try_method(obj, "SetStats", 0) def add_object(self, root_object, style=None, label=None): """Add a ROOT object to be plotted inside this ROOTPlot Args: object: ROOT object to be added style: either Style or None, default None label: str to appear in the legend or None, default None """ # Prepare lists if not self.objects: self.objects = [] if not self.styles: self.styles = [] if not self.labels: self.labels = [] # clone the ROOT object and don't touch the original one self.objects.append(clone_root(root_object)) self.styles.append(style) self.labels.append(label) def __create_frame(self, **kwargs): """Make the frame used to plot the axes Args: kwargs: dict x_axis_title: str for x-axis title y_axis_title: str for x-axis title use_any_titles: bool whether or not to use any titles defined for any of the ROOT objects in case title for x- or y-axis are not specified by the user """ self._axes[0].title = kwargs.pop("x_axis_title", self._axes[0].title) self._axes[1].title = kwargs.pop("y_axis_title", self._axes[1].title) use_any_titles = kwargs.pop("use_any_titles", True) # pylint: disable=protected-access x_force_limits = False if self._share_x: self._axes[0].limits[0] = self._share_x._axes[0].limits[0] self._axes[0].limits[1] = self._share_x._axes[0].limits[1] x_force_limits = True y_force_limits = False if self._share_y: self._axes[1].limits[0] = self._share_y._axes[1].limits[0] self._axes[1].limits[1] = self._share_y._axes[1].limits[1] y_force_limits = True # pylint: enable=protected-access # Find the x- and y-limits for this plot x_low, x_up, y_low, y_up, z_low, z_up = \ find_boundaries(self.objects, self._axes[0].limits[0], self._axes[0].limits[1], self._axes[1].limits[0], self._axes[1].limits[1], self._axes[2].limits[0], self._axes[2].limits[1], reserve_ndc_top=kwargs.pop("reserve_ndc_top", None), reserve_ndc_bottom=kwargs.pop("reserve_ndc_bottom", None), x_force_limits=x_force_limits, y_force_limits=y_force_limits, x_log=self._axes[0].is_log, y_log=self._axes[1].is_log, y_account_for_errors=self._axes[1].account_for_errors) # add titles to axes if not specified by the user by trying to # use those which are set for any ROOT object if use_any_titles and \ (not self._axes[0].title or not self._axes[1].title): for obj in self.objects: if not self._axes[0].title and hasattr(obj, "GetXaxis"): self._axes[0].title = obj.GetXaxis().GetTitle() if not self._axes[1].title and hasattr(obj, "GetYaxis"): self._axes[1].title = obj.GetYaxis().GetTitle() if self._axes[0].title and self._axes[1].title: break # Finally create the frame for this plot frame_title = self._title if self._title else "" frame_string = f"{frame_title};{self._axes[0].title};{self._axes[1].title}" self.frame = self.pad.DrawFrame(x_low, y_low, x_up, y_up, frame_string) self._axes[0].limits[0] = x_low self._axes[0].limits[1] = x_up self._axes[1].limits[0] = y_low self._axes[1].limits[1] = y_up self._axes[2].limits[0] = z_low self._axes[2].limits[1] = z_up # Give the frame a unique name, for now just because we do it for # ROOT related object self.frame.SetName(f"{self.name}_frame") def __adjust_text_size(self, size): return int(max(1, size * self._parent_figure_spec.size[1])) def __adjust_text_size_x(self, size): """Helper method to adjust text sizes for x-axis' title and labels TODO That has to be revised """ return size / (self._rel_coordinates[2] - self._rel_coordinates[0]) def __adjust_text_size_y(self, size): """Helper method to adjust text sizes for y-axis' title and labels TODO That has to be revised """ return size / (self._rel_coordinates[3] - self._rel_coordinates[1]) def __adjust_row_margin(self, margin): """Helper method to recompute the relative row margin in the TPad based on the relative values of the user defined for the FigureSpec """ return margin / (self._rel_coordinates[3] - self._rel_coordinates[1]) def __adjust_column_margin(self, margin): """Helper method to recompute the relative column margin in the TPad based on the relative values of the user defined for the FigureSpec """ return margin / (self._rel_coordinates[2] - self._rel_coordinates[0]) def __adjust_tick_size_x(self, size): """Helper method to adjust the x-tick lengths accordingly """ return size / ((self.pad.GetUxmax() - self.pad.GetUxmin()) / \ (self.pad.GetX2()-self.pad.GetX1()) * \ (self._rel_coordinates[3] - self._rel_coordinates[1])) def __adjust_tick_size_y(self, size): """Helper method to adjust the y-tick lengths accordingly """ return size / ((self.pad.GetUymax() - self.pad.GetUymin()) / \ (self.pad.GetY2()-self.pad.GetY1()) * \ (self._rel_coordinates[2] - self._rel_coordinates[0])) def __adjust_legend_coordinates(self, coords): """Helper method to recompute the legend coordinates properly given the margins TODO This seems to work but needs to be revised to be sure """ x_low = map_value(coords[0], 0, 1, self.__adjust_column_margin(self._column_margins[0]), 1 - self.__adjust_column_margin(self._column_margins[1])) x_up = map_value(coords[2], 0, 1, self.__adjust_column_margin(self._column_margins[0]), 1 - self.__adjust_column_margin(self._column_margins[1])) y_low = map_value(coords[1], 0, 1, self.__adjust_row_margin(self._row_margins[0]), 1 - self.__adjust_row_margin(self._row_margins[1])) y_up = map_value(coords[3], 0, 1, self.__adjust_row_margin(self._row_margins[0]), 1 - self.__adjust_row_margin(self._row_margins[1])) return [x_low, y_low, x_up, y_up] def __adjust_coordinate_x(self, x_in): return map_value(x_in, 0, 1, self.__adjust_column_margin(self._column_margins[0]), 1 - self.__adjust_column_margin(self._column_margins[1])) def __adjust_coordinate_y(self, y_in): return map_value(y_in, 0, 1, self.__adjust_row_margin(self._row_margins[0]), 1 - self.__adjust_row_margin(self._row_margins[1])) def __create_legends(self): # pylint: disable=too-many-branches """Create legend(s) So far it only makes one legend but it should be possible to create multiple in general """ # count how many labels we want to have n_labels = sum(l is not None for l in self.labels) # just return if there are no labels if not n_labels: return # Adjust to account for number of columns n_labels = int(n_labels / self._legend_spec.n_columns) # Get the same legend positioning relative to the axes in the plot # adjust legend position starting from top right as default coordinates = [0.5, 0.7, 1, 0.89] if not isinstance(self._legend_spec.position, str): # assume relative coordinates coordinates = self._legend_spec.position # TODO that is rough if coordinates[3] < 0.5 and not self._legend_spec.principal_position: self._legend_spec.principal_position = "bottom" elif not self._legend_spec.principal_position: self._legend_spec.principal_position = "top" else: # deal with text string position # however, catch this one... if "left" in self._legend_spec.position and "right" in self._legend_spec.position: raise ValueError("Choose EITHER \"left\" OR \"right\" for legend positioning") if "bottom" in self._legend_spec.position and "top" in self._legend_spec.position: raise ValueError("Choose EITHER \"bottom\" OR \"top\" for legend positioning") # and now really adjust the coordinates if "left" in self._legend_spec.position: coordinates[0], coordinates[2] = (0.1, 0.5) if "bottom" in self._legend_spec.position: coordinates[1] = 0.1 if not self._legend_spec.principal_position: self._legend_spec.principal_position = "bottom" elif not self._legend_spec.principal_position: self._legend_spec.principal_position = "top" coordinates = self.__adjust_legend_coordinates(coordinates) if self._legend_spec.principal_position == "bottom": coordinates[3] = coordinates[1] + 0.05 * n_labels else: coordinates[1] = coordinates[3] - 0.05 * n_labels self.root_legend = TLegend(*coordinates) self.root_legend.SetNColumns(self._legend_spec.n_columns) # Make legend transparent and remove border self.root_legend.SetFillStyle(0) self.root_legend.SetLineWidth(0) # TODO this has to synced correctly with the line height of the legend which # atm is 0.05 (see above) self.root_legend.SetTextFont(63) self.root_legend.SetTextSize(self.__adjust_text_size(self._legend_spec.text_size)) for obj, lab in zip(self.objects, self.labels): if lab is None: continue self.root_legend.AddEntry(obj, lab) def __style_frame(self): """Style the frame axes """ # recomupte tick lengths self.frame.GetXaxis().SetTickLength(self.__adjust_tick_size_x(self._axes[0].tick_size)) self.frame.GetYaxis().SetTickLength(self.__adjust_tick_size_y(self._axes[1].tick_size)) # limit number of digits, TODO maybe make it configurable in the future self.frame.GetYaxis().SetMaxDigits(4) if self._share_x: # no labels or title in case of shared x-axis self.frame.GetXaxis().SetTitleSize(0) self.frame.GetXaxis().SetLabelSize(0) else: axis = self.frame.GetXaxis() axis.SetLabelFont(63) axis.SetTitleFont(63) axis.SetTitleSize(self.__adjust_text_size(self._axes[0].title_size)) axis.SetLabelSize(self.__adjust_text_size(self._axes[0].label_size)) # TODO properly compute offsets, the following is a wild guess for now #margin = self.__adjust_row_margin(self._row_margins[0]) if self._axes[0].title_offset is not None: axis.SetTitleOffset(self.__adjust_column_margin(self._axes[0].title_offset)) #axis.SetTitleOffset(19 * margin) if self._share_y: # no labels or title in case of shared y-axis self.frame.GetYaxis().SetTitleSize(0) self.frame.GetYaxis().SetLabelSize(0) else: axis = self.frame.GetYaxis() axis.SetLabelFont(63) axis.SetTitleFont(63) axis.SetTitleSize(self.__adjust_text_size(self._axes[1].title_size)) axis.SetLabelSize(self.__adjust_text_size(self._axes[1].label_size)) # TODO properly compute offsets, the following is a wild guess for now #margin = self.__adjust_column_margin(self._column_margins[0]) if self._axes[1].title_offset is not None: axis.SetTitleOffset(self.__adjust_column_margin(self._axes[1].title_offset)) #axis.SetTitleOffset(23 * margin) def __draw_text(self): """Draw potential text into pad """ if not self._texts: return if not self.pave_boxes: self.pave_boxes = [] for text in self._texts: pave_box = TPaveText(self.__adjust_coordinate_x(text.x_low), self.__adjust_coordinate_y(text.y_low), self.__adjust_coordinate_x(1), self.__adjust_coordinate_y(text.y_low + text.size), "brNDC") pave_box.SetLineWidth(0) pave_box.AddText(text.text) pave_box.SetBorderSize(0) pave_box.SetFillStyle(0) pave_box.SetTextAlign(10) pave_box.SetTextFont(63) pave_box.SetTextSizePixels(self.__adjust_text_size(text.size)) self.pave_boxes.append(pave_box) pave_box.Draw() def create(self, name, **kwargs): """Create this plot Args: name: str for TPad name kwargs: dict reserve_ndc_top: float to reserve relative space for the legend at the top """ if not self.objects: return # remember if another TPad was active before prev_pad = gPad.cd() if gPad and gPad.GetName() != name else None self.name = name self.pad = TPad(name, "", *self._rel_coordinates) self.pad.Draw() if self._axes[1].is_log: self.pad.SetLogy() if self._axes[0].is_log: self.pad.SetLogx() # but for now change to this TPad self.pad.cd() # This HAS to come before the frame creation self.pad.SetLeftMargin(self.__adjust_column_margin(self._column_margins[0])) self.pad.SetRightMargin(self.__adjust_column_margin(self._column_margins[1])) self.pad.SetBottomMargin(self.__adjust_row_margin(self._row_margins[0])) self.pad.SetTopMargin(self.__adjust_row_margin(self._row_margins[1])) # Set ticks on either side, might be customisable in the future self.pad.SetTickx(1) self.pad.SetTicky(1) # style objects and create legend self.__style_objects(**kwargs) self.__create_legends() if self.root_legend and self._legend_spec.principal_position == "top": kwargs["reserve_ndc_top"] = \ 1 - map_value(self.root_legend.GetY1(), self.__adjust_row_margin(self._row_margins[0]), 1 - self.__adjust_row_margin(self._row_margins[1]), 0, 1) elif self.root_legend and self._legend_spec.principal_position == "bottom": kwargs["reserve_ndc_bottom"] = \ map_value(self.root_legend.GetY2(), self.__adjust_row_margin(self._row_margins[0]), 1 - self.__adjust_row_margin(self._row_margins[1]), 0, 1) # Create the frame now everything is in place self.__create_frame(**kwargs) # adjust some frame axis properties self.__style_frame() # draw objects and legends self.__draw_objects() self.__draw_lines() self.__draw_legends() self.__draw_text() if prev_pad: # Don't spoil what the user might want to do afterwards prev_pad.cd()
class Plot(object): """Structural class for representing, accessing, and maintaining references to ROOT graphical elements forming a plot, potentially with a ratio subplot. """ # Plotting 'constants' for the plot class. Ideally, one would allow these # to be flexible, but unfortunately ROOT's coordinate system is extremely # inconsistent and fragile, so it is best to fix these values here. You # can change them dynamically with Plot.Whatever = value, but it is # probably best to leave them alone. # TODO: 600x600 and 800x600 are the ATLAS default for square and # rectangular plots respectively. Fix this when everything is calm. PLOT_WIDTH = 1280 # px PLOT_HEIGHT = 1024 # px #PLOT_MARGINS = (0.125, 0.05, 0.1, 0.1) # Left, Right, Bottom, Top PLOT_MARGINS = (0.125, 0.05, 0.1, 0.07) # Left, Right, Bottom, Top PLOT_MARGINS_WITH_RATIO = (0.125, 0.05, 0.025, 0.1) PLOT_RATIO_MARGINS = (0.125, 0.05, 0.325, 0.05) PLOT_TITLE_X = 0.5 PLOT_TITLE_Y = 0.95 PLOT_TITLE_TEXT_SIZE = 0.04 PLOT_TITLE_TEXT_COLOR = 1 PLOT_TITLE_TEXT_FONT = 42 PLOT_HEADER_HEIGHT = 400 # px PLOT_LEGEND_LEFT = 0.45 PLOT_LEGEND_RIGHT = 0.95 PLOT_LEGEND_BOTTOM = 0.7 PLOT_LEGEND_BOTTOM_WITH_RATIO = 0.63 PLOT_LEGEND_TOP = 0.90 PLOT_LEGEND_TOP_WITH_RATIO = 0.86 PLOT_LEGEND_TEXT_SIZE = 0.025 PLOT_LEGEND_TEXT_SIZE_WITH_RATIO = 0.03 PLOT_LEGEND_ROW_SIZE = 0.04 PLOT_LEGEND_ROW_SIZE_WITH_RATIO = 0.045 PLOT_LEGEND_N_COLUMNS = 1 PLOT_LEGEND_PIVOT_COLUMNS = True PLOT_STAT_LEFT = 0.55 PLOT_STAT_LEFT_WITH_RATIO = 0.60 PLOT_STAT_RIGHT = 0.85 PLOT_STAT_RIGHT_WITH_RATIO = 0.93 PLOT_STAT_BOTTOM = 0.15 PLOT_STAT_BOTTOM_WITH_RATIO = 0.07 PLOT_STAT_TOP = 0.4 PLOT_STAT_TOP_WITH_RATIO = 0.45 PLOT_STAT_TEXT_FONT = 42 PLOT_STAT_TEXT_SIZE = 0.03 PLOT_STAT_TEXT_SIZE_WITH_RATIO = 0.04 PLOT_RATIO_FRACTION = 0.3 # fraction of canvas height PLOT_X_AXIS_TITLE_SIZE = 0.042 PLOT_X_AXIS_TITLE_SIZE_WITH_RATIO = 0.14 PLOT_X_AXIS_TITLE_OFFSET = 0.95 PLOT_X_AXIS_TITLE_OFFSET_WITH_RATIO = 0.96 PLOT_X_AXIS_LABEL_SIZE_WITH_RATIO = 0.12 PLOT_Y_AXIS_LABEL_OFFSET = 0.01 PLOT_Y_AXIS_TITLE_SIZE = 0.042 PLOT_Y_AXIS_TITLE_SIZE_WITH_RATIO = 0.06 PLOT_Y_AXIS_TITLE_OFFSET = 1.0 PLOT_Y_AXIS_TITLE_OFFSET_WITH_RATIO = 0.95 PLOT_Y_AXIS_LABEL_SIZE_WITH_RATIO = 0.05 PLOT_RATIO_Y_AXIS_TITLE_SIZE = 0.12 PLOT_RATIO_Y_AXIS_TITLE_OFFSET = 0.40 PLOT_RATIO_Y_AXIS_LABEL_SIZE = 0.12 PLOT_RATIO_Y_AXIS_LABEL_OFFSET = PLOT_Y_AXIS_LABEL_OFFSET PLOT_RATIO_Y_AXIS_NDIVISIONS = 504 PLOT_RATIO_Y_AXIS_MINIMUM = 0.6 PLOT_RATIO_Y_AXIS_MAXIMUM = 1.4 PLOT_ERROR_BAND_FILL_STYLE = 3254 # Diagonal lines PLOT_ERROR_BAND_FILL_COLOR = 13 # Gray PLOT_ERROR_BAND_LINE_WIDTH = 0 PLOT_ERROR_BAND_LINE_COLOR = 0 PLOT_RATIO_ERROR_BAND_FILL_STYLE = 3254 # Diagonal lines PLOT_RATIO_ERROR_BAND_FILL_COLOR = 807 # Orange PLOT_RATIO_ERROR_BAND_LINE_WIDTH = 0 PLOT_RATIO_ERROR_BAND_LINE_COLOR = 0 # Stamp settings PLOT_ATLAS_STAMP_TEXT_SIZE = 0.035 PLOT_ATLAS_STAMP_TEXT_SIZE_WITH_RATIO = 0.05 PLOT_ATLAS_STAMP_TEXT_COLOR = 1 PLOT_ATLAS_STAMP_TEXT_FONT = 42 PLOT_ATLAS_STAMP_LEFT = 0.18 PLOT_ATLAS_STAMP_TOP = 0.875 PLOT_ATLAS_STAMP_TOP_WITH_RATIO = 0.82 # Stamp specializations PLOT_ATLAS_STAMP_ATLAS_TEXT_FONT = 72 PLOT_ATLAS_STAMP_ATLAS_LABEL_LEFT = 0.28 PLOT_ATLAS_STAMP_LUMINOSITY_OFFSET = 0.036 PLOT_ATLAS_STAMP_LUMINOSITY_OFFSET_WITH_RATIO = 0.05 PLOT_ATLAS_STAMP_LUMINOSITY_SIZE = 0.062 PLOT_ATLAS_STAMP_LUMINOSITY_SIZE_WITH_RATIO = 0.085 def __init__(self, title='', x_title=None, y_title=None, plot_header=True, ratio=False, x_range=None, y_max=None, y_log_scale=False): """Initializes a new instance of the Plot class. Args: title: The title to set for the histogram plot_header: Whether or not to include whitespace at the top of the plot for the ATLAS label and legend ratio: Whether or not to include a ratio plot x_range: A tuple of (x_min, x_max) y_max: The maximum Y axis value y_log_scale: Use log scale for Y axis """ # Store the title self._title = title self._x_title, self._y_title = x_title, y_title # Store whether or not the user wants to create a plot header self._plot_header = plot_header # Calculate a unique name for the plot components name = _rand_uuid() # Create a canvas self._canvas = TCanvas(name + '_canvas', name, int(self.PLOT_WIDTH), int(self.PLOT_HEIGHT)) SetOwnership(self._canvas, False) # Create the main plot and draw it self._plot = TPad(name + '_plot', name, 0.0, (self.PLOT_RATIO_FRACTION if ratio else 0.0), 1.0, 1.0) SetOwnership(self._plot, False) self._plot.SetMargin( *(self.PLOT_MARGINS_WITH_RATIO if ratio else self.PLOT_MARGINS)) self._plot.Draw() # HACK: Draw the plot title. # https://root.cern.ch/phpBB3/viewtopic.php?t=18282. Wonderful. self._draw_title() # Store ranges self._x_range = x_range if y_max is not None: self._set_maximum_value(y_max) # Store log scale self._y_log_scale = y_log_scale # Switch back to the context of the canvas self._canvas.cd() # Create a ratio plot and draw it if requested if ratio: self._ratio_plot = TPad(name + '_ratio', name, 0.0, 0.0, 1.0, self.PLOT_RATIO_FRACTION) SetOwnership(self._ratio_plot, False) self._ratio_plot.SetMargin(*self.PLOT_RATIO_MARGINS) self._ratio_plot.SetGridy(True) self._ratio_plot.Draw() else: self._ratio_plot = None # Track whether or not we've already drawn to the main pad self._drawn = False # Track whether or not we've already drawn to the ratio pad self._ratio_drawn = False # Track that object which sets up the axes in the main plot self._axes_object = None # Create a structure to track any histograms we generate internally # which need to be added to any legends created self._legend_extras = [] # Create lists of the cloned drawables, just to be certain self._drawables = [] self._ratio_drawables = [] def save(self, path, extensions=['pdf']): """Saves this plot to file. Args: path: The path where the plot should be saved. """ # Force an update of the canvas self._canvas.Update() # Save to file for e in extensions: self._canvas.SaveAs(path + '.' + e) def _get_maximum_value(self): """Returns the currently set maximum value (possibly None). """ if hasattr(self, '_maximum_value'): return self._maximum_value return None def _set_maximum_value(self, value): """Sets the current maximum value, possibly including room for a plot header. Args: value: The value to set """ # Check if the current value is not None, and if so, throw an error # because this property should not be set twice if self._get_maximum_value() is not None: raise RuntimeError('maximum value should not be set twice') # If the value is None, ignore it if value is None: return # If the user wants a plot header, then add space for one if self._plot_header: # Grab the plot pad height (in pixels) plot_height = (self.PLOT_HEIGHT * (self._plot.GetY2() - self._plot.GetY1())) # Adjust the height value *= (plot_height + self.PLOT_HEADER_HEIGHT) / plot_height # Set the value self._maximum_value = value def draw(self, *drawables_styles_options): """Plots a collection of plottables to the main plot pad. All TH1 objects are drawn with error bars. THStack elements are only drawn with an error band if one is provided. This method may only be called once Args: drawables_styles_options: Each argument of this function must be of the form (object, style, options), where object is one of the following: - A TH1 object - A TH2 object - A THStack object - A tuple of the form (THStack, TGraph) where the latter represents error bars - A TGraph object - A TLine object style is a tuple of the form (line_color, fill_color, marker_style), and options is a string which will be used for the options argument of the object's Draw method. Plottables will be rendered in the order provided. Axes drawing options (e.g. 'a' or 'same' should not be provided and will be set automatically). A TLine may not be the first drawable element. """ # Make sure there are drawables if len(drawables_styles_options) == 0: raise ValueError('must provide at least one plottable') # Check if we've already drawn if self._drawn: raise RuntimeError('cannot draw twice to a plot') self._drawn = True # Remove None-valued drawables drawables_styles_options = tuple( ((d, s, o) for d, s, o in drawables_styles_options if valid_drawable(d))) # Extract drawables drawables, _, _ = zip(*drawables_styles_options) # Check if there is a maximum value set, and if not, set it if self._get_maximum_value() is None: self._set_maximum_value(maximum_value(drawables)) # Move to the context of the plot pad self._plot.cd() # Iterate through and draw drawables based on type first = True for drawable, style, option in drawables_styles_options: # Check if this a tuple of histogram, error_band if isinstance(drawable, tuple): drawable, error_band = drawable else: error_band = None # Make a clone of the drawable so we don't modify it o = clone(drawable) SetOwnership(o, False) # Add it to the list of drawables self._drawables.append(o) # Set the title appropriately if not is_line(o): o.SetTitle(drawable.GetTitle()) # Style the drawable before it is drawn if style is not None: if is_line(drawable) or is_function(drawable): if isinstance(style, dict): style_line(o, **style) else: style_line(o, *style) else: if isinstance(style, dict): style_histogram(o, **style) else: style_histogram(o, *style) # Set the maximum value of the drawable if supported # HACK: I wish this could go into _handle_axes, but apparently it # can't because ROOT sucks and this has to be set on EVERY # drawable, not just the one with the axes. if is_scatter(o): o.SetMinimum(1 if self._y_log_scale else 0) if is_histo(o) or is_graph(o) or is_stack(o) or is_function(o): o.SetMaximum(self._get_maximum_value()) # With TGraph, this is sometimes necessary. Perhaps with TH1 # too. I'm not sure what happens if we set log scale, but # we'll cross that bridge then. o.SetMinimum(1 if self._y_log_scale else 0) # Include axes if we need to. Store the x-axis range. if first: if is_line(o): raise ValueError('TLine may not be first drawable') if is_graph(o): option += 'a' else: option += 'same' first = False # Draw the drawable o.Draw(option) # TODO: This method of plotting the stats box is a huge hack. We # should plot fit functions separately, and move the stats box # plotting to its own function. Home grow everything, the only # way you can make ROOT work. if is_graph(o) or is_histo(o): if len(o.GetListOfFunctions()) > 0: # HACK: Need to call Update() to paint the fit stats self._plot.Update() stats = o.FindObject("stats") if stats: stats.SetTextFont(Plot.PLOT_STAT_TEXT_FONT) stats.SetTextSize( (Plot.PLOT_STAT_TEXT_SIZE_WITH_RATIO if self._ratio_plot else Plot.PLOT_STAT_TEXT_SIZE)) stats.SetX1NDC( (Plot.PLOT_STAT_LEFT_WITH_RATIO if self._ratio_plot else Plot.PLOT_STAT_LEFT)) stats.SetY1NDC( (Plot.PLOT_STAT_BOTTOM_WITH_RATIO if self._ratio_plot else Plot.PLOT_STAT_BOTTOM)) stats.SetX2NDC( (Plot.PLOT_STAT_RIGHT_WITH_RATIO if self._ratio_plot else Plot.PLOT_STAT_RIGHT)) stats.SetY2NDC( (Plot.PLOT_STAT_TOP_WITH_RATIO if self._ratio_plot else Plot.PLOT_STAT_TOP)) # Handle axes if not is_line(o): self._handle_axes(o, option) # If there is an error band, draw it if error_band is not None: self._draw_error_band(error_band) if self._y_log_scale: self._plot.SetLogy(1) # TODO: Verify this. It breaks 2D plotting. # HACK: Need to force a redraw of plot axes due to issue with ROOT: # http://root.cern.ch/phpBB3/viewtopic.php?f=3&t=14034 #self._plot.RedrawAxis() def _handle_axes(self, drawable, option): """If there is no object currently registered as the owner of the axes drawn on the main plot, then this will set it. Args: drawable: The graph, histogram or stack whose axes were ALREADY drawn option: The option with which to draw the axes """ # If we already have an axes object, ignore this one if self._axes_object is not None: return # Grab the histogram used for axes style/range manipulation if is_stack(drawable) or is_graph(drawable): axes_histogram = drawable.GetHistogram() else: axes_histogram = drawable self._axes_object = axes_histogram # Grab the histogram used for title manipulation if is_stack(drawable): title_histogram = drawable.GetHists()[0] else: title_histogram = drawable # Grab axes x_axis, y_axis = axes_histogram.GetXaxis(), axes_histogram.GetYaxis() # Grab titles from first histogram if not set explicitly if self._x_title is None: self._x_title = title_histogram.GetXaxis().GetTitle() if self._y_title is None: self._y_title = title_histogram.GetYaxis().GetTitle() if self._x_range is not None: #x_axis.SetRangeUser(*self._x_range) x_axis.SetLimits(*self._x_range) # Style x-axis, or hide it if this plot has a ratio plot if self._ratio_plot: x_axis.SetLabelOffset(999) x_axis.SetTitleOffset(999) else: x_axis.SetTitle(self._x_title) x_axis.SetTitleSize(self.PLOT_X_AXIS_TITLE_SIZE) x_axis.SetTitleOffset(self.PLOT_X_AXIS_TITLE_OFFSET) # Style y-axis if self._ratio_plot: y_axis.SetLabelSize(self.PLOT_Y_AXIS_LABEL_SIZE_WITH_RATIO) y_axis.SetLabelOffset(self.PLOT_Y_AXIS_LABEL_OFFSET) y_axis.SetTitle(self._y_title) y_axis.SetTitleSize( (self.PLOT_Y_AXIS_TITLE_SIZE_WITH_RATIO if self._ratio_plot else self.PLOT_Y_AXIS_TITLE_SIZE)) y_axis.SetTitleOffset( (self.PLOT_Y_AXIS_TITLE_OFFSET_WITH_RATIO if self._ratio_plot else self.PLOT_Y_AXIS_TITLE_OFFSET)) # Redraw the drawable with the new style drawable.Draw(option) def _draw_error_band(self, error_band): """Draws an error band on top of histogram objects. Args: error_band: The error band to draw (a TGraphAsymmErrors) """ # Style it # HACK: Setting the marker style to 0 specifies this should be filled # in the legend error_band.SetMarkerStyle(0) error_band.SetMarkerSize(0) error_band.SetFillStyle(self.PLOT_ERROR_BAND_FILL_STYLE) error_band.SetFillColor(self.PLOT_ERROR_BAND_FILL_COLOR) error_band.SetLineWidth(self.PLOT_ERROR_BAND_LINE_WIDTH) error_band.SetLineColor(self.PLOT_ERROR_BAND_LINE_COLOR) # Draw it error_band.Draw('e2same') # Add it to the list of things we need to add to the legend self._legend_extras.append(error_band) def draw_ratio_histogram(self, histogram, draw_unity=True, error_band=None): """Draws a ratio histogram to the ratio pad. Args: histogram: The ratio histogram to draw (use ratio_histogram) draw_unity: Whether or not to draw a line at 1 error_band: An error band to draw under the ratio histogram (see owls_hep.uncertainty.ratio_uncertainty_band) The histogram X axis title is set by draw_histogram if not set explicitly. draw_ratio_histogram should therefore be called after draw_histogram. """ # Check if we've already drawn if self._ratio_drawn: raise RuntimeError('cannot draw twice to a plot') self._ratio_drawn = True # Switch to the context of the ratio pad self._ratio_plot.cd() # Clone the histogram histogram = histogram.Clone(_rand_uuid()) SetOwnership(histogram, False) # Style it x_axis, y_axis = histogram.GetXaxis(), histogram.GetYaxis() x_axis.SetTitleSize(self.PLOT_X_AXIS_TITLE_SIZE_WITH_RATIO) x_axis.SetTitleOffset(self.PLOT_X_AXIS_TITLE_OFFSET_WITH_RATIO) x_axis.SetLabelSize(self.PLOT_X_AXIS_LABEL_SIZE_WITH_RATIO) x_axis.SetTitle(self._x_title) if self._x_range: x_axis.SetLimits(*self._x_range) #x_axis.SetRangeUser(*self._x_range) else: x_axis.SetLimits(self._axes_object.GetXaxis().GetXmin(), self._axes_object.GetXaxis().GetXmax()) #x_axis.SetRangeUser(self._axes_object.GetXaxis().GetXmin(), #self._axes_object.GetXaxis().GetXmax()) y_axis.SetTitleSize(self.PLOT_RATIO_Y_AXIS_TITLE_SIZE) y_axis.SetTitleOffset(self.PLOT_RATIO_Y_AXIS_TITLE_OFFSET) y_axis.SetLabelSize(self.PLOT_RATIO_Y_AXIS_LABEL_SIZE) y_axis.SetLabelOffset(self.PLOT_RATIO_Y_AXIS_LABEL_OFFSET) y_axis.SetRangeUser(self.PLOT_RATIO_Y_AXIS_MINIMUM, self.PLOT_RATIO_Y_AXIS_MAXIMUM) y_axis.SetNdivisions(self.PLOT_RATIO_Y_AXIS_NDIVISIONS, False) # Draw it # NOTE: Have to specify E0 or points out of the vertical range won't # have their error bars drawn: # https://root.cern.ch/phpBB3/viewtopic.php?f=3&t=13329 # histogram.Draw('e0p') # NOTE: Or live with it and get rid of points from zero value bins histogram.Draw('ep') # Draw a line at unity if requested if draw_unity: # Calculate the line coordinates line_min = histogram.GetBinLowEdge(1) max_bin = histogram.GetNbinsX() line_max = (histogram.GetBinLowEdge(max_bin) + histogram.GetBinWidth(max_bin)) # Create and draw the line unit_line = TLine(line_min, 1.0, line_max, 1.0) SetOwnership(unit_line, False) unit_line.SetLineColor(2) # Red unit_line.SetLineWidth(2) unit_line.Draw('same') # If an error band was provided, draw it and add it to our legend # elements if error_band: # Keep ownership of the error band SetOwnership(error_band, False) # Style it error_band.SetMarkerSize(0) error_band.SetFillStyle(self.PLOT_RATIO_ERROR_BAND_FILL_STYLE) error_band.SetFillColor(self.PLOT_RATIO_ERROR_BAND_FILL_COLOR) error_band.SetLineWidth(self.PLOT_RATIO_ERROR_BAND_LINE_WIDTH) error_band.SetLineColor(self.PLOT_RATIO_ERROR_BAND_LINE_COLOR) # Draw it error_band.Draw('e2same') # Now, if we've drawn unity or an error band, redraw our ratio # histogram so that its point lie on top of the unity line or error # band, but use 'same' so that the axes/ticks don't cover the red line if draw_unity or error_band: # histogram.Draw('e0psame') histogram.Draw('epsame') def draw_ratios(self, drawables_styles_options, draw_unity=True, y_range=None, y_title=None): """Draws a drawable to the ratio pad. Args: drawable: The drawable to draw draw_unity: Whether or not to draw a line at 1 """ # Check if we've already drawn if self._ratio_drawn: raise RuntimeError('cannot draw twice to a plot') self._ratio_drawn = True # Switch to the context of the ratio pad self._ratio_plot.cd() # Iterate through and draw drawables based on type first = True for drawable, style, option in drawables_styles_options: # Make a clone of the drawable so we don't modify it o = clone(drawable) SetOwnership(o, False) # Add it to the list of drawables self._ratio_drawables.append(o) # Set the title appropriately if not is_line(o): o.SetTitle(drawable.GetTitle()) # Style the drawable before it is drawn if style is not None: if is_line(drawable) or is_function(drawable): style_line(o, *style) else: style_histogram(o, *style) if not is_line(o): if y_range is not None: o.SetMinimum(y_range[0]) o.SetMaximum(y_range[1]) # Include axes if we need if first: if is_line(o): raise ValueError('TLine may not be first drawable') x_axis, y_axis = o.GetXaxis(), o.GetYaxis() if is_graph(o): option += 'a' else: option += 'same' first = False # Draw the drawable o.Draw(option) x_axis.SetTitleSize(self.PLOT_X_AXIS_TITLE_SIZE_WITH_RATIO) x_axis.SetTitleOffset(self.PLOT_X_AXIS_TITLE_OFFSET_WITH_RATIO) x_axis.SetLabelSize(self.PLOT_X_AXIS_LABEL_SIZE_WITH_RATIO) x_axis.SetTitle(self._x_title) if self._x_range: x_axis.SetLimits(*self._x_range) #x_axis.SetRangeUser(*self._x_range) else: x_axis.SetLimits(self._axes_object.GetXaxis().GetXmin(), self._axes_object.GetXaxis().GetXmax()) #x_axis.SetRangeUser(self._axes_object.GetXaxis().GetXmin(), #self._axes_object.GetXaxis().GetXmax()) y_axis.SetTitleSize(self.PLOT_RATIO_Y_AXIS_TITLE_SIZE) y_axis.SetTitleOffset(self.PLOT_RATIO_Y_AXIS_TITLE_OFFSET) y_axis.SetLabelSize(self.PLOT_RATIO_Y_AXIS_LABEL_SIZE) y_axis.SetLabelOffset(self.PLOT_RATIO_Y_AXIS_LABEL_OFFSET) y_axis.SetNdivisions(self.PLOT_RATIO_Y_AXIS_NDIVISIONS, False) if y_title is not None: y_axis.SetTitle(y_title) self._ratio_plot.Update() def _draw_title(self): """Draws a title on the plot. """ title = TLatex() title.SetTextColor(self.PLOT_TITLE_TEXT_COLOR) title.SetTextFont(self.PLOT_TITLE_TEXT_FONT) title.SetNDC() title.SetTextSize(self.PLOT_TITLE_TEXT_SIZE) title.SetTextAlign(22) title.DrawLatex(self.PLOT_TITLE_X, self.PLOT_TITLE_Y, self._title) def draw_atlas_label(self, luminosity=None, sqrt_s=None, custom_label=None, atlas_label=None): """Draws an ATLAS stamp on the plot, with an optional categorization label. It is recommended that you construct the Plot with plot_header = True in order to make space for the label. Args: luminosity: The integrated luminosity, in pb^-1 sqrt_s: The center of mass energy, in MeV label: The label to put after 'ATLAS', None to exclude the 'ATLAS' categorization entirely """ # Change context to the plot pad self._plot.cd() # Create the latex object # TODO: Consider using TPaveText to overwrite drawn graphs and # histograms. At least for scatter plots. # TODO: Increase readability: Create two sets of constants - # one for plots with ratio and one for plots without. Select the # correct one in draw_ratio_histogram. stamp = TLatex() # Style it stamp.SetTextColor(self.PLOT_ATLAS_STAMP_TEXT_COLOR) stamp.SetTextSize( (self.PLOT_ATLAS_STAMP_TEXT_SIZE_WITH_RATIO if self._ratio_plot else self.PLOT_ATLAS_STAMP_TEXT_SIZE)) stamp.SetTextFont(self.PLOT_ATLAS_STAMP_TEXT_FONT) stamp.SetNDC() top = (self.PLOT_ATLAS_STAMP_TOP_WITH_RATIO if self._ratio_plot else self.PLOT_ATLAS_STAMP_TOP) # Print an ATLAS label on top if atlas_label is not None: # Draw the label stamp.SetTextFont(self.PLOT_ATLAS_STAMP_ATLAS_TEXT_FONT) stamp.DrawLatex(self.PLOT_ATLAS_STAMP_LEFT, top, 'ATLAS') stamp.SetTextFont(self.PLOT_ATLAS_STAMP_TEXT_FONT) stamp.DrawLatex(self.PLOT_ATLAS_STAMP_ATLAS_LABEL_LEFT, top, atlas_label) top -= (self.PLOT_ATLAS_STAMP_TEXT_SIZE_WITH_RATIO if self. _ratio_plot else self.PLOT_ATLAS_STAMP_TEXT_SIZE) * 1.3 # Draw the luminosity and sqrt(s) if luminosity is not None or sqrt_s is not None: text = '' if sqrt_s is not None: text += '#sqrt{{s}} = {0:.0f} TeV'.format(sqrt_s / 1.0e6) if luminosity is not None: text += ', ' if luminosity is not None: if luminosity >= 1000.0: text += '{0:.1f} fb^{{-1}}'.format(luminosity / 1000.0) elif luminosity > 100.0: text += '{0:.2f} fb^{{-1}}'.format(luminosity / 1000.0) else: text += '{0:.1f} pb^{{-1}}'.format(luminosity) stamp.DrawLatex(self.PLOT_ATLAS_STAMP_LEFT, top, text) top -= (self.PLOT_ATLAS_STAMP_TEXT_SIZE_WITH_RATIO if self. _ratio_plot else self.PLOT_ATLAS_STAMP_TEXT_SIZE) * 1.3 # If requested, draw the custom label or the 'ATLAS' label, # preferring the former if custom_label is not None: # Draw each line of text, decreasing top for each step for text in [t for t in custom_label if t is not None]: stamp.DrawLatex(self.PLOT_ATLAS_STAMP_LEFT, top, text) top -= (self.PLOT_ATLAS_STAMP_TEXT_SIZE_WITH_RATIO if self. _ratio_plot else self.PLOT_ATLAS_STAMP_TEXT_SIZE) * 1.3 def draw_pave(self, texts, position): """Draw a text box at the position and fill it with text. Args: texts: String or N-tuple of strings. positon: Absolute position of the form (x1, x2, y1, y2) as an N-tuple of floats, or a string with one of the values "{top,bottom}{left,right}" """ if isinstance(position, basestring): if position == 'topleft': position = (0.15, 0.35, 0.85, 0.75) elif position == 'topright': position = (0.95, 0.75, 0.95, 0.85) elif position == 'bottomleft': position = (0.05, 0.25, 0.05, 0.15) elif position == 'topright': position = (0.95, 0.75, 0.05, 0.15) # Switch to the context of the main plot self._plot.cd() # Create the pave self._pave = TPaveText(position[0], position[2], position[1], position[3], 'NDC') SetOwnership(self._pave, False) # Add the text if isinstance(texts, basestring): self._pave.AddText(texts) else: for t in texts: self._pave.AddText(t) # Draw the pave self._pave.Draw() def draw_legend(self, use_functions=False, legend_entries=None): """Draws a legend onto the plot with the specified histograms. It is recommended that you construct the Plot with plot_header = True in order to make space for the legend. Args: drawables: The elements to include in the legend (via AddEntry) use_functions: Add associated functions to the legend """ # Check if we already have a legend if hasattr(self, '_legend'): raise RuntimeError('legend already exists on this plot') # Check if the plot has been drawn if not self._drawn: raise RuntimeError('plot must be drawn before the legend') # Remove None-valued drawables drawables = tuple((d for d in self._drawables if d is not None)) # Remove TLine objects drawables = tuple((d for d in drawables if not is_line(d))) # If we shouldn't add functions to the legend, remove them if not use_functions: drawables = tuple((d for d in drawables if not isinstance(d, TF1))) # Use only certain entries if legend_entries is not None: def get_drawable_by_title(title): for d in drawables: if d.GetTitle() == title: return d return None drawables = tuple((get_drawable_by_title(e.GetTitle()) for e in legend_entries if e is not None)) drawables = tuple((d for d in drawables if d is not None)) # Switch to the context of the main plot self._plot.cd() # Create the legend self._legend = TLegend(self.PLOT_LEGEND_LEFT, (self.PLOT_LEGEND_BOTTOM_WITH_RATIO if self._ratio_plot else self.PLOT_LEGEND_BOTTOM), self.PLOT_LEGEND_RIGHT, (self.PLOT_LEGEND_TOP_WITH_RATIO if self._ratio_plot else self.PLOT_LEGEND_TOP)) SetOwnership(self._legend, False) # Style it self._legend.SetTextSize( (self.PLOT_LEGEND_TEXT_SIZE_WITH_RATIO if self._ratio_plot else self.PLOT_LEGEND_TEXT_SIZE)) self._legend.SetBorderSize(0) self._legend.SetFillStyle(0) # transparent self._legend.SetNColumns(self.PLOT_LEGEND_N_COLUMNS) # Create a chained list of all drawables. We decompose THStack # objects in reverse order, i.e. top-to-bottom. drawables = \ list(chain(*(drawable_iterable(h, True, True) for h in drawables))) # Add anything to this list that we created internally drawables.extend(self._legend_extras) # Because ROOT draws legend entries from left-to-right across rows and # not top-to-bottom along columns, we need to do a bit of a pivot on # the list so that the histograms appear in the vertical order of the # stack if self.PLOT_LEGEND_PIVOT_COLUMNS: n_entries = len(drawables) n_col = self.PLOT_LEGEND_N_COLUMNS n_row = int(ceil(float(n_entries) / n_col)) self._legend.SetY1(self._legend.GetY2() - n_row * (self.PLOT_LEGEND_ROW_SIZE_WITH_RATIO if self. _ratio_plot else self.PLOT_LEGEND_ROW_SIZE)) legend_order = [] for r in xrange(0, n_row): for c in xrange(0, n_col): if (r * n_col + c) == n_entries: # Don't need an outer break, this would only happen on the # last row if n_row * n_col != n_entries break legend_order.append(drawables[r + c * n_row]) else: legend_order = drawables # Add the drawables for drawable in legend_order: SetOwnership(drawable, False) title = drawable.GetTitle() # HACK: Convention: legend for drawables with a non-default # marker style (data) to be drawn as line with point, and with # empty fill (signal) to be drawn as line if drawable.GetMarkerStyle() != 0: self._legend.AddEntry(drawable, title, 'ep') elif drawable.GetFillColor() == 0: self._legend.AddEntry(drawable, title, 'l') else: self._legend.AddEntry(drawable, title, 'f') # Draw the legend self._legend.Draw()