diff --git a/src/cmsstyle/cmsstyle.py b/src/cmsstyle/cmsstyle.py index a0bf7b4..47974b5 100644 --- a/src/cmsstyle/cmsstyle.py +++ b/src/cmsstyle/cmsstyle.py @@ -17,12 +17,12 @@ import re from contextlib import contextmanager -from dataclasses import dataclass -from typing import Any, Iterable import os from ._utils import env_type -# This global variables for the module should not be accessed directy! Use the utilities below. +# This global variables for the module should not +# be accessed directy! Use the utilities below. + cms_lumi = "Run 2, 138 fb^{#minus1}" cms_energy = "13 TeV" @@ -98,7 +98,7 @@ def SetLumi(lumi, unit="fb", run="Run 2", round_lumi=-1): lumi (float): The integrated luminosity value. May be skipped if set to None. unit (str, optional): The integrated luminosity unit. Defaults to "fb". run (str, optional): The LHC run to which the sample refers to. - round_lumi (int, optional): Number of decimal digits to present the number. If no 0, 1 nor 2, no rounding is done. + round_lumi (int, optional): Number of decimal digits to present the number. If no 0,1 nor 2, no rounding is done """ global cms_lumi @@ -395,12 +395,14 @@ def getPettroffColor(color): # -> EColor # We try to identify a ROOT color... try: # Some versions don't identify GetColorByName as a valid method (still used in CMSSW) val = rt.TColor.GetColorByName(color) - if (val>=0): return rt.TColor.GetColorByName(color) - except: # We keep for others some basic/common color names + if (val >= 0): + return rt.TColor.GetColorByName(color) + except Exception: # We keep for others some basic/common color names pass return None # Not valid color! + # # # # def getPettroffColorSet(ncolors): """This method returns a list of colors for the given number of colors based on @@ -806,7 +808,7 @@ def CMS_lumi(pad, iPosX=11, scaleLumi=1): posY_ = 1 - t - relPosY * (1 - t - b) - if outOfFrame: # CMS logo and extra text out of the frame + if outOfFrame: # CMS logo and extra text out of the frame if ( len(useCmsLogo) > 0 ): # Using CMS Logo instead of the text label (uncommon and discouraged!) @@ -1018,8 +1020,10 @@ def cmsCanvas( else: y_offset = yTitOffset - if (y_offset<1.5): L += y_offset*50-60 # Some adjustment - elif (y_offset<1.8): L += (y_offset-1.4)*35+25 + if (y_offset < 1.5): + L += y_offset*50-60 # Some adjustment + elif (y_offset < 1.8): + L += (y_offset - 1.4)*35+25 canv = rt.TCanvas(canvName, canvName, 50, 50, W, H) canv.SetFillColor(0) @@ -1696,15 +1700,16 @@ def setRootObjectProperties(obj, **kwargs): getattr(obj, method)(*xval) else: try: - getattr(obj,method)(xval) + getattr(obj, method)(xval) except TypeError: if 'Color' in xkey: # The string may be just a color indicated as a name - getattr(obj,method)(getPettroffColor(xval)) + getattr(obj, method)(getPettroffColor(xval)) else: raise + # # # # -def copyRootObjectProperties (obj,srcobj,proplist,**kwargs): +def copyRootObjectProperties(obj, srcobj, proplist, **kwargs): """This method allows to copy the properties of a ROOT object from a reference source (another ROOT object) using a list of named keyword arguments to call the associated methods. @@ -1730,12 +1735,13 @@ def copyRootObjectProperties (obj,srcobj,proplist,**kwargs): """ for xprp in proplist: # Just proceding with the copy! - getattr(obj,'Set'+xprp)(getattr(srcobj,'Get'+xprp)()) + getattr(obj, 'Set'+xprp)(getattr(srcobj, 'Get'+xprp)()) # If we indicated some additional arguments, we use them to further # configure the object. - if len(kwargs)>0: - setRootObjectProperties(obj,**kwargs) + if len(kwargs) > 0: + setRootObjectProperties(obj, **kwargs) + # # # # def is_valid_hex_color(hexcolor): @@ -1784,7 +1790,7 @@ def cmsReturnMaxY(*args): maxval = 0 for xobj in args: - if (xobj.Class().GetName()=='TEfficiency'): # For efficiencies, we just put 1.2 as maximum! + if (xobj.Class().GetName() == 'TEfficiency'): # For efficiencies, we just put 1.2 as maximum! maxval = 1.19 elif ( xobj.Class().GetName() == "THStack" @@ -1892,6 +1898,7 @@ def plot(self, obj, opt="", **kwargs): self._pad.cd() setRootObjectProperties(obj, **kwargs) obj.Draw(opt) + self._pad.RedrawAxis() self._drawables.append(obj) @@ -1911,19 +1918,11 @@ def __init__(self, ncolumns, nrows, pad_horizontal_margin, pad_vertical_margin): self.pad_vertical_margin = pad_vertical_margin -class LegendItem(object): - """An item to be added to a legend, together with its name and drawing option.""" - def __init__(self, obj, name, opt): - self.obj = obj - self.name = name - self.opt = opt - - class CMSCanvasManager(object): """A manager of the different graphical parts of a canvas.""" def __init__(self, canvas, pads=None, frames=None, - bottom_pad=None, top_pad=None, grid_metadata=None): + bottom_pad=None, top_pad=None, legendtextsize=None, cmslogotextsize=None, ipos=None, grid_metadata=None): """ At minimum, a canvas manager needs a canvas to plot on. Optionally, it can manage different sub-components of a canvas: @@ -1937,6 +1936,9 @@ def __init__(self, canvas, pads=None, frames=None, """ self._canvas = canvas self._frames = frames + self._legendtextsize = legendtextsize + self._cmslogotextsize = cmslogotextsize + self._ipos = ipos if self._frames is not None: if pads is None: @@ -1996,6 +1998,7 @@ def plot_common_legend(self, pad, *args, **kwargs): subtitleFont = kwargs.get("subtitleFont", 52) textalign = kwargs.get("textalign", 13) ipos = kwargs.get("ipos", 0) + legendtextSize = kwargs.get("legendtextSize", 30) pad._pad.cd() horizontal_margin = float(self._grid_metadata.pad_horizontal_margin) / self._grid_metadata.ncolumns @@ -2020,11 +2023,12 @@ def plot_common_legend(self, pad, *args, **kwargs): if n % ncolumns == 0: leg.AddEntry(0, " ", " ") n += 1 - leg.AddEntry(arg.obj, arg.name, arg.opt) + # leg.AddEntry(arg.obj, arg.name, arg.opt) + leg.AddEntry(arg[0], arg[1], arg[2]) n += 1 else: for arg in args: - leg.AddEntry(arg.obj, arg.name, arg.opt) + leg.AddEntry(arg[0], arg[1], arg[2]) pad.plot(leg) @@ -2033,9 +2037,7 @@ def plot_common_legend(self, pad, *args, **kwargs): latex.SetTextFont(titleFont) canvas_height = pad._pad.GetWh() - ymin = pad._pad.GetYlowNDC() - ymax = pad._pad.GetYlowNDC() + pad._pad.GetHNDC() - pad_ndc_height = ymax - ymin + pad_ndc_height = pad._pad.GetHNDC() pad_pixel_height = canvas_height * pad_ndc_height titleSize = titleSize / pad_pixel_height subtitleSize = titleSize * 0.76 @@ -2043,8 +2045,10 @@ def plot_common_legend(self, pad, *args, **kwargs): latex.SetTextSize(titleSize) latex.SetTextAlign(13) + leg.SetTextSize(legendtextSize / pad_pixel_height) + if ipos != 0: - latex.DrawLatex(0.11, 0.60, title) + latex.DrawLatex(0.105, 0.60, title) else: latex.DrawLatex(0.10, 0.97, title) @@ -2052,7 +2056,7 @@ def plot_common_legend(self, pad, *args, **kwargs): latex.SetTextSize(subtitleSize) if ipos != 0: - latex.DrawLatex(0.11, 0.30, subtitle) + latex.DrawLatex(0.105, 0.30, subtitle) else: latex.DrawLatex(0.17, 0.94, subtitle) @@ -2083,9 +2087,7 @@ def plot_text( latex.SetTextAlign(textalign) canvas_height = pad._pad.GetWh() - ymin = pad._pad.GetYlowNDC() - ymax = pad._pad.GetYlowNDC() + pad._pad.GetHNDC() - pad_ndc_height = ymax - ymin + pad_ndc_height = pad._pad.GetHNDC() pad_pixel_height = canvas_height * pad_ndc_height textsize = textsize / pad_pixel_height latex.SetTextSize(textsize) @@ -2112,7 +2114,7 @@ def ylimits(self, limits=None): for nframe in limits: self._frames[nframe].SetMinimum(limits[nframe][0]) self._frames[nframe].SetMaximum(limits[nframe][1]) - #self._frames[nframe].GetYaxis().SetLimits(limits[nframe][0], limits[nframe][1]) + # self._frames[nframe].GetYaxis().SetLimits(limits[nframe][0], limits[nframe][1]) def xlimits(self, limits=None): for nframe in limits: @@ -2193,7 +2195,7 @@ def _subplots_coordinates( 1 - canvas_top_margin - canvas_bottom_margin ) pad_w = width_ratio / sum(width_ratios) - # We skip the first pad as its coordinates are computed already before the for loop + # We skip the first pad as its coordinates are computed already before the for loop if npad != 0 and npad % ncolumns == 0: # This branch is for the start of a new row ylow -= pad_h @@ -2218,6 +2220,7 @@ def _subplots_coordinates( return coordinates, top_pad_coords, bottom_pad_coords + def subplots( ncolumns, nrows, @@ -2230,25 +2233,31 @@ def subplots( canvas_width=2000, canvas_height=2000, axis_title_size=50, - axis_label_size=50 * 0.8 + axis_label_size=50 * 0.8, + logotextsize=50, + legendtextsize=30, + ipos=0 ): """ Creates multiple pads in a canvas according to the input configuration, then returns an object to help manage the canvas and all its graphical parts. Args: - - ncolumns: number of columns in the grid - - nrows: number of rows in the grid - - height_ratios: list of weights for the relative heights of the pads along the columns. Length must be equal to nrows - - width_ratios: list of weights for the relative widths of the pads along the rows. Length must be equal to ncolumns - - canvas_top_margin: margin to remove starting from the top of the canvas to make space for the top pad - - canvas_bottom_margin: margin to remove starting from the bottom of the canvas to make space for the bottom pad - - shared_x_axis: whether the x axis of all columns should be shared - - shared_y_axis: whether the y axis of all columns should be shared - - canvas_width: total width of the canvas - - canvas_height: total height of the canvas - - axis_title_size: reference absolute size for axis titles - - axis_label_size: reference absolute size for axis labels + ncolumns (int): number of columns in the grid + nrows (int): number of rows in the grid + height_ratios (list, optional): list of weights for the relative heights of the pads along the columns. Length must be equal to nrows + width_ratios (list, optional): list of weights for the relative widths of the pads along the rows. Length must be equal to ncolumns + canvas_top_margin (float, optional): margin to remove starting from the top of the canvas to make space for the top pad + canvas_bottom_margin (float, optional): margin to remove starting from the bottom of the canvas to make space for the bottom pad + shared_x_axis (bool, optional) : whether the x axis of all columns should be shared + shared_y_axis (bool, optional): whether the y axis of all columns should be shared + canvas_width (float, optional): total width of the canvas + canvas_height (float, optional): total height of the canvas + axis_title_size (float, optional): reference absolute size for axis titles + axis_label_size (float, optional): reference absolute size for axis labels + logotextsize (float, optional): absolute text size of experiment logo + legendtextsize (float, optional): absolute text size of legend + ipos (int): position of experiment logo """ top_pad = None @@ -2359,9 +2368,7 @@ def subplots( with _managed_tpad_context(canvas): pad.cd() canvas_height = pad.GetWh() - ymin = pad.GetYlowNDC() - ymax = pad.GetYlowNDC() + pad.GetHNDC() - pad_ndc_height = ymax - ymin + pad_ndc_height = pad.GetHNDC() pad_pixel_height = canvas_height * pad_ndc_height labeltextsize = axis_label_size / pad_pixel_height frame.GetXaxis().SetLabelSize(labeltextsize) @@ -2374,18 +2381,14 @@ def subplots( with _managed_tpad_context(canvas): pad.cd() canvas_height = pad.GetWh() - ymin = pad.GetYlowNDC() - ymax = pad.GetYlowNDC() + pad.GetHNDC() - pad_ndc_height = ymax - ymin + pad_ndc_height = pad.GetHNDC() pad_pixel_height = canvas_height * pad_ndc_height labeltextsize = axis_label_size / pad_pixel_height frame.GetYaxis().SetLabelSize(labeltextsize) frame.GetYaxis().SetNdivisions(3, 5, 0, True) titletextsize = axis_title_size / pad_pixel_height frame.GetYaxis().SetTitleSize(titletextsize) - frame.GetYaxis().SetTitleOffset( - 3 * (height_ratios[i // ncolumns] / float(sum(height_ratios))) - ) + frame.GetYaxis().SetTitleOffset(0) return CMSCanvasManager( canvas, @@ -2393,10 +2396,117 @@ def subplots( frames=listofframes, bottom_pad=bottom_pad, top_pad=top_pad, + cmslogotextsize=logotextsize, + legendtextsize=legendtextsize, + ipos=ipos, grid_metadata=GridMetaData( ncolumns, nrows, pad_horizontal_margin, pad_vertical_margin ), ) +def cmsMultiCanvas( + canvName, + nColumns, + nRows, + Xlimits, + Ylimits, + nameXaxis, + nameYaxis, + labelTextSize=50*0.8, + titleTextSize=50, + lumiTextSize=50, + logoTextSize=50 * 0.75 / 0.6, + legendTextSize=30, + heightRatios=None, + widthRatios=None, + canvasTopMargin=0.1, + canvasBottomMargin=0.03, + canvasWidth=2000, + canvasHeight=2000, + iPos=0 +): + """ + Create a plot with multiple pads arranged in a grid, with shared axes, common legend and CMS styling. + + Args: + canvName (str): The name of the canvas. + nColumns (int): number of columns in the grid + nRows (int): number of rows in the grid + XLimits (dict): dictionary containing X axis limits for all plots, with integer keys corresponding to left-to-right top-to-bottom numbering scheme, starting from 0, and a 2-element iterable value, containing x_min and x_max. e.g. {0: [0, 50000], 1: [1, 50000]}. + YLimits (dict): dictionary containing Y axis limits for all plots, with integer keys corresponding to left-to-right top-to-bottom numbering scheme, starting from 0, and a 2-element iterable value, containing y_min and y_max. e.g. {0: [0, 50000], 1: [1, 50000]}. + nameXaxis (str): the label for the x-axis. + nameYaxis (dict): the label for the y-axis, in the form of a dict int:str with integer keys corresponding to left-to-right top-to-bottom numbering scheme, starting from 0. e.g. {0: "Y Label for 1st plot", 4: "Y Label for 5th plot"} + labelTextSize (float, optional): absolute value of textSize of axis labels (same for X and Y). Defaults to 50*0.8. + titleTextSize (float, optional): absolute value for textSize of axis titles (same for X and Y). Defaults to 50. + lumiTextSize (float, optional): absolute value for textSize of lumi text. Defaults to 50. + logoTextSize (float, optional): absolute value for textSize of CMS Logo, extraText is scaled accordingly. Defaults to 50 * 0.75 / 0.6. + legendTextSize (float, optional): Absolute value for text size of legend. Defaults to 30. + heightRatios (list, optional): list of weights for the relative heights of the pads along the columns. Length must be equal to nrows. Defaults to None, which means each plot gets the same height. + widthRatios (list, optional): list of weights for the relative widths of the pads along the rows. Length must be equal to ncolumns. Defaults to None, which means each plot gets the same width. + canvasTopMargin (float, optional): margin to remove starting from the top of the canvas to make space for the top pad. Defaults to 0.1. + canvasBottomMargin (float, optional): margin to remove starting from the bottom of the canvas to make space for the bottom pad. Defaults to 0.03. + canvasWidth (float, optional): total width of the canvas. Defaults to 2000. + canvasHeight (float, optional): total height of the canvas. Defaults to 2000. + iPos (int, optional): The position of the CMS text. 0 (outside of legend box) or 11 (top-left left-aligned inside legend box). Defaults to 11. + + Returns: + CMSCanvasManager: A object to handle the multipad grid. + """ + cvm = subplots( + ncolumns=nColumns, + nrows=nRows, + height_ratios=heightRatios, + width_ratios=widthRatios, + canvas_top_margin=canvasTopMargin, + canvas_bottom_margin=canvasBottomMargin, + axis_label_size=labelTextSize, + axis_title_size=titleTextSize, + canvas_height=canvasHeight, + canvas_width=canvasWidth, + logotextsize=logoTextSize, + legendtextsize=legendTextSize, + ipos=iPos + ) + + cvm.plot_text( + cvm.top_pad, + cms_lumi, + textsize=lumiTextSize + ) + + cvm.plot_text( + cvm.bottom_pad, + nameXaxis, + textsize=titleTextSize, + ) + + cvm.ylabel(labels=nameYaxis) + cvm.ylimits(limits=Ylimits) + cvm.xlimits(limits=Xlimits) + + return cvm + + +def cmsMultiCanvasLeg(cvm, *legend_items): + """ + Create, fill and draw the common legend of a cmsMultiCanvas. + + Args: + cvm (CMSCanvasManager): A CMSCanvasManager object returned by cmsMultiCanvas() method. + *legend_items (list): list of tuples in the form [(ROOTobject1, name for ROOTobject1 as a string, print option for ROOTObject1 as a string), (ROOTobject2, name for ROOTobject2 as a string, print option for ROOTObject2 as a string), ...] + + Returns: + None + """ + cvm.plot_common_legend( + cvm.top_pad, + *legend_items, + textalign=12, + legendtextSize=cvm._legendtextsize, + title=cmsText, + subtitle=extraText, + titleSize=cvm._cmslogotextsize, + ipos=cvm._ipos + ) # ####################################################################### diff --git a/tests/test_subplots.py b/tests/test_subplots.py index 1157224..32a7577 100644 --- a/tests/test_subplots.py +++ b/tests/test_subplots.py @@ -57,46 +57,66 @@ def _create_drawables(): def test_subplots(): """Example of multiple plots in the same canvas, with shared common legend""" + cmsstyle.setCMSStyle() + + # set the luminosity, the COM energy, the Run period to show in the canvases + cmsstyle.SetLumi(34.8, run='2022') + cmsstyle.SetEnergy(13.6) + # default extra text is "Preliminary", set it to an empty string to remove it + cmsstyle.SetExtraText('') + ncolumns = 2 nrows = 6 - cvm = cmsstyle.subplots( - ncolumns=ncolumns, - nrows=nrows, - height_ratios=[2, 1] * (nrows // 2), - canvas_top_margin=0.1, - canvas_bottom_margin=0.03, - axis_label_size = 40 + cvm = cmsstyle.cmsMultiCanvas( + canvName="", + nColumns=ncolumns, + nRows=nrows, + heightRatios=[3, 1] * (nrows // 2), + Xlimits={ + 0: [-2, 2], 1: [-2, 2], + 2: [-2, 2], 3: [-2, 2], + 4: [-2, 2], 5: [-2, 2], + 6: [-2, 2], 7: [-2, 2], + 8: [-2, 2], 9: [-2, 2], + 10: [-2, 2], 11: [-2, 2] + }, + Ylimits={ + 0: [0, 500], 1: [1, 500], + 2: [0.5, 1.5], 3: [0.5, 1.5], + 4: [0, 500], 5: [0, 500], + 6: [0.5, 1.5], 7: [0.5, 1.5], + 8: [0, 500], 9: [0, 500], + 10: [0.5, 1.5], 11: [0.5, 1.5] + }, + nameXaxis="m^{ll} (GeV)", + nameYaxis={0:"Test0", 2:"", 4:"Test4", 6:"", 8:"Test8", 10:""}, + labelTextSize=30, + titleTextSize=40, + lumiTextSize=50, + logoTextSize=50 * 0.75 / 0.6, + legendTextSize=30, + canvasTopMargin=0.1, + canvasBottomMargin=0.03, + canvasHeight=2000, + iPos=11 ) data, hs, h_err, ratio, yerr_root, ref_line, bkg, signal = _create_drawables() - cvm.plot_common_legend( - cvm.top_pad, - cmsstyle.LegendItem(data, "Uncertainty", "pe"), - cmsstyle.LegendItem(bkg, "MC1", "f"), - cmsstyle.LegendItem(signal, "MC2", "f"), - cmsstyle.LegendItem(ratio, "Ratio", "pe"), - cmsstyle.LegendItem(ratio, "Ratio", "pe"), - cmsstyle.LegendItem(signal, "Testing", "f"), - cmsstyle.LegendItem(data, "Data", "pe"), - cmsstyle.LegendItem(bkg, "MC1", "f"), - cmsstyle.LegendItem(signal, "MC2", "f"), - cmsstyle.LegendItem(data, "Hello", "pe"), - cmsstyle.LegendItem(ratio, "BigTitle", "pe"), - textalign=12, - ipos = 11 + _ = cmsstyle.cmsMultiCanvasLeg(cvm, + (data, "Uncertainty", "pe"), + (bkg, "MC1", "f"), + (signal, "MC2", "f"), + (ratio, "Ratio", "pe"), + (ratio, "Ratio", "pe"), + (signal, "Testing", "f"), + (data, "Data", "pe"), + (bkg, "MC1", "f"), + (signal, "MC2", "f"), + (data, "Hello", "pe"), + (ratio, "BigTitle", "pe"), ) - cvm.plot_text( - cvm.top_pad, - "Run 2, 138 fb^{#minus1}", - ) - cvm.plot_text( - cvm.bottom_pad, - "m^{ll} (GeV)", - textsize=50, - ) - cvm.ylabel(labels={0:"Test0", 2:"", 4:"Test4", 6:"", 8:"Test8", 10:""}) row_index = -1 for i, pad in enumerate(cvm.pads): @@ -104,7 +124,7 @@ def test_subplots(): row_index += 1 if row_index % 2 == 0: pad.plot(hs) - pad.plot(h_err, "E2SAME0", FillColor=ROOT.kBlack, LineWidth=1, LineColor=355, FillStyle=3004) + pad.plot(h_err, "E2SAME0", FillColor=ROOT.kBlack, MarkerSize=0, LineWidth=1, LineColor=355, FillStyle=3004) pad.plot(data, "E1X0", FillColor=kBlue) else: pad.plot(yerr_root, "E2SAME0", LineWidth=100, MarkerSize=0, FillColor=ROOT.kBlack, FillStyle=3004) @@ -115,4 +135,4 @@ def test_subplots(): if __name__ == "__main__": - test_subplots() \ No newline at end of file + test_subplots()