plotter.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339
  1. #!/usr/bin/env python3
  2. from collections import defaultdict
  3. from io import BytesIO
  4. from base64 import b64encode
  5. import numpy as np
  6. import matplotlib.pyplot as plt
  7. from markdown import Markdown
  8. import latexipy as lp
  9. from filval.histogram_utils import (hist, hist2d, hist_bin_centers, hist_fit,
  10. hist_normalize, hist_stats)
  11. __all__ = ['Plot',
  12. 'decl_plot',
  13. 'grid_plot',
  14. 'render_plots',
  15. 'generate_dashboard',
  16. 'hist_plot',
  17. 'hist_plot_stack',
  18. 'hist2d_plot']
  19. class Plot:
  20. def __init__(self, subplots, name, title=None, docs="N/A", arg_dicts=None):
  21. self.subplots = subplots
  22. self.name = name
  23. self.title = title
  24. self.docs = docs
  25. self.arg_dicts = arg_dicts if arg_dicts is not None else {}
  26. MD = Markdown(extensions=['mdx_math'],
  27. extension_configs={'mdx_math': {'enable_dollar_delimiter': True}})
  28. lp.latexify(params={'pgf.texsystem': 'pdflatex',
  29. 'text.usetex': True,
  30. 'font.family': 'serif',
  31. 'pgf.preamble': [],
  32. 'font.size': 15,
  33. 'axes.labelsize': 15,
  34. 'axes.titlesize': 13,
  35. 'legend.fontsize': 13,
  36. 'xtick.labelsize': 11,
  37. 'ytick.labelsize': 11,
  38. 'figure.dpi': 150,
  39. 'savefig.transparent': False,
  40. },
  41. new_backend='TkAgg')
  42. def _fn_call_to_dict(fn, *args, **kwargs):
  43. from inspect import signature
  44. pnames = list(signature(fn).parameters)
  45. pvals = list(args) + list(kwargs.values())
  46. return {k: v for k, v in zip(pnames, pvals)}
  47. def _process_docs(fn):
  48. from inspect import getdoc
  49. raw = getdoc(fn)
  50. if raw:
  51. return MD.convert(raw)
  52. else:
  53. return None
  54. def decl_plot(fn):
  55. from functools import wraps
  56. @wraps(fn)
  57. def f(*args, **kwargs):
  58. fn(*args, **kwargs)
  59. argdict = _fn_call_to_dict(fn, *args, **kwargs)
  60. docs = _process_docs(fn)
  61. return argdict, docs
  62. return f
  63. def generate_dashboard(plots, title, output='dashboard.html', template='dashboard.j2', source_file=None):
  64. from jinja2 import Environment, PackageLoader, select_autoescape
  65. from os.path import join, isdir
  66. from os import mkdir
  67. from urllib.parse import quote
  68. env = Environment(
  69. loader=PackageLoader('filval', 'templates'),
  70. autoescape=select_autoescape(['htm', 'html', 'xml']),
  71. )
  72. env.globals.update({'quote': quote,
  73. 'enumerate': enumerate,
  74. 'zip': zip,
  75. })
  76. def get_by_n(objects, n=2):
  77. objects = list(objects)
  78. while objects:
  79. yield objects[:n]
  80. objects = objects[n:]
  81. if source_file is not None:
  82. with open(source_file, 'r') as this_file:
  83. source = this_file.read()
  84. else:
  85. source = "# Not supplied!!"
  86. if not isdir('output'):
  87. mkdir('output')
  88. with open(join('output', output), 'w') as tempout:
  89. templ = env.get_template(template)
  90. tempout.write(templ.render(
  91. plots=get_by_n(plots, 3),
  92. title=title,
  93. source=source,
  94. outdir="figures/"
  95. ))
  96. def _add_stats(hist, title=''):
  97. fmt = r'''\begin{{eqnarray*}}
  98. \sum{{x_i}} &=& {sum:5.3f} \\
  99. \sum{{\Delta x_i \cdot x_i}} &=& {int:5.3G} \\
  100. \mu &=& {mean:5.3G} \\
  101. \sigma^2 &=& {var:5.3G} \\
  102. \sigma &=& {std:5.3G}
  103. \end{{eqnarray*}}'''
  104. txt = fmt.format(**hist_stats(hist), title=title)
  105. txt = txt.replace('\n', ' ')
  106. plt.text(0.7, 0.9, txt,
  107. bbox={'facecolor': 'white',
  108. 'alpha': 0.7,
  109. 'boxstyle': 'square,pad=0.8'},
  110. transform=plt.gca().transAxes,
  111. verticalalignment='top',
  112. horizontalalignment='left',
  113. size='small')
  114. if title:
  115. plt.text(0.72, 0.97, title,
  116. bbox={'facecolor': 'white',
  117. 'alpha': 0.8},
  118. transform=plt.gca().transAxes,
  119. verticalalignment='top',
  120. horizontalalignment='left')
  121. def grid_plot(subplots):
  122. if any(len(row) != len(subplots[0]) for row in subplots):
  123. raise ValueError("make_plot requires a rectangular list-of-lists as "
  124. "input. Fill empty slots with None")
  125. def calc_rowspan(fig, row, col):
  126. span = 1
  127. for r in range(row + 1, len(fig)):
  128. if fig[r][col] == "FU":
  129. span += 1
  130. else:
  131. break
  132. return span
  133. def calc_colspan(fig, row, col):
  134. span = 1
  135. for c in range(col + 1, len(fig[row])):
  136. if fig[row][c] == "FL":
  137. span += 1
  138. else:
  139. break
  140. return span
  141. rows = len(subplots)
  142. cols = len(subplots[0])
  143. argdicts = defaultdict(list)
  144. docs = defaultdict(list)
  145. for i in range(rows):
  146. for j in range(cols):
  147. cell = subplots[i][j]
  148. if cell in ("FL", "FU", None):
  149. continue
  150. if not isinstance(cell, list):
  151. cell = [cell]
  152. colspan = calc_colspan(subplots, i, j)
  153. rowspan = calc_rowspan(subplots, i, j)
  154. plt.subplot2grid((rows, cols), (i, j),
  155. colspan=colspan, rowspan=rowspan)
  156. for plot in cell:
  157. plot_fn, args, kwargs = plot
  158. this_args, this_docs = plot_fn(*args, **kwargs)
  159. argdicts[(i, j)].append(this_args)
  160. docs[(i, j)].append(this_docs)
  161. return argdicts, docs
  162. def render_plots(plots, exts=('png',), scale=1.0, to_disk=True):
  163. for plot in plots:
  164. print(f'Building plot {plot.name}')
  165. plot.data = None
  166. if to_disk:
  167. with lp.figure(plot.name.replace(' ', '_'), directory='output/figures',
  168. exts=exts,
  169. size=(scale * 10, scale * 10)):
  170. argdicts, docs = grid_plot(plot.subplots)
  171. else:
  172. out = BytesIO()
  173. with lp.mem_figure(out,
  174. ext=exts[0],
  175. size=(scale * 10, scale * 10)):
  176. argdicts, docs = grid_plot(plot.subplots)
  177. out.seek(0)
  178. plot.data = b64encode(out.read()).decode()
  179. plot.argdicts = argdicts
  180. plot.docs = docs
  181. def add_decorations(axes, luminosity, energy):
  182. cms_prelim = r'{\raggedright{}\textsf{\textbf{CMS}}\\ \emph{Preliminary}}'
  183. axes.text(0.01, 0.98, cms_prelim,
  184. horizontalalignment='left',
  185. verticalalignment='top',
  186. transform=axes.transAxes)
  187. lumi = ""
  188. energy_str = ""
  189. if luminosity is not None:
  190. lumi = r'${} \mathrm{{fb}}^{{-1}}$'.format(luminosity)
  191. if energy is not None:
  192. energy_str = r'({} TeV)'.format(energy)
  193. axes.text(1, 1, ' '.join([lumi, energy_str]),
  194. horizontalalignment='right',
  195. verticalalignment='bottom',
  196. transform=axes.transAxes)
  197. def hist_plot(h, *args, norm=None, include_errors=False,
  198. log=False, xlim=None, ylim=None, fit=None,
  199. grid=False, stats=False, **kwargs):
  200. """ Plots a 1D ROOT histogram object using matplotlib """
  201. from inspect import signature
  202. if norm:
  203. h = hist_normalize(h, norm)
  204. values, errors, edges = h
  205. scale = 1. if norm is None else norm / np.sum(values)
  206. values = [val * scale for val in values]
  207. errors = [val * scale for val in errors]
  208. left, right = np.array(edges[:-1]), np.array(edges[1:])
  209. x = np.array([left, right]).T.flatten()
  210. y = np.array([values, values]).T.flatten()
  211. ax = plt.gca()
  212. ax.set_xlabel(kwargs.pop('xlabel', ''))
  213. ax.set_ylabel(kwargs.pop('ylabel', ''))
  214. title = kwargs.pop('title', '')
  215. if xlim is not None:
  216. ax.set_xlim(xlim)
  217. if ylim is not None:
  218. ax.set_ylim(ylim)
  219. # elif not log:
  220. # axes.set_ylim((0, None))
  221. ax.plot(x, y, *args, linewidth=1, **kwargs)
  222. if include_errors:
  223. ax.errorbar(hist_bin_centers(h), values, yerr=errors,
  224. color='k', marker=None, linestyle='None',
  225. barsabove=True, elinewidth=.7, capsize=1)
  226. if log:
  227. ax.set_yscale('log')
  228. if fit:
  229. f, p0 = fit
  230. popt, pcov = hist_fit(h, f, p0)
  231. fit_xs = np.linspace(x[0], x[-1], 100)
  232. fit_ys = f(fit_xs, *popt)
  233. ax.plot(fit_xs, fit_ys, '--g')
  234. arglabels = list(signature(f).parameters)[1:]
  235. label_txt = "\n".join('{:7s}={: 0.2G}'.format(label, value)
  236. for label, value in zip(arglabels, popt))
  237. ax.text(0.60, 0.95, label_txt, va='top', transform=ax.transAxes,
  238. fontsize='medium', family='monospace', usetex=False)
  239. if stats:
  240. _add_stats(h, title)
  241. else:
  242. ax.set_title(title)
  243. ax.grid(grid, color='#E0E0E0')
  244. def hist2d_plot(h, **kwargs):
  245. """ Plots a 2D ROOT histogram object using matplotlib """
  246. try:
  247. values, errors, xs, ys = h
  248. except (TypeError, ValueError):
  249. values, errors, xs, ys = hist2d(h)
  250. plt.xlabel(kwargs.pop('xlabel', ''))
  251. plt.ylabel(kwargs.pop('ylabel', ''))
  252. plt.title(kwargs.pop('title', ''))
  253. plt.pcolormesh(xs, ys, values, )
  254. # axes.colorbar() TODO: Re-enable this
  255. def hist_plot_stack(hists: list, labels: list = None):
  256. """
  257. Creates a stacked histogram in the current axes.
  258. :param hists: list of histogram
  259. :param labels:
  260. :return:
  261. """
  262. if len(hists) == 0:
  263. return
  264. if len(set([len(hist[0]) for hist in hists])) != 1:
  265. raise ValueError("all histograms must have the same number of bins")
  266. if labels is None:
  267. labels = [None for _ in hists]
  268. if len(labels) != len(hists):
  269. raise ValueError("Label mismatch")
  270. bottoms = [0 for _ in hists[0][0]]
  271. for hist, label in zip(hists, labels):
  272. centers = []
  273. widths = []
  274. heights = []
  275. for left, right, content in zip(hist[2][:-1], hist[2][1:], hist[0]):
  276. centers.append((right + left) / 2)
  277. widths.append(right - left)
  278. heights.append(content)
  279. plt.bar(centers, heights, widths, bottoms, label=label)
  280. for i, content in enumerate(hist[0]):
  281. bottoms[i] += content