python/plotting/ntuplePlotting.py

0001 import collections
0002 import itertools
0003
0004 import ROOT
0005
0006 import Validation.RecoTrack.plotting.plotting as plotting
0007
0008 def saveHistograms(tdirectory, histos):
0009     for h in histos:
0010         h.SetDirectory(tdirectory)
0011
0012 def applyStyle(h, color, markerStyle):
0013     h.SetMarkerStyle(markerStyle)
0014     h.SetMarkerColor(color)
0015     h.SetMarkerSize(1.2)
0016     h.SetLineColor(color)
0017     h.SetLineWidth(2)
0018
0019 # https://stackoverflow.com/questions/6076270/python-lambda-function-in-list-comprehensions
0020 _defaultStyles = [(lambda c, m: (lambda h: applyStyle(h, c, m)))(color, ms) for color, ms in itertools.izip(plotting._plotStylesColor, plotting._plotStylesMarker)]
0021
0022 _ratioFactor = 1.25
0023
0024 def draw(name, histos, styles=_defaultStyles, legendLabels=[], **kwargs):
0025     width = 600
0026     height = 600
0027     ratioFactor = 1.25
0028
0029     args = {}
0030     args.update(kwargs)
0031     if not "ratioFactor" in args:
0032         args["ratioFactor"] = _ratioFactor
0033     ratio = args.get("ratio", False)
0034
0035     if ratio:
0036         height = int(height*ratioFactor)
0037     c = plotting._createCanvas(name, width, height)
0038     if ratio:
0039         plotting._modifyPadForRatio(c, ratioFactor)
0040
0041     frame = drawSingle(c, histos, styles, **args)
0042
0043     if len(legendLabels) > 0:
0044         if len(legendLabels) != len(histos):
0045             raise Exception("Got %d histos but %d legend labels" % (len(histos), len(legendLabels)))
0046         lx1 = 0.6
0047         lx2 = 0.9
0048         ly1 = 0.7
0049         ly2 = 0.85
0050
0051         lx1 += legendDx
0052         lx2 += legendDx
0053         ly1 += legendDy
0054         ly2 += legendDy
0055
0056         lx2 += legendDw
0057         ly1 -= legendDh
0058
0059         legend = ROOT.TLegend(lx1, ly1, lx2, ly2)
0060         legend.SetLineColor(1)
0061         legend.SetLineWidth(1)
0062         legend.SetLineStyle(1)
0063         legend.SetFillColor(0)
0064         legend.SetMargin(0.07)
0065         legend.SetBorderSize(0)
0066
0067         for h, l in zip(histos, legendLabels):
0068             legend.AddEntry(h, l, "L")
0069
0070         legend.Draw()
0071
0072     frame._pad.cd()
0073
0074     c.Update()
0075     c.RedrawAxis()
0076     c.SaveAs(name+".png")
0077     c.SaveAs(name+".pdf")
0078
0079
0080 def drawSingle(pad, histos, styles=_defaultStyles,
0081                nrows=1,
0082                xtitle=None, ytitle=None,
0083                drawOpt="HIST",
0084                legendDx=0, legendDy=0, legendDw=0, legendDh=0,
0085                xmin=None, ymin=0, xmax=None, ymax=None, xlog=False, ylog=False,
0086                xgrid=True, ygrid=True,
0087                ratio=False, ratioYmin=0.5, ratioYmax=1.5, ratioYTitle=plotting._ratioYTitle, ratioFactor=1.25):
0088
0089     bounds = plotting._findBounds(histos, ylog, xmin, xmax, ymin, ymax)
0090     if ratio:
0091         ratioBounds = (bounds[0], ratioYmin, bounds[2], ratioYmax)
0092         frame = plotting.FrameRatio(pad, bounds, ratioBounds, ratioFactor, ratioYTitle=ratioYTitle, nrows=nrows)
0093         #frame._frameRatio.GetYaxis().SetLabelSize(0.12)
0094     else:
0095         frame = plotting.Frame(pad, bounds, nrows=nrows)
0096
0097     if xtitle is not None:
0098         frame.setXTitle(xtitle)
0099     if ytitle is not None:
0100         frame.setYTitle(ytitle)
0101
0102     frame.setLogx(xlog)
0103     frame.setLogy(ylog)
0104     frame.setGridx(xgrid)
0105     frame.setGridy(ygrid)
0106
0107     if ratio:
0108         frame._pad.cd()
0109     for i, h in enumerate(histos):
0110         st = styles[i%len(styles)]
0111         st(h)
0112         h.Draw(drawOpt+" same")
0113
0114     ratios = None
0115     if ratio:
0116         frame._padRatio.cd()
0117         ratios = plotting._calculateRatios(histos)
0118         for r in ratios[1:]:
0119             r.draw()
0120         frame._pad.cd()
0121
0122     return frame
0123
0124
0125 def drawMany(name, histoDicts, styles=_defaultStyles, opts={}, ncolumns=4):
0126     if len(histoDicts) == 0:
0127         return
0128
0129     histoNames = histoDicts[0].keys()
0130     ratio = False
0131     ratioFactor = _ratioFactor
0132     for opt in opts.values():
0133         if "ratio" in opt:
0134             ratio = True
0135         if "ratioFactor" in opt:
0136             ratioFactor = max(ratioFactor, opt["ratioFactor"])
0137
0138     nhistos = len(histoNames)
0139     nrows = int((nhistos+ncolumns-1)/ncolumns)
0140
0141     width = 500*ncolumns
0142     height = 500*nrows
0143     if ratio:
0144         height = int(_ratioFactor*height)
0145
0146     canvas = plotting._createCanvas(name, width, height)
0147     canvas.Divide(ncolumns, nrows)
0148
0149     histos = collections.defaultdict(list)
0150
0151     for d in histoDicts:
0152         for n, h in d.items():
0153             histos[n].append(h)
0154
0155     for i, histoName in enumerate(histoNames):
0156         pad = canvas.cd(i+1)
0157
0158         args = {}
0159         args.update(opts.get(histoName, {}))
0160         if "ratio" in args:
0161             if not "ratioFactor" in args:
0162                 args["ratioFactor"] = _ratioFactor # use the default, not the max
0163             plotting._modifyPadForRatio(pad, args["ratioFactor"])
0164
0165         frame = drawSingle(pad, histos[histoName], styles, nrows, **args)
0166         frame._pad.cd()
0167         frame._pad.Update()
0168         frame._pad.RedrawAxis()
0169
0170     canvas.SaveAs(name+".png")
0171     canvas.SaveAs(name+".pdf")