class Frame: markers = { "opencircle": (24, 1), "smallopencircle": (24, 0.5), "fullcircle": (20, 1), "smallfullcircle": (20, 0.5), "opensquare": (25, 1), "smallopensquare": (25, 0.5), "fullsquare": (21, 1), "smallfullsquare": (21, 0.5), "opentriangle": (26, 1), "smallopentriangle": (26, 0.5), "fulltriangle": (22, 1), "smallfulltriangle": (22, 0.5), "opencross": (28, 1), "smallopencross": (28, 0.5), "fullcross": (34, 1), "smallfullcross": (34, 0.5), "opendiamond": (27, 1), "smallopendiamond": (27, 0.5), "fulldiamond": (33, 1), "smallfulldiamond": (33, 0.5) } #http://root.cern.ch/root/html/TAttMarker.html def __init__(self, title, xtitle="", xmin=0, xmax=1., ytitle="", ymin=0, ymax=0., option="", nxbin=2, nybin=2, logy=False, logx=False): randomid = "%d" % (random.randint(10000000000, 99999999999)) self.ymin = ymin self.ymax = ymax self.xmin = xmin self.xmax = xmax self.palette = (1, 2, 4, 8, 9, 42, 46, 38, 31) self.logy = logy if logy and self.ymin <= 0: self.ymin = 1.e-5 self.logx = logx if logx and self.xmin <= 0: self.xmin = 1.e-5 if ymax > ymin: self.hF = TH2F(randomid, title, nxbin, xmin, xmax, nybin, ymin, ymax) else: self.hF = TH2F(randomid, title, nxbin, xmin, xmax, nybin, ymin, ymin + 1) ymax = ymin + 1 self.yaxis2 = None self.hF.GetXaxis().SetTitle(xtitle) self.hF.GetYaxis().SetTitle(ytitle) if logy: gPad.SetLogy() if logx: gPad.SetLogx() self.hF.SetStats(0) self.hF.Draw() self.pad = gPad.cd() # yes, the cd() is needed, gPad is special self.items = [] self.options = [] self.labels = [] # list of options (e.g. logy), comma or white-space separated self.frame_options = option.replace(",", " ").split() self.tls = [] self.tlines = [] def process_options(self, h, opts): """ FIXME, work in progress, to be used in add and drawLegend consistently """ if opts == "": optstring = h.GetDrawOption() if optstring == "same": optsring = "L" else: optstring = "" for o in opts.split("+"): if o == "hist": optstring += "L" elif o.lower() in self.markers: optstring += "P" elif o == "L": optstring += "L" elif o == "C": optstring += "C" else: print "unknown option : ", o print "optstring ", optstring, "opts", opts def auto_color(self): return self.palette[(len(self.items) - 1) % len(self.palette)] def auto_marker_color(self): markers = ("fullcircle", "opencircle", "fullsquare", "opensquare", "fulltriangle", "opentriangle") n = len(self.items) - 1 m = n % len(markers) c = (m + n // len(markers)) % len(self.palette) return "small" + markers[m], self.palette[c] def add(self, item, drawOption="", label=None, color=None, size=None, f=None, axis=1, add_to_legend=True): self.items.append(item) self.options.append(drawOption) if add_to_legend == False: self.labels.append(None) elif label: self.labels.append(label) else: try: if item.file.label: self.labels.append(item.file.label) elif item.file.info: self.labels.append(item.file.info) else: self.labels.append("") except AttributeError: # histogram may not have a file attribute self.labels.append("") if color is None and drawOption == "auto": drawOption, color = self.auto_marker_color() self.options[-1] = drawOption elif color is None: color = self.auto_color() item.SetLineColor(color) item.SetMarkerColor(color) if size is not None: item.SetMarkerSize(size) if drawOption == "P" and item.GetMarkerStyle() == 1: item.SetMarkerStyle(20) if drawOption.lower() in self.markers: item.SetMarkerStyle(self.markers[drawOption.lower()][0]) item.SetMarkerSize(self.markers[drawOption.lower()][1]) drawOption = "P" if drawOption.startswith( "L+") and drawOption[2:].lower() in self.markers: item.SetMarkerStyle(self.markers[drawOption[2:].lower()][0]) item.SetMarkerSize(self.markers[drawOption[2:].lower()][1]) drawOption = "PL" if drawOption == "fill": drawOption = "hist" if color: item.SetFillColor(color) if drawOption == "box": item.SetFillColor(0) item.SetLineColor(color) if f is not None: f(item) if axis == 2 and self.yaxis2 is not None: yscale = (self.ymax - self.ymin) / (self.ymax2 - self.ymin2) # rescale for the second axis if type(item) is TGraph: xdata = [x for x in item.GetX()] ydata = [ self.ymin + (y - self.ymin2) * yscale for y in item.GetY() ] for n in range(len(ydata)): item.SetPoint(n, xdata[n], ydata[n]) self.pad.cd() item.Draw("same0 " + drawOption) self.pad.Update() def addLatex(self, x, y, text, textsize=0.3, textcolor=1): tl = TLatex(x, y, text) tl.SetTextSize(textsize) tl.SetTextColor(textcolor) tl.Draw("same") self.tls.append(tl) return tl def addLine(self, x1, y1, x2, y2, arrow="", arrowsize=0.1): if arrow == "": tl = TLine(x1, y1, x2, y2) elif arrow == "<->": tl = TArrow(x1, y1, x2, y2, arrowsize, "<|>") elif arrow == "->": print x2, y2, arrowsize tl = TArrow(x1, y1, x2, y2, arrowsize, "|>") tl.SetAngle(40) tl.SetLineWidth(1) tl.SetFillColor(1) tl.Draw("same") self.tlines.append(tl) return tl def drawLegend(self, xl=0.8, yl=0.8, xh=0.95, yh=0.95, title="", columns=1): self.pad.cd() self.legend = TLegend(xl, yl, xh, yh, title) self.legend.SetNColumns(columns) self.legend.SetFillColor(0) for h, l, o in zip(self.items, self.labels, self.options): if l is None or l == "": continue if o == "": o = h.GetDrawOption() if o == "same": o = "L" if o == "hist" or o == "H": o = "L" elif o.lower() in self.markers: o = "P" self.legend.AddEntry(h, l, o) self.legend.Draw() self.Draw() def DrawLegend(self, xl=0.8, yl=0.8, xh=0.95, yh=0.95, title="", columns=1): self.drawLegend(xl, yl, xh, yh, title, columns) def Draw(self): ymax = self.ymax if self.ymax <= self.ymin: for h in self.items: try: y = h.GetBinContent(h.GetMaximumBin()) * 1.10 except AttributeError: try: y = max(h.GetY()) except AttributeError: y = ymax if y > ymax: ymax = y if ymax <= self.ymin: ymax = self.ymin + 1 self.hF.GetYaxis().SetLimits(self.ymin, ymax) if "logy" in self.frame_options: self.pad.SetLogy() if "logx" in self.frame_options: self.pad.SetLogx() if "grid" in self.frame_options or "gridxy" in self.frame_options: self.pad.SetGrid() if "gridx" in self.frame_options: self.pad.SetGridx() if "gridy" in self.frame_options: self.pad.SetGridy() self.pad.Update() def Print(self, filename=None): if filename: self.pad.Print(filename) def addVerticalAxis(self, title, ymin, ymax): self.yaxis2 = TGaxis(self.xmax, self.ymin, self.xmax, self.ymax, ymin, ymax, 510, "+L") self.yaxis2.SetName("axis") self.yaxis2.SetLabelOffset(0.01) self.yaxis2.SetTitle(title) self.yaxis2.Draw() self.ymin2 = ymin self.ymax2 = ymax return self.yaxis2 def getXaxis(self): """ return the x-axis (for setting attributes) """ return self.hF.GetXaxis() def getYaxis(self): """ return the y-axis (for setting attributes) """ return self.hF.GetYaxis() def getYaxis2(self): """ return the second y-axis (note: this is a TGaxis, not a TAxis)""" return self.axis2
def plot_dir(dlist, binfo, outfile, prefix, infix, color=[kBlack], legend=[], norm=False): c = TCanvas() c.SetGrid(0) bnarr = binfo.keys() bnarr.sort() bncode = reduce(lambda x, y: x + y, bnarr) namex = bnarr[-2] namey = bnarr[-1] xarr = binfo[namex] yarr = binfo[namey] c.GetPad(0).SetMargin(0.2, 0.15, 0.2, 0.15) c.Divide(len(xarr) - 1, len(yarr) - 1, 0.0, 0.0) #CanvasPartition(c,len(xarr)-1,len(yarr)-1) Ntot = 1 for k in range(len(bnarr) - 2): Ntot *= len(binfo[bnarr[k]]) - 1 for k in range(Ntot): indarr = [] fact = 1 for n in range(len(bnarr) - 2): edg = binfo[bnarr[n]] indarr.append((k / fact) % (len(edg) - 1)) fact *= len(binfo[bnarr[n]]) - 1 for j in range(len(yarr) - 1, 0, -1): YMAX = -100000000. YMIN = 100000000. row = j for i in range(len(xarr) - 1): col = i + 1 bcode = reduce(lambda x, y: str(x) + str(y), indarr) bcode = str(bcode) + str(col - 1) + str(row - 1) cnt = 0 for d in dlist: h = d.Get(prefix[cnt] + bncode + "_" + bcode) maxb = h.GetMaximumBin() minb = h.GetMinimumBin() ymaxtest = h.GetBinContent(maxb) ymintest = h.GetBinContent(minb) # if True: # print "before::" + d.GetPath() + "/" + h.GetName() + "::row::" + str(row) + ":: ymax / integral / bin / content(bin): " + str(ymaxtest) + " / " + str(h.Integral()) + " / " + str( maxb) + " / " + str( h.GetBinContent(maxb)) if norm and h.Integral() != 0: ymaxtest = ymaxtest / h.Integral() ymintest = ymintest / h.Integral() if norm and h.Integral() == 0: ymaxtest = 0 ymintest = 0 if ymaxtest > YMAX: YMAX = ymaxtest if ymintest < YMIN: YMIN = ymintest cnt += 1 # if norm: # print bncode + "::row::" + str(row) + "::ymax: " + str(YMAX) for i in range(len(xarr) - 1): col = i + 1 pad = c.cd(col + (len(yarr) - 1 - row) * (len(xarr) - 1)) pad.SetGridy(0) #pad.SetGridx(0) pad.SetTicks(0, 0) bcode = reduce(lambda x, y: str(x) + str(y), indarr) bcode = str(bcode) + str(col - 1) + str(row - 1) cnt = 0 leg = TLegend(0.77, 0.64, 0.95, 0.89) leg.SetTextSize(0.15) for d in dlist: h = d.Get(prefix[cnt] + bncode + "_" + bcode) h.GetYaxis().SetRangeUser(YMIN, YMAX * 1.1) if "ac" in prefix[cnt] or "Acceptance" in d.GetPath(): h.GetYaxis().SetRangeUser(0, 1) h.GetYaxis().SetLabelSize(0.15) h.GetXaxis().SetLabelSize(0.1) h.GetYaxis().SetLabelSize(0.15) xtitle = str( binfo[namex][col - 1]) + '<' + namex + '<' + str( binfo[namex][col]) + " " + prefix[cnt].split( '_')[0][1:] h.GetXaxis().SetTitle(xtitle) h.GetXaxis().CenterTitle() h.GetXaxis().SetTitleSize(0.12) ytitle = str( binfo[namey][row - 1]) + '<' + namey + '<' + str( binfo[namey][row]) h.GetXaxis().SetNdivisions(505) h.GetYaxis().SetNdivisions(505) h.SetLineWidth(3) h.SetLineColor(color[cnt]) h.SetMarkerColor(color[cnt]) if (legend != []): le = leg.AddEntry(h, legend[cnt], 'lp') le.SetTextColor(color[cnt]) if cnt == 0: if not norm: h.DrawCopy() elif (h.Integral() != 0): hnorm = h.DrawNormalized("e") hnorm.GetYaxis().SetRangeUser(YMIN, YMAX * 1.1) else: h.DrawCopy() if 'hM_' in h.GetName(): rhopeak = TLine(0.77, 0, 0.77, YMAX) rhopeak.SetLineStyle(kDashed) rhopeak.SetLineWidth(1) rhopeak.SetLineColor(kRed) rhopeak.DrawLine(0.77, 0, 0.77, YMAX) pad.Update() pt = pad.FindObject("title") pt.SetX1NDC(0.15) pt.SetTextSizePixels(12) if col == len(xarr) - 1: yaxis = TGaxis(pad.GetUxmax(), pad.GetUymin(), pad.GetUxmax(), pad.GetUymax(), 0, 1, 1, "U+-") yaxis.SetTitle(ytitle + " ") yaxis.SetTitleSize(0.12) yaxis.SetTitleOffset(-0.25) yaxis.SetName(h.GetName() + "_ax_y") yaxis.Draw("") pad.GetListOfPrimitives().Add(yaxis) pad.Update() else: if not norm: h.DrawCopy("same") elif (h.Integral() != 0): hnorm = h.DrawNormalized("samee") hnorm.GetYaxis().SetRangeUser(YMIN, YMAX * 1.1) else: h.DrawCopy("samee") cnt += 1 if (legend != []): leg.Draw() pad.GetListOfPrimitives().Add(leg) outfile.cd() bcode = reduce(lambda x, y: str(x) + str(y), indarr) cncode = reduce(lambda x, y: x + y, bnarr[-3]) c.Write("can_" + infix + "_" + str(cncode) + "_" + str(bcode), TObject.kOverwrite)