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)
