TTTT/python/utils.py

import io
import os
import sys
import itertools as it
from os.path import dirname, join, abspath, normpath, getctime
from math import ceil, floor, sqrt
from subprocess import run, PIPE
from IPython.display import Image

import pydotplus.graphviz as pdp
import ROOT

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

PDG = {1:   'd',   -1:  'd̄',
       2:   'u',   -2:  'ū',
       3:   's',   -3:  's̄',
       4:   'c',   -4:  'c̄',
       5:   'b',   -5:  'b̄',
       6:   't',   -6:  't̄',

       11:  'e-',  -11: 'e+',
       12:  'ν_e', -12: 'ῡ_e',

       13:  'μ-',  -13: 'μ+',
       14:  'ν_μ', -14: 'ῡ_μ',

       15:  'τ-',  -15: 'τ+',
       16:  'ν_τ', -16: 'ῡ_τ',

       21:  'gluon',
       22:  'γ',
       23:  'Z0',
       24:  'W+',  -24: 'W-',
       25:  'H',
       }

SINGLE_PLOT_SIZE = (600, 450)
MAX_WIDTH = 1800

SCALE = .75
CAN_SIZE_DEF = (int(1600*SCALE), int(1200*SCALE))
CANVAS = ROOT.TCanvas("c1", "", *CAN_SIZE_DEF)
ROOT.gStyle.SetPalette(112)  # set the "virdidis" color map

def clear():
    CANVAS.Clear()
    CANVAS.SetCanvasSize(*CAN_SIZE_DEF)

def get_color(val, max_val, min_val = 0):
    val = (val-min_val)/(max_val-min_val)
    val = round(val * (ROOT.gStyle.GetNumberOfColors()-1))
    col_idx = ROOT.gStyle.GetColorPalette(val)
    col = ROOT.gROOT.GetColor(col_idx)
    r = floor(256*col.GetRed())
    g = floor(256*col.GetGreen())
    b = floor(256*col.GetBlue())
    gs = (r + g + b)//3
    text_color = 'white' if gs < 100 else 'black'
    return '#{:02x}{:02x}{:02x}'.format(r, g, b), text_color


def show_event(dataset, idx):
    ids = list(dataset.GenPart_pdgId[idx])
    stats = list(dataset.GenPart_status[idx])
    energies = list(dataset.GenPart_energy[idx])
    links = list(dataset.GenPart_motherIndex[idx])
    max_energy = max(energies)
    g = pdp.Dot()
    for i, id_ in enumerate(ids):
        color, text_color = get_color(energies[i], max_energy)
        shape = "ellipse" if stats[i] in (1, 23) else "invhouse"
        label = "{}({})".format(PDG[id_], stats[i])
        # label = PDG[id_]+"({:03e})".format(energies[i])
        g.add_node(pdp.Node(str(i), label=label,
                   style="filled",
                   shape=shape,
                   fontcolor=text_color,
                   fillcolor=color))
    for i, mother in enumerate(links):
        if mother != -1:
            g.add_edge(pdp.Edge(str(mother), str(i)))
    return Image(g.create_gif())


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 normalize_columns(hist2d):
    normHist = ROOT.TH2D(hist2d)
    cols, rows = hist2d.GetNbinsX(), hist2d.GetNbinsY()
    for col in range(1, cols+1):
        sum_ = 0
        for row in range(1, rows+1):
            sum_ += hist2d.GetBinContent(col, row)
        if sum_ == 0:
            continue
        for row in range(1, rows+1):
            norm = hist2d.GetBinContent(col, row) / sum_
            normHist.SetBinContent(col, row, norm)
    return normHist


class HistCollection:

    def __init__(self, sample_name, input_filename):
        self.sample_name = sample_name
        self.input_filename = input_filename
        self.output_filename = self.input_filename.replace(".root", "_result.root")
        self.conditional_recompute()
        self.load_objects()

        # 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 conditional_recompute(self):

        def recompute():
            print("Running analysis for sample: ", self.sample_name)
            if run([EXE_PATH, "-s", "-f", self.input_filename]).returncode != 0:
                raise RuntimeError(("Failed running analysis code."
                                    " See log file for more information"))

        if run(["make"], cwd=join(PRJ_PATH, "build"), stdout=PIPE, stderr=PIPE).returncode != 0:
            raise RuntimeError("Failed recompiling analysis code")
        if (not os.path.isfile(self.output_filename) or (getctime(EXE_PATH) > getctime(self.output_filename))):
            recompute()
        else:
            print("Loading unchanged result file ", self.output_filename)

    def load_objects(self):
        file = ROOT.TFile.Open(self.output_filename)
        l = file.GetListOfKeys()
        self.map = {}
        for i in range(l.GetSize()):
            name = l.At(i).GetName()
            new_name = ":".join((self.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()

    @classmethod
    def calc_shape(cls, n_plots):
        if n_plots*SINGLE_PLOT_SIZE[0] > MAX_WIDTH:
            shape_x = MAX_WIDTH//SINGLE_PLOT_SIZE[0]
            shape_y = ceil(n_plots / shape_x)
            return (shape_x, shape_y)
        else:
            return (n_plots, 1)


    def draw(self, shape=None):
        objs = [obj for obj in self.map.values() if hasattr(obj, "Draw")]
        if shape is None:
            n_plots = len(objs)
            shape = self.calc_shape(n_plots)
        CANVAS.Clear()
        CANVAS.SetCanvasSize(shape[0]*SINGLE_PLOT_SIZE[0], shape[1]*SINGLE_PLOT_SIZE[1])
        CANVAS.Divide(*shape)
        i = 1
        for hist in objs:
            CANVAS.cd(i)
            try:
                hist.SetStats(False)
            except AttributeError:
                pass
            hist.Draw(self.get_draw_option(hist))
            i += 1
        CANVAS.Draw()

    @staticmethod
    def get_draw_option(obj):
        obj_type = type(obj)
        if obj_type in (ROOT.TH1F, ROOT.TH1I, ROOT.TH1D):
            return ""
        elif obj_type in (ROOT.TH2F, ROOT.TH2I, ROOT.TH2D):
            return "COLZ"
        elif obj_type in (ROOT.TGraph,):
            return "A*"
        else:
            return None

    @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

    @classmethod
    def stack_hist(cls,
                   hist_name,
                   title="",
                   enable_fill=False,
                   normalize_to=0,
                   draw=False,
                   draw_canvas=True,
                   draw_option="",
                   make_legend=False,
                   _stacks={}):
        labels, hists = cls.get_hist_set(hist_name)
        if draw_canvas:
            CANVAS.Clear()
            CANVAS.SetCanvasSize(SINGLE_PLOT_SIZE[0],
                                 SINGLE_PLOT_SIZE[1])

        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" + draw_option)
            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(True)
            stack.Add(hist_copy)
        if draw:
            stack.Draw(draw_option)
            if make_legend:
                CANVAS.BuildLegend(0.75, 0.75, 0.95, 0.95, "")
        # prevent stack from getting garbage collected
        _stacks[stack.GetName()] = stack
        if draw_canvas:
            CANVAS.Draw()
        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
        CANVAS.SetCanvasSize(SINGLE_PLOT_SIZE[0]*shape[0],
                             SINGLE_PLOT_SIZE[1]*shape[1])
        CANVAS.Divide(*shape)
        for i, hist_name, title in zip(range(1, n_hists+1), hist_names, titles):
            CANVAS.cd(i)
            cls.stack_hist(hist_name, title=title, draw=True,
                           draw_canvas=False, **kwargs)
        CANVAS.cd(n_hists).BuildLegend(0.75, 0.75, 0.95, 0.95, "")

    pts = []
    @classmethod
    def hist_array_single(cls,
                          hist_name,
                          title=None,
                          **kwargs):
        n_hists = len(cls.collections)
        shape = cls.calc_shape(n_hists)
        CANVAS.SetCanvasSize(SINGLE_PLOT_SIZE[0]*shape[0],
                             SINGLE_PLOT_SIZE[1]*shape[1])
        CANVAS.Divide(*shape)
        labels, hists = cls.get_hist_set(hist_name)
        for i, label, hist in zip(range(1, n_hists+1), labels, hists):
            pt = ROOT.TPaveText(300, 3, 400, 3.5)
            CANVAS.cd(i)
            hist.SetStats(False)
            hist.Draw(cls.get_draw_option(hist))
            pt.AddText("Dataset: "+label)
            pt.Draw()
            cls.pts.append(pt)