|
|
@@ -1,4 +1,5 @@
|
|
|
#!/usr/bin/env python3
|
|
|
+import math
|
|
|
import matplotlib as mpl
|
|
|
mpl.rc('font', **{'family': 'sans-serif', 'sans-serif': ['Helvetica']})
|
|
|
mpl.rc('font', **{'family': 'serif', 'serif': ['Palatino']})
|
|
|
@@ -6,6 +7,26 @@ mpl.rc('text', usetex=True)
|
|
|
mpl.rc('savefig', dpi=120)
|
|
|
|
|
|
|
|
|
+def add_decorations(axes, luminosity, energy):
|
|
|
+ cms_prelim = r'{\raggedright{}\textsf{\textbf{CMS}}\\ \emph{Preliminary}}'
|
|
|
+ axes.text(0.01, 0.98, cms_prelim,
|
|
|
+ horizontalalignment='left',
|
|
|
+ verticalalignment='top',
|
|
|
+ transform=axes.transAxes)
|
|
|
+
|
|
|
+ lumi = ""
|
|
|
+ energy_str = ""
|
|
|
+ if luminosity is not None:
|
|
|
+ lumi = r'${} \mathrm{{fb}}^{{-1}}$'.format(luminosity)
|
|
|
+ if energy is not None:
|
|
|
+ energy_str = r'({} TeV)'.format(energy)
|
|
|
+
|
|
|
+ axes.text(1, 1, ' '.join([lumi, energy_str]),
|
|
|
+ horizontalalignment='right',
|
|
|
+ verticalalignment='bottom',
|
|
|
+ transform=axes.transAxes)
|
|
|
+
|
|
|
+
|
|
|
class StackHist:
|
|
|
|
|
|
def __init__(self, title=""):
|
|
|
@@ -22,30 +43,29 @@ class StackHist:
|
|
|
self.data = None
|
|
|
|
|
|
@staticmethod
|
|
|
- def to_bin_list(th1, scale=1):
|
|
|
+ def to_bin_list(th1):
|
|
|
bins = []
|
|
|
for i in range(th1.GetNbinsX()):
|
|
|
center = th1.GetBinCenter(i + 1)
|
|
|
width = th1.GetBinWidth(i + 1)
|
|
|
content = th1.GetBinContent(i + 1)
|
|
|
- bins.append((center-width/2, center+width/2, content*scale))
|
|
|
+ bins.append((center-width/2, center+width/2, content))
|
|
|
return bins
|
|
|
|
|
|
- def add_mc_background(self, th1, label, lumi=None):
|
|
|
- self.backgrounds.append((label, lumi, self.to_bin_list(th1)))
|
|
|
+ def add_mc_background(self, th1, label, lumi=None, plot_color=''):
|
|
|
+ self.backgrounds.append((label, lumi, self.to_bin_list(th1), plot_color))
|
|
|
|
|
|
- def set_mc_signal(self, th1, label, lumi=None, stack=True, scale=1):
|
|
|
- if scale != 1:
|
|
|
- label = r"{}$\times{:02d}$".format(label, scale)
|
|
|
- self.signal = (label, lumi, self.to_bin_list(th1, scale))
|
|
|
+ def set_mc_signal(self, th1, label, lumi=None, stack=True, scale=1, plot_color=''):
|
|
|
+ self.signal = (label, lumi, self.to_bin_list(th1), plot_color)
|
|
|
self.signal_stack = stack
|
|
|
+ self.signal_scale = scale
|
|
|
|
|
|
- def set_data(self, th1, lumi=None):
|
|
|
- self.data = ('data', lumi, self.to_bin_list(th1))
|
|
|
+ def set_data(self, th1, lumi=None, plot_color=''):
|
|
|
+ self.data = ('data', lumi, self.to_bin_list(th1), plot_color)
|
|
|
self.luminosity = lumi
|
|
|
|
|
|
def _verify_binning_match(self):
|
|
|
- bins_count = [len(bins) for label, lumi, bins in self.backgrounds]
|
|
|
+ bins_count = [len(bins) for _, _, bins, _ in self.backgrounds]
|
|
|
if self.signal is not None:
|
|
|
bins_count.append(len(self.signal[2]))
|
|
|
if self.data is not None:
|
|
|
@@ -53,37 +73,40 @@ class StackHist:
|
|
|
n_bins = bins_count[0]
|
|
|
if any(bin_count != n_bins for bin_count in bins_count):
|
|
|
raise ValueError("all histograms must have the same number of bins")
|
|
|
- return n_bins
|
|
|
-
|
|
|
- def _add_decorations(self, axes):
|
|
|
- cms_prelim = r'{\raggedright{}\textsf{\textbf{CMS}}\\ \emph{Preliminary}}'
|
|
|
- axes.text(0.01, 0.99, cms_prelim,
|
|
|
- horizontalalignment='left',
|
|
|
- verticalalignment='top',
|
|
|
- transform=axes.transAxes)
|
|
|
-
|
|
|
- lumi = ""
|
|
|
- energy = ""
|
|
|
- if self.luminosity is not None:
|
|
|
- lumi = r'${} \mathrm{{fb}}^{{-1}}$'.format(self.luminosity)
|
|
|
- if self.energy is not None:
|
|
|
- energy = r'({} TeV)'.format(self.energy)
|
|
|
-
|
|
|
- axes.text(1, 1, ' '.join([lumi, energy]),
|
|
|
- horizontalalignment='right',
|
|
|
- verticalalignment='bottom',
|
|
|
- transform=axes.transAxes)
|
|
|
-
|
|
|
- def draw(self, axes):
|
|
|
- n_bins = self._verify_binning_match()
|
|
|
- bottoms = [0]*n_bins
|
|
|
+ self.n_bins = n_bins
|
|
|
+
|
|
|
+ def save(self, fname):
|
|
|
+ from matplotlib.transforms import Bbox
|
|
|
+ def full_extent(ax, pad=0.0):
|
|
|
+ """Get the full extent of an axes, including axes labels, tick labels, and
|
|
|
+ titles."""
|
|
|
+ # For text objects, we need to draw the figure first, otherwise the extents
|
|
|
+ # are undefined.
|
|
|
+ ax.figure.canvas.draw()
|
|
|
+ items = ax.get_xticklabels() + ax.get_yticklabels()
|
|
|
+ items += [ax, ax.title, ax.xaxis.label, ax.yaxis.label]
|
|
|
+ # items += [ax, ax.title]
|
|
|
+ bbox = Bbox.union([item.get_window_extent() for item in items])
|
|
|
+
|
|
|
+ return bbox.expanded(1.0 + pad, 1.0 + pad)
|
|
|
+
|
|
|
+ extents = []
|
|
|
+ for axes in self.axeses:
|
|
|
+ extents.append(full_extent(axes).transformed(axes.figure.dpi_scale_trans.inverted()))
|
|
|
+ extent = Bbox.union(extents)
|
|
|
+ axes.figure.savefig('figures/'+fname, bbox_inches=extent)
|
|
|
+
|
|
|
+ def do_draw(self, axes):
|
|
|
+ self.axeses = [axes]
|
|
|
+ self._verify_binning_match()
|
|
|
+ bottoms = [0]*self.n_bins
|
|
|
|
|
|
if self.logx:
|
|
|
axes.set_xscale('log')
|
|
|
if self.logy:
|
|
|
axes.set_yscale('log')
|
|
|
|
|
|
- def draw_bar(label, lumi, bins, stack=True, **kwargs):
|
|
|
+ def draw_bar(label, lumi, bins, plot_color, scale=1, stack=True, **kwargs):
|
|
|
if stack:
|
|
|
lefts = []
|
|
|
widths = []
|
|
|
@@ -93,12 +116,14 @@ class StackHist:
|
|
|
widths.append(right-left)
|
|
|
if lumi is not None:
|
|
|
content *= self.luminosity/lumi
|
|
|
+ content *= scale
|
|
|
heights.append(content)
|
|
|
|
|
|
- axes.bar(lefts, heights, widths, bottoms, label=label, **kwargs)
|
|
|
+ axes.bar(lefts, heights, widths, bottoms, label=label, color=plot_color, **kwargs)
|
|
|
for i, (_, _, content) in enumerate(bins):
|
|
|
if lumi is not None:
|
|
|
content *= self.luminosity/lumi
|
|
|
+ content *= scale
|
|
|
bottoms[i] += content
|
|
|
else:
|
|
|
xs = [bins[0][0] - (bins[0][1]-bins[0][0])/2]
|
|
|
@@ -107,31 +132,102 @@ class StackHist:
|
|
|
width2 = (right-left)/2
|
|
|
if lumi is not None:
|
|
|
content *= self.luminosity/lumi
|
|
|
+ content *= scale
|
|
|
xs.append(left-width2)
|
|
|
ys.append(content)
|
|
|
xs.append(right-width2)
|
|
|
ys.append(content)
|
|
|
xs.append(bins[-1][0] + (bins[-1][1]-bins[-1][0])/2)
|
|
|
ys.append(0)
|
|
|
- axes.plot(xs, ys, label=label, **kwargs)
|
|
|
+ axes.plot(xs, ys, label=label, color=plot_color, **kwargs)
|
|
|
|
|
|
if self.signal is not None and self.signal_stack:
|
|
|
- draw_bar(*self.signal, hatch='/')
|
|
|
+ label, lumi, bins, plot_color = self.signal
|
|
|
+ if self.signal_scale != 1:
|
|
|
+ label = r"{}$\times{:d}$".format(label, self.signal_scale)
|
|
|
+ draw_bar(label, lumi, bins, plot_color, scale=self.signal_scale, hatch='/')
|
|
|
|
|
|
for background in self.backgrounds:
|
|
|
draw_bar(*background)
|
|
|
|
|
|
if self.signal is not None and not self.signal_stack:
|
|
|
- draw_bar(*self.signal, stack=False, color='k')
|
|
|
+ # draw_bar(*self.signal, stack=False, color='k')
|
|
|
+ label, lumi, bins, plot_color = self.signal
|
|
|
+ if self.signal_scale != 1:
|
|
|
+ label = r"{}$\times{:d}$".format(label, self.signal_scale)
|
|
|
+ draw_bar(label, lumi, bins, plot_color, scale=self.signal_scale, stack=False)
|
|
|
|
|
|
axes.set_title(self.title)
|
|
|
axes.set_xlabel(self.xlabel)
|
|
|
axes.set_ylabel(self.ylabel)
|
|
|
axes.set_xlim(*self.xlim)
|
|
|
# axes.set_ylim(*self.ylim)
|
|
|
- axes.set_ylim(None, max(bottoms)*1.2)
|
|
|
+ if self.logy:
|
|
|
+ axes.set_ylim(None, math.exp(math.log(max(bottoms))*1.4))
|
|
|
+ else:
|
|
|
+ axes.set_ylim(None, max(bottoms)*1.2)
|
|
|
axes.legend(frameon=True, ncol=2)
|
|
|
- self._add_decorations(axes)
|
|
|
+ add_decorations(axes, self.luminosity, self.energy)
|
|
|
+
|
|
|
+ def draw(self, axes, save=True, **kwargs):
|
|
|
+ self.do_draw(axes, **kwargs)
|
|
|
+ if save:
|
|
|
+ self.save(self.title+".png")
|
|
|
+
|
|
|
+
|
|
|
+class StackHistWithSignificance(StackHist):
|
|
|
+
|
|
|
+ def __init__(self, *args, **kwargs):
|
|
|
+ super().__init__(*args, **kwargs)
|
|
|
+
|
|
|
+ def do_draw(self, axes, bin_significance=True, low_cut_significance=False, high_cut_significance=False):
|
|
|
+ bottom_box, _, top_box = axes.get_position().splity(0.28, 0.30)
|
|
|
+ axes.set_position(top_box)
|
|
|
+ super().do_draw(axes)
|
|
|
+ axes.set_xticks([])
|
|
|
+ rhs_color = '#cc6600'
|
|
|
+
|
|
|
+ bottom = axes.get_figure().add_axes(bottom_box)
|
|
|
+ bottom_rhs = bottom.twinx()
|
|
|
+ self.axeses = [axes, bottom, bottom_rhs]
|
|
|
+ bgs = [0]*self.n_bins
|
|
|
+ for (_, _, bins, _) in self.backgrounds:
|
|
|
+ for i, (left, right, value) in enumerate(bins):
|
|
|
+ bgs[i] += value
|
|
|
+
|
|
|
+ sigs = [0]*self.n_bins
|
|
|
+ if bin_significance:
|
|
|
+ xs = []
|
|
|
+ for i, (left, right, value) in enumerate(self.signal[2]):
|
|
|
+ sigs[i] += value
|
|
|
+ xs.append(left)
|
|
|
+ xs, ys = zip(*[(x, sig/(sig+bg)) for x, sig, bg in zip(xs, sigs, bgs) if (sig+bg)>0])
|
|
|
+ bottom.plot(xs, ys, '.k')
|
|
|
+
|
|
|
+ if high_cut_significance:
|
|
|
+ # s/(s+b) for events passing a minimum cut requirement
|
|
|
+ min_bg = [sum(bgs[i:]) for i in range(self.n_bins)]
|
|
|
+ min_sig = [sum(sigs[i:]) for i in range(self.n_bins)]
|
|
|
+ min_xs, min_ys = zip(*[(x, sig/math.sqrt(sig+bg)) for x, sig, bg in zip(xs, min_sig, min_bg)
|
|
|
+ if (sig+bg) > 0])
|
|
|
+ bottom_rhs.plot(min_xs, min_ys, '->', color=rhs_color)
|
|
|
+
|
|
|
+ if low_cut_significance:
|
|
|
+ # s/(s+b) for events passing a maximum cut requirement
|
|
|
+ max_bg = [sum(bgs[:i]) for i in range(self.n_bins)]
|
|
|
+ max_sig = [sum(sigs[:i]) for i in range(self.n_bins)]
|
|
|
+ max_xs, max_ys = zip(*[(x, sig/math.sqrt(sig+bg)) for x, sig, bg in zip(xs, max_sig, max_bg)
|
|
|
+ if (sig+bg) > 0])
|
|
|
+ bottom_rhs.plot(max_xs, max_ys, '-<', color=rhs_color)
|
|
|
+
|
|
|
+ bottom.set_ylabel(r'$S/(S+B)$')
|
|
|
+ bottom.set_xlim(axes.get_xlim())
|
|
|
+ bottom.set_ylim((0, 1.1))
|
|
|
+ if low_cut_significance or high_cut_significance:
|
|
|
+ bottom_rhs.set_ylabel(r'$S/\sqrt{S+B}$')
|
|
|
+ bottom_rhs.yaxis.label.set_color(rhs_color)
|
|
|
+ bottom_rhs.tick_params(axis='y', colors=rhs_color, size=4, width=1.5)
|
|
|
+ # bottom.grid()
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
Caleb Fangmeier