Plotter.py

import copy
import importlib
import os
import ROOT
import sys
import yaml


def mk_smart_hist(hist, hist_desc):
    use_log_y = hist_desc.get("use_log_y", False)
    use_log_x = hist_desc.get("use_log_x", False)
    max_y_sf = hist_desc.get("max_y_sf", 1.5)
    divide_by_bin_width = hist_desc.get("divide_by_bin_width", False)
    smart_hist = ROOT.root_ext.SmartHistogram("TH1D")(
        hist, use_log_y, use_log_x, max_y_sf, divide_by_bin_width
    )
    if "x_title" in hist_desc:
        smart_hist.GetXaxis().SetTitle(hist_desc["x_title"])
    if "y_title" in hist_desc:
        smart_hist.GetYaxis().SetTitle(hist_desc["y_title"])
    return smart_hist


def to_str(value):
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, list):
        return ",".join(map(to_str, value))
    return str(value)


def ToItem(entry):
    map = ROOT.std.map("std::string", "std::string")()
    for key, value in entry.items():
        map[key] = to_str(value)
    item = ROOT.root_ext.draw_options.Item()
    item.properties = ROOT.analysis.PropertyList(map)
    return item


def LoadHistOptions(entry):
    item = ToItem(entry)
    return ROOT.root_ext.draw_options.Histogram(item)


def LoadPageOptions(entry):
    item = ToItem(entry)
    return ROOT.root_ext.draw_options.Page(item)


class Plotter(object):
    initialized = False

    def read_config(self, config):
        if type(config) == str:
            with open(config, "r") as f:
                return yaml.safe_load(f)
        elif type(config) == dict:
            return config
        elif type(config) == list:
            return config
        else:
            raise RuntimeError(f"Unknown config type: {type(config)}")

    def __init__(self, page_cfg, page_cfg_custom, hist_cfg):
        if not Plotter.initialized:
            ROOT.gROOT.SetBatch(True)
            python_dir = os.path.dirname(os.path.abspath(__file__))
            header_dir = os.path.join(python_dir, "include")
            ROOT.gInterpreter.Declare(f'#include "{header_dir}/PdfPrinter.h"')
            ROOT.gInterpreter.Declare(
                f'#include "{header_dir}/StackedPlotDescriptor.h"'
            )
            Plotter.initialized = True

        self.page_cfg = self.read_config(page_cfg)
        if page_cfg_custom:
            self.page_cfg.update(self.read_config(page_cfg_custom))
        self.hist_cfg = self.read_config(hist_cfg)

        self.sgn_hist_opt = LoadHistOptions(self.page_cfg["sgn_hist"])
        self.bkg_hist_opt = LoadHistOptions(self.page_cfg["bkg_hist"])
        self.data_hist_opt = LoadHistOptions(self.page_cfg["data_hist"])
        self.bkg_unc_hist_opt = LoadHistOptions(self.page_cfg["bkg_unc_hist"])
        self.page = LoadPageOptions(self.page_cfg["page_setup"])

    def plot(self, hist_name, histograms, output_file, want_data=True, custom=None, scale=None):
        page_cfg = copy.deepcopy(self.page_cfg)
        if custom:
            for key, value in custom.items():
                page_cfg[key]["text"] = value

        items = ROOT.root_ext.draw_options.ItemCollection()
        for element in page_cfg["page_setup"]["text_boxes"] + [
            page_cfg["page_setup"]["legend"]
        ]:
            items[element] = ToItem(page_cfg[element])
        printer = ROOT.root_ext.PdfPrinter(output_file, items, self.page)

        desc = ROOT.root_ext.StackedPlotDescriptor(
            self.page,
            self.sgn_hist_opt,
            self.bkg_hist_opt,
            self.data_hist_opt,
            self.bkg_unc_hist_opt,
        )
        smart_hists = {}
        # for input in self.inputs_cfg:
        # print(histograms.items())
        for hist_key, (
            hist,
            plot_name,
            plot_color,
            process_group,
        ) in histograms.items():
            smart_hists[hist.GetName()] = mk_smart_hist(hist, self.hist_cfg[hist_name])
            if process_group == "signals":
                signal_scale = 1.0
                if scale:
                    if type(scale) in [ int, float ]:
                        signal_scale = scale
                    elif type(scale) == dict and hist.GetName() in scale:
                        signal_scale = scale[hist.GetName()]
                desc.AddSignalHistogram(
                    smart_hists[hist.GetName()],
                    plot_name,
                    ROOT.root_ext.Color.Parse(plot_color),
                    signal_scale,
                )
            elif process_group == "backgrounds":
                desc.AddBackgroundHistogram(
                    smart_hists[hist.GetName()],
                    plot_name,
                    ROOT.root_ext.Color.Parse(plot_color),
                )
            elif process_group == "data":
                desc.AddDataHistogram(
                    smart_hists[hist.GetName()],
                    plot_name,
                )
            else:
                raise RuntimeError(f'Unknown process group: {process_group}')

            # hist_type = input.get('type', 'background')
            # if hist_type == 'signal':
            #   desc.AddSignalHistogram(smart_hists[name], input['title'], ROOT.root_ext.Color.Parse(input['color']),
            #                           input.get('scale', 1.))
            # elif hist_type == 'background':
            #   desc.AddBackgroundHistogram(smart_hists[name], input['title'], ROOT.root_ext.Color.Parse(input['color']))
            # elif hist_type == 'data':
            #   if want_data:
            #     desc.AddDataHistogram(smart_hists[name], input['title'])
            # else:
            #   raise RuntimeError(f'Unknown histogram type: {hist_type}')
        printer.Print(hist_name, desc, True)


def load(module_file):
    if not os.path.exists(module_file):
        raise RuntimeError(f"Cannot find path to {module_file}.")

    module_name, module_ext = os.path.splitext(module_file)
    spec = importlib.util.spec_from_file_location(module_name, module_file)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Plotting tool.")
    parser.add_argument(
        "--page-cfg", required=True, type=str, help="Page style config file"
    )
    parser.add_argument(
        "--page-cfg-custom",
        required=False,
        default=None,
        type=str,
        help="Additional page style customisations",
    )
    parser.add_argument(
        "--hist-cfg", required=True, type=str, help="Histogram style config file"
    )
    parser.add_argument(
        "--inputs-cfg", required=True, type=str, help="Inputs style config file"
    )
    parser.add_argument("--hist-name", required=True, type=str, help="Histogram name")
    parser.add_argument(
        "--custom",
        required=False,
        default=None,
        type=str,
        help="Customizations in format key1=value1,key2=value2,...",
    )
    parser.add_argument("--output", required=True, type=str, help="Output pdf file.")
    parser.add_argument(
        "--hist-maker", required=True, type=str, help="Path to histogram maker module."
    )
    parser.add_argument(
        "--verbose", required=False, type=int, default=0, help="verbosity level"
    )
    parser.add_argument(
        "hist_maker_args", type=str, nargs="*", help="hist maker arguments"
    )
    args = parser.parse_args()

    plotter = Plotter(
        page_cfg=args.page_cfg,
        page_cfg_custom=args.page_cfg_custom,
        hist_cfg=args.hist_cfg,
        inputs_cfg=args.inputs_cfg,
    )
    hist_maker = load(args.hist_maker)
    hists = hist_maker.make_histograms(
        *args.hist_maker_args, hist_name=args.hist_name, hist_cfg=plotter.hist_cfg
    )
    custom = (
        None
        if args.custom is None
        else dict(item.split("=") for item in args.custom.split(","))
    )
    plotter.plot(args.hist_name, hists, args.output, custom=custom)