utils.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325
  1. import io
  2. import os
  3. import sys
  4. import itertools as it
  5. from os.path import dirname, join, abspath, normpath, getctime
  6. from math import ceil, floor, sqrt
  7. from collections import deque
  8. from subprocess import run, PIPE
  9. from IPython.display import Image
  10. import pydotplus.graphviz as pdp
  11. import ROOT
  12. PRJ_PATH = normpath(join(dirname(abspath(__file__)), "../"))
  13. EXE_PATH = join(PRJ_PATH, "build/main")
  14. PDG = {1: 'd', -1: 'd̄',
  15. 2: 'u', -2: 'ū',
  16. 3: 's', -3: 's̄',
  17. 4: 'c', -4: 'c̄',
  18. 5: 'b', -5: 'b̄',
  19. 6: 't', -6: 't̄',
  20. 11: 'e-', -11: 'e+',
  21. 12: 'ν_e', -12: 'ῡ_e',
  22. 13: 'μ-', -13: 'μ+',
  23. 14: 'ν_μ', -14: 'ῡ_μ',
  24. 15: 'τ-', -15: 'τ+',
  25. 16: 'ν_τ', -16: 'ῡ_τ',
  26. 21: 'gluon',
  27. 22: 'γ',
  28. 23: 'Z0',
  29. 24: 'W+', -24: 'W-',
  30. 25: 'H',
  31. }
  32. SINGLE_PLOT_SIZE = (600, 450)
  33. MAX_WIDTH = 1800
  34. SCALE = .75
  35. CAN_SIZE_DEF = (int(1600*SCALE), int(1200*SCALE))
  36. CANVAS = ROOT.TCanvas("c1", "", *CAN_SIZE_DEF)
  37. ROOT.gStyle.SetPalette(112) # set the "virdidis" color map
  38. def clear():
  39. CANVAS.Clear()
  40. CANVAS.SetCanvasSize(*CAN_SIZE_DEF)
  41. def get_color(val, max_val, min_val = 0):
  42. val = (val-min_val)/(max_val-min_val)
  43. val = round(val * (ROOT.gStyle.GetNumberOfColors()-1))
  44. col_idx = ROOT.gStyle.GetColorPalette(val)
  45. col = ROOT.gROOT.GetColor(col_idx)
  46. r = floor(256*col.GetRed())
  47. g = floor(256*col.GetGreen())
  48. b = floor(256*col.GetBlue())
  49. gs = (r + g + b)//3
  50. text_color = 'white' if gs < 100 else 'black'
  51. return '#{:02x}{:02x}{:02x}'.format(r, g, b), text_color
  52. def show_event(dataset, idx):
  53. ids = list(dataset.GenPart_pdgId[idx])
  54. stats = list(dataset.GenPart_status[idx])
  55. energies = list(dataset.GenPart_energy[idx])
  56. links = list(dataset.GenPart_motherIndex[idx])
  57. max_energy = max(energies)
  58. g = pdp.Dot()
  59. for i, id_ in enumerate(ids):
  60. color, text_color = get_color(energies[i], max_energy)
  61. shape = "ellipse" if stats[i] in (1, 23) else "invhouse"
  62. label = "{}({})".format(PDG[id_], stats[i])
  63. # label = PDG[id_]+"({:03e})".format(energies[i])
  64. g.add_node(pdp.Node(str(i), label=label,
  65. style="filled",
  66. shape=shape,
  67. fontcolor=text_color,
  68. fillcolor=color))
  69. for i, mother in enumerate(links):
  70. if mother != -1:
  71. g.add_edge(pdp.Edge(str(mother), str(i)))
  72. return Image(g.create_gif())
  73. class OutputCapture:
  74. def __init__(self):
  75. self.my_stdout = io.StringIO()
  76. self.my_stderr = io.StringIO()
  77. def get_stdout(self):
  78. self.my_stdout.seek(0)
  79. return self.my_stdout.read()
  80. def get_stderr(self):
  81. self.my_stderr.seek(0)
  82. return self.my_stderr.read()
  83. def __enter__(self):
  84. self.stdout = sys.stdout
  85. self.stderr = sys.stderr
  86. sys.stdout = self.my_stdout
  87. sys.stderr = self.my_stderr
  88. def __exit__(self, *args):
  89. sys.stdout = self.stdout
  90. sys.stderr = self.stderr
  91. self.stdout = None
  92. self.stderr = None
  93. def normalize_columns(hist2d):
  94. normHist = ROOT.TH2D(hist2d)
  95. cols, rows = hist2d.GetNbinsX(), hist2d.GetNbinsY()
  96. for col in range(1, cols+1):
  97. sum_ = 0
  98. for row in range(1, rows+1):
  99. sum_ += hist2d.GetBinContent(col, row)
  100. if sum_ == 0:
  101. continue
  102. for row in range(1, rows+1):
  103. norm = hist2d.GetBinContent(col, row) / sum_
  104. normHist.SetBinContent(col, row, norm)
  105. return normHist
  106. class HistCollection:
  107. def __init__(self, sample_name, input_filename):
  108. self.sample_name = sample_name
  109. self.input_filename = input_filename
  110. self.output_filename = self.input_filename.replace(".root", "_result.root")
  111. self.conditional_recompute()
  112. self.load_objects()
  113. # Now add these histograms into the current ROOT directory (in memory)
  114. # and remove old versions if needed
  115. for obj in self.map.values():
  116. try:
  117. old_obj = ROOT.gDirectory.Get(obj.GetName())
  118. ROOT.gDirectory.Remove(old_obj)
  119. ROOT.gDirectory.Add(obj)
  120. except AttributeError:
  121. pass
  122. HistCollection.add_collection(self)
  123. def conditional_recompute(self):
  124. def recompute():
  125. print("Running analysis for sample: ", self.sample_name)
  126. if run([EXE_PATH, "-s", "-f", self.input_filename]).returncode != 0:
  127. raise RuntimeError(("Failed running analysis code."
  128. " See log file for more information"))
  129. if run(["make"], cwd=join(PRJ_PATH, "build"), stdout=PIPE, stderr=PIPE).returncode != 0:
  130. raise RuntimeError("Failed recompiling analysis code")
  131. if (not os.path.isfile(self.output_filename) or (getctime(EXE_PATH) > getctime(self.output_filename))):
  132. recompute()
  133. else:
  134. print("Loading unchanged result file ", self.output_filename)
  135. def load_objects(self):
  136. file = ROOT.TFile.Open(self.output_filename)
  137. l = file.GetListOfKeys()
  138. self.map = {}
  139. for i in range(l.GetSize()):
  140. name = l.At(i).GetName()
  141. new_name = ":".join((self.sample_name, name))
  142. obj = file.Get(name)
  143. try:
  144. obj.SetName(new_name)
  145. obj.SetDirectory(0) # disconnects Object from file
  146. except AttributeError:
  147. pass
  148. self.map[name] = obj
  149. setattr(self, name, obj)
  150. file.Close()
  151. @classmethod
  152. def calc_shape(cls, n_plots):
  153. if n_plots*SINGLE_PLOT_SIZE[0] > MAX_WIDTH:
  154. shape_x = MAX_WIDTH//SINGLE_PLOT_SIZE[0]
  155. shape_y = ceil(n_plots / shape_x)
  156. return (shape_x, shape_y)
  157. else:
  158. return (n_plots, 1)
  159. def draw(self, shape=None):
  160. objs = [obj for obj in self.map.values() if hasattr(obj, "Draw")]
  161. if shape is None:
  162. n_plots = len(objs)
  163. shape = self.calc_shape(n_plots)
  164. CANVAS.Clear()
  165. CANVAS.SetCanvasSize(shape[0]*SINGLE_PLOT_SIZE[0], shape[1]*SINGLE_PLOT_SIZE[1])
  166. CANVAS.Divide(*shape)
  167. i = 1
  168. for hist in objs:
  169. CANVAS.cd(i)
  170. try:
  171. hist.SetStats(False)
  172. except AttributeError:
  173. pass
  174. hist.Draw(self.get_draw_option(hist))
  175. i += 1
  176. CANVAS.Draw()
  177. @staticmethod
  178. def get_draw_option(obj):
  179. obj_type = type(obj)
  180. if obj_type in (ROOT.TH1F, ROOT.TH1I, ROOT.TH1D):
  181. return ""
  182. elif obj_type in (ROOT.TH2F, ROOT.TH2I, ROOT.TH2D):
  183. return "COLZ"
  184. elif obj_type in (ROOT.TGraph,):
  185. return "A*"
  186. else:
  187. return None
  188. @classmethod
  189. def get_hist_set(cls, attrname):
  190. labels, hists = zip(*[(sample_name, getattr(h, attrname))
  191. for sample_name, h in cls.collections.items()])
  192. return labels, hists
  193. @classmethod
  194. def add_collection(cls, hc):
  195. if not hasattr(cls, "collections"):
  196. cls.collections = {}
  197. cls.collections[hc.sample_name] = hc
  198. @classmethod
  199. def stack_hist(cls,
  200. hist_name,
  201. title="",
  202. enable_fill=False,
  203. normalize_to=0,
  204. draw=False,
  205. draw_canvas=True,
  206. draw_option="",
  207. make_legend=False,
  208. _stacks={}):
  209. labels, hists = cls.get_hist_set(hist_name)
  210. if draw_canvas:
  211. CANVAS.Clear()
  212. CANVAS.SetCanvasSize(SINGLE_PLOT_SIZE[0],
  213. SINGLE_PLOT_SIZE[1])
  214. colors = it.cycle([ROOT.kRed, ROOT.kBlue, ROOT.kGreen])
  215. stack = ROOT.THStack(hist_name+"_stack", title)
  216. if labels is None:
  217. labels = [hist.GetName() for hist in hists]
  218. if type(normalize_to) in (int, float):
  219. normalize_to = [normalize_to]*len(hists)
  220. ens = enumerate(zip(hists, labels, colors, normalize_to))
  221. for i, (hist, label, color, norm) in ens:
  222. hist_copy = hist
  223. hist_copy = hist.Clone(hist.GetName()+"_clone" + draw_option)
  224. hist_copy.SetTitle(label)
  225. if enable_fill:
  226. hist_copy.SetFillColorAlpha(color, 0.75)
  227. hist_copy.SetLineColorAlpha(color, 0.75)
  228. if norm:
  229. integral = hist_copy.Integral()
  230. hist_copy.Scale(norm/integral, "nosw2")
  231. hist_copy.SetStats(True)
  232. stack.Add(hist_copy)
  233. if draw:
  234. stack.Draw(draw_option)
  235. if make_legend:
  236. CANVAS.BuildLegend(0.75, 0.75, 0.95, 0.95, "")
  237. # prevent stack from getting garbage collected
  238. _stacks[stack.GetName()] = stack
  239. if draw_canvas:
  240. CANVAS.Draw()
  241. return stack
  242. @classmethod
  243. def stack_hist_array(cls,
  244. hist_names,
  245. titles,
  246. shape=None, **kwargs):
  247. n_hist = len(hist_names)
  248. if shape is None:
  249. if n_hist <= 4:
  250. shape = (1, n_hist)
  251. else:
  252. shape = (ceil(sqrt(n_hist)),)*2
  253. CANVAS.SetCanvasSize(SINGLE_PLOT_SIZE[0]*shape[0],
  254. SINGLE_PLOT_SIZE[1]*shape[1])
  255. CANVAS.Divide(*shape)
  256. for i, hist_name, title in zip(range(1, n_hist+1), hist_names, titles):
  257. CANVAS.cd(i)
  258. cls.stack_hist(hist_name, title=title, draw=True,
  259. draw_canvas=False, **kwargs)
  260. CANVAS.cd(n_hist).BuildLegend(0.75, 0.75, 0.95, 0.95, "")
  261. pts = deque([], 50)
  262. @classmethod
  263. def hist_array_single(cls,
  264. hist_name,
  265. title=None,
  266. **kwargs):
  267. n_hist = len(cls.collections)
  268. shape = cls.calc_shape(n_hist)
  269. CANVAS.SetCanvasSize(SINGLE_PLOT_SIZE[0]*shape[0],
  270. SINGLE_PLOT_SIZE[1]*shape[1])
  271. CANVAS.Divide(*shape)
  272. labels, hists = cls.get_hist_set(hist_name)
  273. def pave_loc():
  274. hist.Get
  275. for i, label, hist in zip(range(1, n_hist+1), labels, hists):
  276. CANVAS.cd(i)
  277. hist.SetStats(False)
  278. hist.Draw(cls.get_draw_option(hist))
  279. pt = ROOT.TPaveText(0.70, 0.87, 0.85, 0.95, "NDC")
  280. pt.AddText("Dataset: "+label)
  281. pt.Draw()
  282. cls.pts.append(pt)