import io import sys import itertools as it from os.path import dirname, join, abspath, normpath from math import ceil, sqrt from subprocess import run 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', } def show_event(dataset, idx): ids = list(dataset.GenPart_pdgId[idx]) nrgs = list(dataset.GenPart_energy[idx]) links = list(dataset.GenPart_motherIndex[idx]) max_nrg = max(nrgs) nrgs_scaled = [nrg/max_nrg for nrg in nrgs] g = pdp.Dot() for i, id_ in enumerate(ids): color = ",".join(map(str, [nrgs_scaled[i], .7, .8])) g.add_node(pdp.Node(str(i), label=PDG[id_], style="filled", 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 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: single_plot_size = (600, 450) max_width = 1800 scale = .75 x_size = int(1600*scale) y_size = int(1200*scale) canvas = ROOT.TCanvas("c1", "", x_size, y_size) # @property # def canvas(self): # cls = self.__class__ # if not hasattr(cls, "_canvas"): # scale = .75 # x_size = int(1600*scale) # y_size = int(1200*scale) # cls._canvas = ROOT.TCanvas("c1", "", x_size, y_size) # return cls._canvas 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_plots = len([obj for obj in self.map.values() if hasattr(obj, "Draw") ]) if n_plots*self.single_plot_size[0] > self.max_width: shape_x = self.max_width//self.single_plot_size[0] shape_y = ceil(n_plots / shape_x) shape = (shape_x, shape_y) self.canvas.Clear() self.canvas.SetCanvasSize(shape[0]*self.single_plot_size[0], shape[1]*self.single_plot_size[1]) 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 @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: cls.canvas.Clear() cls.canvas.SetCanvasSize(cls.single_plot_size[0], cls.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") 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, "") # prevent stack from getting garbage collected _stacks[stack.GetName()] = stack if draw_canvas: cls.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 cls.canvas.SetCanvasSize(cls.single_plot_size[0]*shape[0], cls.single_plot_size[1]*shape[1]) 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, draw_canvas=False, **kwargs) cls.canvas.cd(n_hists).BuildLegend(0.75, 0.75, 0.95, 0.95, "")