Skip to content

Instantly share code, notes, and snippets.

@ram1123

ram1123/Python_PlottingHelperFunctions.md

Last active January 18, 2018 18:18
Show Gist options
  • Save ram1123/774b5a2f1556b25777e36c71ef0c3963 to your computer and use it in GitHub Desktop.
Save ram1123/774b5a2f1556b25777e36c71ef0c3963 to your computer and use it in GitHub Desktop.
Download ZIP
Some function collection that helps in plotting
Raw
Python_PlottingHelperFunctions.md

Time check

import ROOT as r
r.gROOT.SetBatch(True)
clock = r.TStopwatch()
clock.Start()
print "Note the running time..."
clock.Stop()

print 'Tree loop profiling stats:'
print 'Real Time used:', clock.RealTime(),"seconds"
print 'CPU Time used:', clock.CpuTime()

Read file and its tree

tfile = r.TFile("root:://cmseos.fnal.gov//eos/uscms/store/user/rasharma/LHE_GEN_Analyzer_Output/WPlepWMhad_aQGC.root")
ttree = tfile.Get("otree")

Get Ratio of two histograms

def createRatio(h1, h2, xlabel, col):
    h3 = h1.Clone("h3")
    h3.SetLineColor(col)
    h3.SetMarkerColor(col)
    h3.SetMarkerStyle(21)
    h3.SetTitle("")
    h3.SetMinimum(0.5)
    h3.SetMaximum(1.5)
    # Set up plot for markers and errors
    #h3.Sumw2()
    h3.SetStats(0)
    h3.Divide(h2)

    # Adjust y-axis settings
    y = h3.GetYaxis()
    y.SetTitle("ratio (SM/aQGC)")
    y.SetNdivisions(505)
    y.SetTitleSize(20)
    y.SetTitleFont(43)
    y.SetTitleOffset(1.55)
    y.SetLabelFont(43)
    y.SetLabelSize(15)

    # Adjust x-axis settings
    x = h3.GetXaxis()
    x.SetTitle(xlabel)
    x.SetTitleSize(20)
    x.SetTitleFont(43)
    x.SetTitleOffset(4.0)
    x.SetLabelFont(43)
    x.SetLabelSize(15)

    return h3

Create Canvas Pads

def createCanvasPads():
    c = r.TCanvas("c", "canvas", 600, 700)
    # Upper histogram plot is pad1
    pad1 = r.TPad("pad1", "pad1", 0.0, 0.3, 1.0, 1.0)
    pad1.SetBottomMargin(0)  # joins upper and lower plot
    pad1.SetLeftMargin(0.1)
    pad1.SetRightMargin(0.03)
    pad1.SetLogy(1)
    pad1.SetGridx()
    pad1.SetTickx(1)
    pad1.SetTicky(1)
    pad1.Draw()
    # Lower ratio plot is pad2
    c.cd()  # returns to main canvas before defining pad2
    pad2 = r.TPad("pad2", "pad2", 0.0, 0.00, 1.0, 0.3)
    #pad2.SetLogy(1)
    pad2.SetTopMargin(0)  # joins upper and lower plot
    pad2.SetBottomMargin(0.25)
    pad2.SetLeftMargin(0.1)
    pad2.SetRightMargin(0.03)
    pad2.SetGridx()
    pad2.SetTickx(1)
    pad2.SetTicky(1)
    pad2.Draw()

    return c, pad1, pad2

Get histogram from tree

def getHistFromFile (tree, var, weight, nbin, xmin, xmax, xlabel, col):  
    hist = r.TH1F("hist", xlabel, nbin, xmin, xmax)
    tree.Draw(var+ ">>hist","LHEWeights["+str(weight)+"]","")
    #tree.Draw(var+ ">>hist","LHEWeights["+str(weight)+"]","",1000)
    if not hist:
        print 'Failed to get hist from file'
        exit(0)
    hist.GetXaxis().SetTitle(xlabel)
    ylabel = "Events / %s GeV" % int(hist.GetBinWidth(1))
    hist.GetYaxis().SetTitle(ylabel)
    hist.SetLineColor(col)
    hist.SetMarkerColor(col)
    hist.SetDirectory(r.gROOT) # detach "hist" from the file
    return hist

Get Legends

def getLegends(pos,ncol,nvar,fontsize):
    if pos == "tr":
        legend = ROOT.TLegend(.70-(0.15*(ncol-1)), 0.80-(0.02*(nvar/ncol-1)) ,.95 ,.920)
    elif pos == "tl":
        legend = ROOT.TLegend(0.11, 0.85-(0.02*(nvar/ncol-1)) ,.80+(0.15*(ncol-1)) ,.920)
    elif pos == "tc":
        legend = ROOT.TLegend(0.30, 0.80-(0.02*(nvar/ncol-1)) ,.80 ,.920)
    elif pos == "bl":
        legend = ROOT.TLegend(0.11, 0.25+(0.02*(nvar/ncol-1)) ,.60+(0.15*(ncol-1)) ,.130)
    elif pos == "bc":
        legend = ROOT.TLegend(0.30, 0.25+(0.02*(nvar/ncol-1)) ,.80 ,.130)
    elif pos == "br":
        legend = ROOT.TLegend(.45-(0.15*(ncol-1)), 0.25+(0.02*(nvar/ncol-1)) ,.95 ,.130)
    else:
        print "Invalid default position: Switching to default legend position top-right"
        legend = ROOT.TLegend(.85-(0.15*(ncol-1)), 0.88-(0.02*(nvar/ncol-1)) ,.50 ,.950)
    legend.SetFillColor(ROOT.kWhite)
    legend.SetFillStyle(0)
    legend.SetTextSize(fontsize)
    legend.SetNColumns(ncol)
    return legend

Get Canvas

def getCanvas():
    H_ref = 600; 
    W_ref = 600; 
    W = W_ref
    H = H_ref

    T = 0.08*H_ref
    B = 0.12*H_ref 
    L = 0.18*W_ref
    R = 0.01*W_ref

    canvas = ROOT.TCanvas("c2","c2",50,50,W,H)
    canvas.SetFillColor(0)
    canvas.SetBorderMode(0)
    canvas.SetFrameFillStyle(0)
    canvas.SetFrameBorderMode(0)
    canvas.SetLeftMargin( L/W *0.6 )
    canvas.SetRightMargin( R/W*5 )
    canvas.SetTopMargin( T/H )
    canvas.SetBottomMargin( B/H )
    canvas.SetTickx(0)
    canvas.SetTicky(0) 
    return canvas

Set histogram attributes

def setHistAttributes (hist, plot_info, line_color, fill_color):
    #hist.SetFillColor(fill_color)
    hist.SetLineColor(line_color)
    hist.SetMarkerColor(line_color)
    #colorArray = [1, 632, 416, 600, 400, 616, 432, 920, 800, 880, 820, 840, 860, 900, 940, 960, 980, 640, 660, 680, 700, 720, 740, 760]
    #stl=1
    #if(line_color==1):
    #   stl=1
    #if(line_color==632):
    #   stl=2
    #if(line_color==416):
    #   stl=3
    #if(line_color==600):
    #   stl=4
    #hist.SetMarkerStyle(stl)
    #hist.SetMarkerSize(1.5)
    #hist.SetLineWidth(2)
    hist.SetTitleOffset(1.3, "x")
    hist.SetTitleOffset(1.3, "y")
    if plot_info["rebin"] != 0:
        if type(hist) != "<class '__main__.TH2F'>":
            hist.Rebin(plot_info["rebin"])
        else:
            print 'Rebin only defined for 1D hist. Use --rebin2D instead.'
    if plot_info["xmin"] < plot_info["xmax"]:
        hist.GetXaxis().SetRangeUser(plot_info["xmin"], plot_info["xmax"])
    if plot_info["ymin"] < plot_info["ymax"]:
        #hist.GetYaxis().SetRangeUser(plot_info["ymin"], plot_info["ymax"])
        #hist.SetMinimum(plot_info["ymin"])
        hist.SetMaximum(plot_info["ymax"])

    #hist.GetYaxis().SetTitle(plot_info["ylabel"])
Raw
SaveAllCanvas.py
import ROOT as r
file = r.TFile.Open("GE11-X-S-CERN-0007_cutADC500_hitTime4-20_clustTime5-20.root")
dir = file.Get("Summary")
# To list contents of dir
#dir.ls()
# Get one histogram/canvas in histogram h
#h = dir.Get("canv_GE11-X-S-CERN-0007_HitPos_AllEta_Segmented")
#Draw h
#h.Draw()
# Save histogram
#h.SaveAs("test.pdf")
# Get all the keys in a directory "Summary"
keys = dir.GetListOfKeys()
for k in keys:
#print k.GetName()
if (k.GetName().find("canv") != -1):
print k.GetName()
h = dir.Get(k.GetName())
h.Draw()
name = r.TString(k.GetName())+r.TString(".pdf")
#Tstring method Data() converts it to char* https://root.cern.ch/doc/master/classTString.html#a7e4ada8bf2407b461f3ee58b71c5959f
h.SaveAs(name.Data())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment