import math
import ROOT
from .commonHelpers.logger import logger
logger = logger.getChild(__name__)
from . import pyrootHelpers as PH
from . import globalStyle as gst
from .labels import scaleYPosTopMargin
[docs]class Legend(ROOT.TLegend):
"""
:param scaleLegend: Scale Legend by this factor
:param scaleLegendX: Scale Legend by this factor in x-direction
:param scaleLegendY: Scale Legend by this factor in x-direction
:param drawstring: Drawstring for
`ROOT TLegend <https://root.cern.ch/doc/master/classTLegend.html>`_
:param nColumns: Number of columns
:param xOffset: Shift legend in x direction
:param yOffset: Shift legend in y direction
:param forceDynamicFontSize: Always use dynamic font size (scale text to fit legend)
:param noDynamicFontSize: Never use dynamic font size (by default used for legends
containing long titles)
:param addEventCount: Add the total event counts to the legend
:param eventCountCutflow: Use the content of the last bin instead of the integral
when showing total event counts
:param eventCountGen: In addition to the integral, also show the number of raw events
"""
def __init__(self,
scaleLegend=1,
scaleLegendX=None,
scaleLegendY=None,
drawstring='',
nColumns=1,
xOffset=0,
yOffset=0,
forceDynamicFontSize=False,
noDynamicFontSize=False,
addEventCount=False,
eventCountCutflow=False,
eventCountGen=False):
self.scaleLegend = scaleLegend
if scaleLegendX is None:
self.scaleLegendX = scaleLegend
else:
self.scaleLegendX = scaleLegendX
if scaleLegendY is None:
self.scaleLegendY = scaleLegend
else:
self.scaleLegendY = scaleLegendY
self.drawstring = drawstring
self.nColumns = nColumns
self.xOffset = xOffset
self.yOffset = yOffset
self.forceDynamicFontSize = forceDynamicFontSize
self.noDynamicFontSize = noDynamicFontSize
if self.forceDynamicFontSize and self.noDynamicFontSize:
raise ValueError("Can't use forceDynamicFontSize and noDynamicFontSize at once")
self.objects = []
self.zIndex = 2
self.eventCount = addEventCount
self.cutflow = eventCountCutflow
self.genEvents = eventCountGen
[docs] def addEntry(self, obj, option):
logger.debug('adding {} with label {} and option {}'.format(obj,obj.legend, option))
try:
hide = obj.hide
except AttributeError:
hide = False
try:
drawLegend = obj.drawLegend
except AttributeError:
drawLegend = True
if not hide and drawLegend:
legend = obj.legend
if legend is None:
legend = ''
self.objects.append((obj, legend, option))
[docs] def addEventCount(self, label, obj):
if self.cutflow:
label += gst.legendEventCountFormat.format(obj.GetBinContent(obj.GetNbinsX()))
else:
label += gst.legendEventCountFormat.format(obj.Integral())
if self.genEvents and not self.cutflow:
label += gst.legendEventCountFormatRaw.format(obj.GetEntries())
return label
[docs] def hasLongTitleEntry(self):
for obj, label, option in self.objects:
if not label:
label = obj.GetTitle()
if self.eventCount:
label = self.addEventCount(label, obj)
if len(label) > gst.legendLongTitleThreshold:
return True
return False
[docs] def draw(self):
dynamicFontSize = False
margin = 0.4
colwidth = 0.15/ROOT.gPad.GetAbsWNDC()
lineheight = 0.055/ROOT.gPad.GetAbsHNDC()
if self.hasLongTitleEntry():
dynamicFontSize = True
margin = 0.3
if not self.forceDynamicFontSize:
if not self.nColumns == 1:
logger.warning("Legend has a long title - setting nColumns to 1")
self.nColumns = 1
colwidth = 0.2/ROOT.gPad.GetAbsWNDC()
n_entries = len(self.objects)
n_columns = self.nColumns
n_lines = math.ceil(float(n_entries)/float(n_columns))
x_size = n_columns * colwidth * self.scaleLegendX
y_size = n_lines * lineheight * self.scaleLegendY
# create actual legend here, hope that's enough otherwise we need to know the size beforehand
yMax = scaleYPosTopMargin(gst.legendYMax)
yMin = yMax-y_size
if yMin < 0.1:
logger.debug("yMin to small ({}) - Scaling legend".format(yMin))
yMin = 0.1
dynamicFontSize = True
xMin = gst.legendXMin-(n_columns-1)*colwidth
xMax = xMin+x_size
xMin = xMin+self.xOffset
xMax = xMax+self.xOffset
yMin = yMin+self.yOffset
yMax = yMax+self.yOffset
if self.forceDynamicFontSize:
dynamicFontSize = True
if self.noDynamicFontSize:
dynamicFontSize = False
super(Legend, self).__init__(xMin, yMin, xMax, yMax)
self.SetNColumns(n_columns)
self.SetTextFont(gst.legendFont)
# interestingly, this mainly determines the width of the boxes
self.SetMargin(margin)
self.SetBorderSize(gst.legendBorderSize)
if gst.legendTextSize is not None and not dynamicFontSize:
self.SetTextSize(gst.legendTextSize/ROOT.gPad.GetAbsHNDC()*self.scaleLegend)
self.SetFillColor(0)
self.SetFillColorAlpha(0, 0)
for obj, label, option in self.objects:
if self.eventCount:
label = self.addEventCount(label, obj)
self.AddEntry(obj, label, option)
self.Draw(self.drawstring)
return self