diff --git a/src/cmsstyle/cmsstyle.py b/src/cmsstyle/cmsstyle.py index 80f73ff..c52363e 100644 --- a/src/cmsstyle/cmsstyle.py +++ b/src/cmsstyle/cmsstyle.py @@ -63,6 +63,11 @@ def AppendAdditionalInfo(text): kSquare = True kRectangular = False +# Petroff color schemes for 6, 8 and 10 colors, respectively +petroff_6 = ["#5790fc", "#f89c20", "#e42536", "#964a8b", "#9c9ca1", "#7a21dd"] +petroff_8 = ["#1845fb", "#ff5e02", "#c91f16", "#c849a9", "#adad7d", "#86c8dd", "#578dff", "#656364"] +petroff_10 = ["#3f90da", "#ffa90e", "#bd1f01", "#94a4a2", "#832db6", "#a96b59", "#e76300", "#b9ac70", "#717581", "#92dadd"] + # Define an alternative color palette and a function to set it MyPalette = None @@ -98,6 +103,11 @@ def SetAlternative2DColor(hist=None, style=None, alpha=1): hist.SetContour(len(MyPalette)) +def SetCMSPalette(): + """Allow to directly set the official 2D CMS palette""" + cmsStyle.SetPalette(rt.kViridis) + #cmsStyle.SetPalette(rt.kCividis) + def GetPalette(hist): """Allow to retrieve palette option. Must update the pad to access the palette""" UpdatePad() @@ -666,6 +676,49 @@ def cmsDrawLine(line, lcolor=rt.kRed, lstyle=rt.kSolid, lwidth=2): line.SetLineWidth(lwidth) line.Draw("SAME") +import re + +def is_valid_hex_color(hex_color): + hex_color_pattern = re.compile(r'^#(?:[0-9a-fA-F]{3}){1,2}$') + + return bool(hex_color_pattern.match(hex_color)) + + +def cmsDrawStack(stack, legend, MC, data = None, palette = None): + """Draws stacked histograms and data on a pre-defined stackplot and fills a pre-defined legend, using a user-defined or default list (palette) of hex colors""" + + is_user_palette_valid = False + + if palette != None: + is_user_palette_valid = all(is_valid_hex_color(color) for color in palette) + if is_user_palette_valid: + palette_ = palette + if len(MC.keys()) > len(palette_): + print("Length of provided palette is smaller than the number of histograms to be drawn, wrap around is enabled") + else: + print("Invalid palette elements provided, default palette will be used") + + if palette == None or is_user_palette_valid == False: + if len(MC.keys()) < 7: + palette_ = petroff_6 + elif len(MC.keys()) < 9: + palette_ = petroff_8 + else: + palette_ = petroff_10 + if len(MC.keys()) > len(palette_): + print("Length of largest default palette is smaller than the number of histograms to be drawn, wrap around is enabled") + + for n, item in enumerate(MC.items()): + item[1].SetLineColor(rt.TColor.GetColor(palette_[n%len(palette_)])) + item[1].SetFillColor(rt.TColor.GetColor(palette_[n%len(palette_)])) + stack.Add(item[1]) + legend.AddEntry(item[1], item[0], "f") + stack.Draw("HIST SAME") + + if data != None: + cmsDraw(data, "P", mcolor=rt.kBlack) + legend.AddEntry(data, "Data", "lp") + def ScaleText(name, scale=0.75): return "#scale[" + str(scale) + "]{" + str(name) + "}" diff --git a/tests/example_palette.py b/tests/example_palette.py new file mode 100644 index 0000000..0634dc6 --- /dev/null +++ b/tests/example_palette.py @@ -0,0 +1,101 @@ +import os, ROOT +import numpy as np +import cmsstyle as CMS + +CMS.SetExtraText("Simulation") + +class Plotter: + def __init__(self): + self.outputPath = "./pdfs_palette" + os.makedirs(self.outputPath, exist_ok=True) + self.CreateHistograms() + + def CreateHistograms(self): + self.bkgs = [] + for i in range(0,6): + h = ROOT.TH1F("bkg{}".format(i), "bkg{}".format(i), 100, 0, 10) + for _ in range(16666): + h.Fill(np.random.normal(5, 1)) + self.bkgs.append(h) + + self.hist2d = ROOT.TH2F("hist2d", "2D Histogram", 25, 0, 5, 25, 0, 5) + for i in range(200000): + x = np.random.normal(2.5, 1) + y = np.random.normal(2.5, 1) + self.hist2d.Fill(x, y) + self.hist2d.Scale(10.0 / self.hist2d.Integral()) + + def Plot(self, square, iPos): + canv_name = f'example_{"square" if square else "rectangle"}_pos{iPos}' + CMS.SetLumi("") + CMS.SetEnergy("13") + CMS.ResetAdditionalInfo() + + stack = ROOT.THStack("stack", "Stacked") + + canv = CMS.cmsCanvas( + canv_name, + 0, + 10, + 1e-3, + 4300, + "X", + "", + square=square, + extraSpace=0.01, + iPos=iPos, + ) + leg = CMS.cmsLeg(0.81, 0.89 - 0.05 * 7, 0.99, 0.89, textSize=0.04) + + # Put samples together and draw them stacked + hist_dict = {} + names = ["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K"] + for n, hist in enumerate(self.bkgs): + hist_dict[names[n]] = hist + CMS.cmsDrawStack(stack, leg, hist_dict) + + # Takes care of fixing overlay and closing object + CMS.SaveCanvas(canv, os.path.join(self.outputPath, canv_name + ".pdf")) + + def Plot2D(self, square, iPos): + canv_name = f'example_2D_{"square" if square else "rectangle"}_pos{iPos}' + # Allow to reduce the size of the lumi info + scaleLumi = 0.80 if square else None + canv = CMS.cmsCanvas( + canv_name, + 0, + 5, + 0, + 5, + "X", + "Y", + square=square, + extraSpace=0.01, + iPos=iPos, + with_z_axis=True, + scaleLumi=scaleLumi, + ) + + self.hist2d.GetZaxis().SetTitle("Events normalised") + self.hist2d.GetZaxis().SetTitleOffset(1.4 if square else 1.2) + self.hist2d.Draw("same colz") + + # Set the CMS official palette + CMS.SetCMSPalette() + + # Allow to adjust palette position + CMS.UpdatePalettePosition(self.hist2d, canv) + + CMS.SaveCanvas(canv, os.path.join(self.outputPath, canv_name + ".pdf")) + +def main(): + plotter = Plotter() + plotter.Plot(square=CMS.kSquare, iPos=0) + plotter.Plot(square=CMS.kRectangular, iPos=0) + plotter.Plot(square=CMS.kSquare, iPos=11) + plotter.Plot(square=CMS.kRectangular, iPos=11) + plotter.Plot2D(square=CMS.kSquare, iPos=0) + plotter.Plot2D(square=CMS.kRectangular, iPos=0) + +if __name__ == "__main__": + main() diff --git a/tests/pdfs_palette/example_2D_rectangle_pos0.pdf b/tests/pdfs_palette/example_2D_rectangle_pos0.pdf new file mode 100644 index 0000000..f471bd0 Binary files /dev/null and b/tests/pdfs_palette/example_2D_rectangle_pos0.pdf differ diff --git a/tests/pdfs_palette/example_2D_square_pos0.pdf b/tests/pdfs_palette/example_2D_square_pos0.pdf new file mode 100644 index 0000000..aea0f4e Binary files /dev/null and b/tests/pdfs_palette/example_2D_square_pos0.pdf differ diff --git a/tests/pdfs_palette/example_rectangle_pos0.pdf b/tests/pdfs_palette/example_rectangle_pos0.pdf new file mode 100644 index 0000000..436959e Binary files /dev/null and b/tests/pdfs_palette/example_rectangle_pos0.pdf differ diff --git a/tests/pdfs_palette/example_rectangle_pos11.pdf b/tests/pdfs_palette/example_rectangle_pos11.pdf new file mode 100644 index 0000000..98f4d45 Binary files /dev/null and b/tests/pdfs_palette/example_rectangle_pos11.pdf differ diff --git a/tests/pdfs_palette/example_square_pos0.pdf b/tests/pdfs_palette/example_square_pos0.pdf new file mode 100644 index 0000000..038cf77 Binary files /dev/null and b/tests/pdfs_palette/example_square_pos0.pdf differ diff --git a/tests/pdfs_palette/example_square_pos11.pdf b/tests/pdfs_palette/example_square_pos11.pdf new file mode 100644 index 0000000..c7d2bbc Binary files /dev/null and b/tests/pdfs_palette/example_square_pos11.pdf differ