Merge pull request #24 from ttedeschi/update-palette

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) + "}"

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()