TTTT/python/utils.py

import io
from os.path import dirname, join, abspath, normpath
import sys
from math import ceil, sqrt
from subprocess import run
import itertools as it
import ROOT

PRJ_PATH = normpath(join(dirname(abspath(__file__)), "../"))
EXE_PATH = join(PRJ_PATH, "build/main")


class OutputCapture:
    def __init__(self):
        self.my_stdout = io.StringIO()
        self.my_stderr = io.StringIO()

    def get_stdout(self):
        self.my_stdout.seek(0)
        return self.my_stdout.read()

    def get_stderr(self):
        self.my_stderr.seek(0)
        return self.my_stderr.read()

    def __enter__(self):
        self.stdout = sys.stdout
        self.stderr = sys.stderr
        sys.stdout = self.my_stdout
        sys.stderr = self.my_stderr

    def __exit__(self, *args):
        sys.stdout = self.stdout
        sys.stderr = self.stderr
        self.stdout = None
        self.stderr = None


def bin_range(n, end=None):
    if end is None:
        return range(1, n+1)
    else:
        return range(n+1, end+1)


def normalize_columns(hist2d):
    normHist = ROOT.TH2D(hist2d)
    cols, rows = hist2d.GetNbinsX(), hist2d.GetNbinsY()
    for col in bin_range(cols):
        sum_ = 0
        for row in bin_range(rows):
            sum_ += hist2d.GetBinContent(col, row)
        if sum_ == 0:
            continue
        for row in bin_range(rows):
            norm = hist2d.GetBinContent(col, row) / sum_
            normHist.SetBinContent(col, row, norm)
    return normHist


class HistCollection:
    def __init__(self, sample_name, input_filename,
                 rebuild_hists=False):
        self.sample_name = sample_name
        if rebuild_hists:
            run([EXE_PATH, "-s", "-f", input_filename])
        output_filename = input_filename.replace(".root", "_result.root")
        file = ROOT.TFile.Open(output_filename)
        l = file.GetListOfKeys()
        self.map = {}
        for i in range(l.GetSize()):
            name = l.At(i).GetName()
            new_name = ":".join((sample_name, name))
            obj = file.Get(name)
            try:
                obj.SetName(new_name)
                obj.SetDirectory(0)  # disconnects Object from file
            except AttributeError:
                pass
            self.map[name] = obj
            setattr(self, name, obj)
        file.Close()
        # Now add these histograms into the current ROOT directory (in memory)
        # and remove old versions if needed
        for obj in self.map.values():
            try:
                old_obj = ROOT.gDirectory.Get(obj.GetName())
                ROOT.gDirectory.Remove(old_obj)
                ROOT.gDirectory.Add(obj)
            except AttributeError:
                pass
        HistCollection.add_collection(self)

    def draw(self, shape=None):
        if shape is None:
            n = int(ceil(sqrt(len(self.map))))
            shape = (n, n)
        self.canvas.Clear()
        self.canvas.Divide(*shape)
        i = 1
        for hist in self.map.values():
            self.canvas.cd(i)
            try:
                hist.SetStats(False)
            except AttributeError:
                pass
            draw_option = ""
            if type(hist) in (ROOT.TH1F, ROOT.TH1I, ROOT.TH1D):
                draw_option = ""
            elif type(hist) in (ROOT.TH2F, ROOT.TH2I, ROOT.TH2D):
                draw_option = "COLZ"
            elif type(hist) in (ROOT.TGraph,):
                draw_option = "A*"
            else:
                print("cannot draw object", hist)
                continue  # Not a drawable type(probably)
            hist.Draw(draw_option)
            i += 1
        self.canvas.Draw()

    @classmethod
    def get_hist_set(cls, attrname):
        labels, hists = zip(*[(sample_name, getattr(h, attrname))
                              for sample_name, h in cls.collections.items()])
        return labels, hists

    @classmethod
    def add_collection(cls, hc):
        if not hasattr(cls, "collections"):
            cls.collections = {}
        cls.collections[hc.sample_name] = hc
        print("collection added: " + hc.sample_name)
        print("collections present: " + ', '.join(list(hc.collections.keys())))

    @property
    def canvas(self):
        cls = self.__class__
        if not hasattr(cls, "_canvas"):
            cls._canvas = ROOT.TCanvas("c1", "", 1600, 1200)
        return cls._canvas

    @canvas.setter
    def canvas(self, canvas):
        cls = self.__class__
        cls._canvas = canvas

    @classmethod
    def stack_hist(cls,
                   hist_name,
                   title="", enable_fill=False,
                   normalize_to=0, draw=False,
                   draw_option="",
                   make_legend=False,
                   _stacks={}):
        labels, hists = cls.get_hist_set(hist_name)

        colors = it.cycle([ROOT.kRed, ROOT.kBlue, ROOT.kGreen])
        stack = ROOT.THStack(hist_name+"_stack", title)
        if labels is None:
            labels = [hist.GetName() for hist in hists]
        if type(normalize_to) in (int, float):
            normalize_to = [normalize_to]*len(hists)
        ens = enumerate(zip(hists, labels, colors, normalize_to))
        for i, (hist, label, color, norm) in ens:
            hist_copy = hist
            hist_copy = hist.Clone(hist.GetName()+"_clone")
            hist_copy.SetTitle(label)
            if enable_fill:
                hist_copy.SetFillColorAlpha(color, 0.75)
                hist_copy.SetLineColorAlpha(color, 0.75)
            if norm:
                integral = hist_copy.Integral()
                hist_copy.Scale(norm/integral, "nosw2")
                hist_copy.SetStats(False)
            stack.Add(hist_copy)
        if draw:
            stack.Draw(draw_option)
            if make_legend:
                cls._canvas.BuildLegend(0.75, 0.75, 0.95, 0.95, "")
            # cls._canvas.Draw()
        # prevent stack from getting garbage collected
        _stacks[stack.GetName()] = stack
        return stack

    @classmethod
    def stack_hist_array(cls,
                         hist_names,
                         titles,
                         shape=None, **kwargs):
        n_hists = len(hist_names)
        if shape is None:
            if n_hists <= 4:
                shape = (1, n_hists)
            else:
                shape = (ceil(sqrt(n_hists)),)*2
        cls._canvas.Divide(*shape)
        for i, hist_name, title in zip(bin_range(n_hists), hist_names, titles):
            cls._canvas.cd(i)
            hists, labels = cls.get_hist_set(hist_name)
            cls.stack_hist(hist_name, title=title, draw=True, **kwargs)
        cls._canvas.cd(1).BuildLegend(0.75, 0.75, 0.95, 0.95, "")