def plot_fit_params(self, wsp): """ Plot all free fit parameters onto canvas Args: wsp (ROOT.RooWorkspace): workspace where the model is stored """ fit_var = get_var(wsp, self.fit_var) cans = OrderedDict() for bin_name in self.bins: frame = fit_var.frame(rf.Title('Fit Results')) plotname = '{}_{}'.format(self.full_model, bin_name) full_pdf = wsp.pdf(plotname) full_pdf.paramOn(frame, rf.Layout(0.1, 0.9, 0.9), rf.Format('NEU', rf.AutoPrecision(2))) can = r.TCanvas(create_random_str(32), 'rcan', 600, 600) can.cd() frame.findObject('{}_paramBox'.format(full_pdf.GetName())).Draw() cans[bin_name] = can #can.SaveAs(pdfname) # returns fit params for each fit return cans
def make_pt_eta_plot(plotvar, cuts, cut_plot_set, savename): """ Make pt-eta plot (2D) for passed variables and overlay the passed cuts """ pt_eta_hist = (50, 0, 20, 24, 0, 2.4) pt_eta = create_histogram(plotvar, pt_eta_hist) pt_eta.SetXTitle('p_{T}') pt_eta.SetYTitle('|#eta|') can = r.TCanvas(create_random_str(), 'c', 600, 600) can.cd() pt_eta.Draw('colz') # lines = {c: [] for c in cuts} # keep the TLines alive long enough for saving lines = [] for cut in cuts.values(): lines.append(draw_line(cut())) for i, cut in enumerate(lines): plot_on_canvas(can, cut, attr=(cut_plot_set[i], )) leg = r.TLegend(0.5, 0.3, 0.8, 0.4) leg.SetTextColor(0) leg.SetBorderSize(0) leg.SetFillStyle(0) for i, cut in enumerate(cuts.keys()): leg.AddEntry(lines[i][0], cut, 'l') leg.Draw() can.Update() # can.SaveAs(savename) return can, pt_eta, lines # for ipython so that they survive
def divide(num, denom, option='', **kwargs): """ Divide to histograms and return the ratio Args: num (ROOT.TH1): numerator histogram denom (ROOT.TH1): denominator histogram option (str, optional): Option string that will be passed to the Divide function. Only meaningful value is 'B', in which case Binomial uncertainties will be calculated instead of (uncorrelated) Poisson uncertainties. Keyword Args: name (str): If not empty this will be set as the name of the ratio TH1 [x|y]label (str): Labels for the x and y-axis. Returns: ratio (ROOT.TH1): Ratio histogram obtained from cloning num and then dividing it by denom """ ratio = num.Clone(kwargs.pop('name', create_random_str(8))) ratio.Divide(num, denom, 1, 1, option) set_labels(ratio, kwargs.pop('xlabel', ''), kwargs.pop('ylabel', '')) return ratio
def get_func(variable): """ Get the TF1 belonging to the passed variable, without fixing any parameters yet """ func_str = ANA_FUNC_STRINGS[variable] return r.TF1(create_random_str(), func_str, *XRANGES[variable])
def _tf1_helper(self, wsp, nameold, vrs): """ helper function that takes all vars in a string and replaces them by their value in the fit NOTE: only works if there is only dependence on the mean value of ONE binning variable (otherwise it wouldn't be a tf1) """ v_names = vrs.split(', ') listp = [('{}'.format(elm), '[' + str(i) + ']') for i, elm in enumerate(v_names)] name = replace_all(nameold, listp) mean_rgx = re.compile(r'<(\w+)>') #TODO: error message if more than one element in this set? mean_var = list(set(mean_rgx.findall(name))) if len(mean_var) != 1: # Only warning here. Normally the TF1 constructor should emit # another error message because the function doesn't compile below logging.warning('Found dependency on more than one variable') mean_var = mean_var[0] name = name.replace('<{}>'.format(mean_var), 'x') f1 = r.TF1(create_random_str(), name, self.binning[self.bin_vars.index(mean_var)][0], self.binning[self.bin_vars.index(mean_var)][-1]) for i, elm in enumerate(v_names): f1.SetParameter(i, get_var(wsp, elm).getVal()) f1.SetParError(i, get_var(wsp, elm).getError()) f1.SetParName(i, elm) return f1, mean_var
def _setup_canvas(can, **kwargs): """ Setup a TCanvas (resp. a TCanvasWrapper) to plot on Returns: TCanvasWrapper """ if can is None: can_name = create_random_str() size = kwargs.get('size', (600, 600)) can = r.TCanvas(can_name, '', 50, 50, *size) if not isinstance(can, TCanvasWrapper): can = TCanvasWrapper(can) if kwargs.get('logy', False): can.SetLogy() if kwargs.get('logx', False): can.SetLogx() if kwargs.get('logz', False): can.SetLogz() if kwargs.get('grid', False): can.SetGrid() return can
def create_histogram(var, hist_sett, **kwargs): """ Create a ROOT histogram from the passed variable(s) Args: var (np.array): Array with maximum of 3 columns containing the variables to plot. hist_set (tuple): Histogram settings, that are directly unpacked into the constructor of the ROOT histogram Keyword Args: name (str, optional): Name to be used for the histogram weights (np.array, optional): weight array with the same number of events as the var array. Each entry corresponds to the weight of the event {x,y,z}_axis (str): axis labels to be set for the histogram Returns: ROOT.TH{1,2,3}D: The histogram with the dimension corresponding to the number of columns of var """ name = kwargs.pop('name', '') if not name: name = create_random_str() # use the number of dimensions from the var to determine which sort of # histogram to use ndim = var.shape if len(ndim) == 1: ndim = 1 else: ndim = ndim[1] if ndim > 3 or ndim < 0: logging.error('Dimension of histogram is {}. Cannot create histogram' .format(ndim)) raise TypeError('Invalid number of dimensions in create_histograms') hist_type = 'TH{}D'.format(ndim) try: hist = getattr(r, hist_type)(name, '', *hist_sett) except TypeError as exc: logging.error('Could not construct TH{}D with passed hist_sett: {}' .format(ndim, hist_sett)) raise exc set_hist_opts(hist) # set axis labels xax, yax, zax = (kwargs.pop(a, '') for a in ['x_axis', 'y_axis', 'z_axis']) if xax: hist.SetXTitle(xax) if yax: hist.SetYTitle(yax) if zax: hist.SetZTitle(zax) fill_hist(hist, var, weights=kwargs.pop('weights', None)) return hist
def make_plot(hists, savename, reference, doleg=False): """ Make plot """ set_TDR_style() can = r.TCanvas(create_random_str(16), '', 50, 50, 600, 600) can.cd() # preliminary collect them and delete entries if necessary legentries = [] with_ref = False if reference is not None: ref_hist = hists[reference] with_ref = True rest_hists = [hists[k] for k in hists if k != reference] pull_hists = [do_ratio(ref_hist, h) for h in rest_hists] legentries = [reference] + [n for n in hists if n != reference] else: rest_hists = hists.values() legentries = hists.keys() if doleg: for i, entry in enumerate(legentries): legentries[i] = entry.replace('jpsi_kin_sel', '').replace('__', '_') else: legentries = [] y_max = get_y_max(hists.values()) * 1.1 leg=setup_legend() if with_ref: pad = r.TPad('ratio_pad', 'ratio_pad', 0, 0.3, 1, 1) r.SetOwnership(pad, False) pad.Draw() pad.cd() # passing all legentries here, since indexing doesn't work with an empty # list and mkplot picks only the first in this case make_ratio_plot(pad, [ref_hist], leg, legentries, yRange=[0, y_max], attr=gen_attributes, drawOpt='E2', legOpt='F') make_ratio_plot(pad, rest_hists, leg, legentries[1:], yRange=[0, y_max], attr=comp_attributes, drawOpt='sameE1', legOpt='PLE') can.cd() pull_pad = r.TPad('pull_pad', 'pull_pad', 0, 0, 1, 0.3) r.SetOwnership(pull_pad, False) pull_pad.Draw() pull_pad.cd() make_ratio_plot(pull_pad, pull_hists, None, [], yRange=[None, None], attr=comp_attributes) else: make_ratio_plot(can, rest_hists, leg, legentries, yRange=[0, y_max], attr=comp_attributes) can.SaveAs(savename)
def setUp(self): self.func = r.TF1('test_func', '[0] * (1 + [1] * x[0] + [2] * x[0]*x[0])', -2, 2) self.func.SetParameters(1, 2.0, -2.0) self.func_x_max = 2 self.func_x_min = -2 self.func_y_min = -11 self.func_y_max = 1.5 # Use TGraphAsymmErrors with random errors self.graph = r.TGraphAsymmErrors(10, np.linspace(0, 4, 10), np.sqrt(np.linspace(4, 16, 10)), np.random.uniform(size=10), np.random.uniform(size=10), np.random.uniform(size=10), np.random.uniform(size=10)) gxlo, gxhi, gylo, gyhi = get_errors(self.graph) graph_x = np.array(self.graph.GetX()) graph_y = np.array(self.graph.GetY()) self.graph_x_max = np.max(graph_x + gxhi) self.graph_x_min = np.min(graph_x - gxlo) self.graph_y_max = np.max(graph_y + gyhi) self.graph_y_min = np.min(graph_y - gylo) self.hist = r.TH1D(create_random_str(8), '', 10, -3, 1) self.hist.Fill(-1) # Fill one event into each bin to have non-trivial results for i in xrange(12): self.hist.Fill(-3.2 + 4.0 / 10 * i) self.hist_x_max = 1 self.hist_x_min = -3 self.hist_y_max = 2 self.hist_y_min = 1 # For 2D histograms we don't need to fill the histogram since the min # and max values do not depend on it self.hist2d = r.TH2D(create_random_str(8), '', 10, -3, 1, 10, 4, 5) self.hist2d_x_max = 1 self.hist2d_x_min = -3 self.hist2d_y_max = 5 self.hist2d_y_min = 4
def test_find_bin_nonreg_binning(self): hist = r.TH1D(create_random_str(8), '', 10, np.linspace(0, 1, 11)**2) binning = hu.get_binning(hist) values = np.random.uniform(0, 1, 1000) exp_idcs = np.array([hist.FindBin(v) for v in values]) exp_idcs -= 1 bin_idcs = hu.find_bin(binning, values) npt.assert_equal(bin_idcs, exp_idcs)
def test_find_bin_reg_binning(self): hist = r.TH1D(create_random_str(8), '', 10, 0, 1) binning = hu.get_binning(hist) values = np.random.uniform(0, 1, 1000) exp_idcs = np.array([hist.FindBin(v) for v in values]) exp_idcs -= 1 # correct for TH1 indexing starting at 1 bin_idcs = hu.find_bin(binning, values) npt.assert_equal(bin_idcs, exp_idcs)
def test_find_bin_warning(self, mock_logger): exp_warn = 'When trying to find the bin indices at least one value '\ 'could not be attributed to a bin in the passed binning' bins = hu.get_binning(hist = r.TH1D(create_random_str(), '', 10, 0, 1)) values = np.array([-0.1, 0.2, 0.3, 0.4]) bin_idcs = hu.find_bin(bins, values) mock_logger.warn.assert_called_with(exp_warn) values = np.array([0.1, 0.2, 1.3, 0.4, 0.5]) bin_idcs = hu.find_bin(bins, values) mock_logger.warn.assert_called_with(exp_warn)
def analytical_ratio(chi1_scen, chi2_scen): lth1 = lth_1(R=chi1_scen.split('_')[-1].replace('o', '/')) lth2 = lth_2(R1=chi2_scen.split('_')[2].replace('o', '/'), R2=chi2_scen.split('_')[4].replace('o', '/')) func_ratio = r.TF1(create_random_str(4), '[0] * (3 + [1]) / (3 + [2]) * (1 + [2] * x[0]*x[0]) / (1 + [1] * x[0]*x[0])', 0, 1) func_ratio.FixParameter(1, float(lth1)) func_ratio.FixParameter(2, float(lth2)) func_ratio.SetParameter(0, 1) return func_ratio, lth2 - lth1
def make_ratio_comp_plot(graph_file, plot_config, variable): """ Make the plot comparing the unscaled and scaled graphs in two different panels """ cgraph = graph_file.Get('central_r_chic2_chic1') vgraphs = get_variation_graphs(graph_file, '', plot_config) vsgraphs = get_variation_graphs(graph_file, '_scaled', plot_config) can = r.TCanvas(create_random_str(), '', 600, 600) can.cd() pad1 = r.TPad('unscaled_pad', '', 0, 0.5, 1, 1) r.SetOwnership(pad1, False) pad1.Draw() pad1.cd() leg = create_legend(0.675, 0.88, 0.16, plot_config, True) leg.SetTextSize(leg.GetTextSize() * 2) pad1 = mkplot(cgraph, drawOpt='PE', can=pad1, attr=CENTRAL_ATTR, yRange=YRANGE[variable], yLabel='unscaled ' + YLABELS['r_chic2_chic1'], leg=leg, legEntries=['nominal'], **VAR_PLOT[variable]) mkplot(shift_graphs(vgraphs, HORIZONTAL_SHIFT[variable]), can=pad1, drawOpt='samePEX0', leg=leg, legEntries=[v[1] for v in plot_config['variations']]) increase_label_size(pad1.pltables[0], 1.8) pad1.pltables[0].SetXTitle('') pad1.pltables[0].SetTitleOffset(0.75, 'Y') can.cd() pad2 = r.TPad('scaled_pad', '', 0, 0, 1, 0.535) pad2.SetBottomMargin(0.19) r.SetOwnership(pad2, False) pad2.Draw() pad2.cd() pad2 = mkplot(cgraph, drawOpt='PE', can=pad2, attr=CENTRAL_ATTR, yRange=YRANGE[variable], yLabel='normalized ' + YLABELS['r_chic2_chic1'], **VAR_PLOT[variable]) mkplot(shift_graphs(vsgraphs, HORIZONTAL_SHIFT[variable]), can=pad2, drawOpt='samePEX0') increase_label_size(pad2.pltables[0], 1.8) pad2.pltables[0].SetTitleOffset(0.9, 'X') pad2.pltables[0].SetTitleOffset(0.75, 'Y') # Make the pads survive can._pads = [pad1, pad2] return can
def rebin(hist, targ_bins): """ Rebin the passed histogram to the desired binning with the possibility to fix the number of desired bins along each axis. Args: hist (ROOT.TH1 or ROOT.THn): Histogram that should be rebinned targ_bins (list of tuples): Each tuple contains an axis identifier as first element and the desired number of bins as second element. Axis identifiers can either be labels (i.e. chars) or indices for TH1 or integer indices for THn. Returns: ROOT.TH1 or ROOT.THn: Histogram of the same type as the input histogram but with different number of bins along some axis. """ orig_bins = _get_nbins(hist) ndims = len(orig_bins) rebin_factors = np.ones(ndims, dtype='i4') for axl, nbins in targ_bins: axl = _get_bin_index(axl) fact = is_divisable(orig_bins[axl], nbins) if fact is None: logging.error('Cannot rebin axis {} to {} bins, because it is not ' 'a divisor of {}'.format(axl, nbins, orig_bins[axl])) return None rebin_factors[axl] = fact if isinstance(hist, r.THn): return hist.Rebin(rebin_factors) # for the cases where the hist is a TH1, we can also set the name to avoid # modifying the passed-in histogram name = hist.GetName() + '_rebin_' + create_random_str(4) # have to change types here to match the signature that ROOT expects # (Could be connected to the indexing into the array?) rebin_factors = rebin_factors.astype(int) if isinstance(hist, r.TH3): return hist.Rebin3D(rebin_factors[0], rebin_factors[1], rebin_factors[2], name) if isinstance(hist, r.TH2): return hist.Rebin2D(rebin_factors[0], rebin_factors[1], name) return hist.Rebin(rebin_factors[0], name)
def _get_slice(hist, direction, ilow, ihigh, name=None): """ Get the slice between index ilow and ihigh """ # TH2.GetProjection[X|Y] includes both passed bins, so we have to take care # that we do not "double count" bins # since 0 is the underflow bin, it is enough to bump the lower index by 1 axis = getattr(hist, 'Get' + OTHER_AXIS[direction] + 'axis')() vlow = axis.GetBinLowEdge(ilow + 1) vhigh = axis.GetBinUpEdge(ihigh) proj_f = getattr(hist, 'Projection' + direction) proj = proj_f('_'.join([name if name is not None else create_random_str(4), direction, str(vlow), str(vhigh)]), ilow + 1, ihigh) set_hist_opts(proj) return proj, (vlow, vhigh)
def r_costh(set_vals, fix_lambdas=True): """ R(costh) with possible fixed values """ val_to_idx = {'norm': 0, 'lth1': 1, 'lth2': 2} func = r.TF1(create_random_str(), '[0] * (1 + [2] * x[0]*x[0]) / (1 + [1] * x[0]*x[0])', 0, 1) def_vals = {'norm': 1.0, 'lth1': 0, 'lth2': 0} def_vals.update(set_vals) for name, val in def_vals.iteritems(): func.SetParameter(val_to_idx[name], val) if fix_lambdas: for ipar in [1, 2]: func.FixParameter(ipar, func.GetParameter(ipar)) return func
def _project_THn(hist, axes): """ Internal function for THn projections. For these it should be more straight forward to get the expected behavior from ROOT """ axes = np.array(list(make_iterable(axes)), dtype='i8') all_axes = list(xrange(hist.GetNdimensions())) sum_axes = tuple(a for a in all_axes if a not in axes) # Set ranges for axes that are integrated over to not include overflow for ax_idx in sum_axes: axis = hist.GetAxis(ax_idx) axis.SetRange(1, axis.GetNbins()) # create a name that does not interfere with the options name = replace_all(create_random_str(), (('e', '_'), ('E', '_'), ('a', '_'), ('A', '_'), ('o', '_'), ('O', '_'))) # If at all possible return a TH1 (because they are easier to handle) if len(axes) == 3: proj_hist = hist.Projection(axes[0], axes[1], axes[2], "E_" + name) elif len(axes) == 2: proj_hist = hist.Projection(axes[0], axes[1], "E_" + name) elif len(axes) == 1: proj_hist = hist.Projection(axes[0], "E_" + name) else: proj_hist = hist.ProjectionND(axes.shape[0], axes.astype('i4'), "E_" + name) # Make sure that the uncertainties are set / calculated # For some reason this seems to work when a TH1 is returned, but for THn # this status is somehow "lost", which can lead to trouble later if isinstance(proj_hist, r.TH1): if proj_hist.GetSumw2().fN == 0: proj_hist.Sumw2() if isinstance(proj_hist, r.THnBase): if not proj_hist.GetCalculateErrors(): proj_hist.Sumw2() return proj_hist
def plot_fit_params(self, wsp, pdfname, snapname=''): """ Plot all free fit parameters onto canvas and save as a pdf. Args: wsp (ROOT.RooWorkspace): workspace where the model is stored pdfname (str): Name of the created pdf under which the plot will be stored snapname (str, optional): Name of snapshot that will be loaded before plotting """ if snapname: wsp.loadSnapshot(snapname) frame = get_var(wsp, self.mname).frame(rf.Title('Fit Results')) full_pdf = wsp.pdf(self.full_model) full_pdf.paramOn(frame, rf.Layout(0.1, 0.9, 0.9), rf.Format('NEU', rf.AutoPrecision(2))) can = r.TCanvas(create_random_str(32), 'rcan', 600, 600) can.cd() frame.findObject('{}_paramBox'.format(full_pdf.GetName())).Draw() can.SaveAs(pdfname)
def w_costh_phi(set_vals=None, fix_norm=False, fix_lambdas=False): """ Angular 2d distribution as TF2, with possibly fixed values. Args: set_vals (dict, optional): Dictionary containing the values to which lth (lamba_theta), lph (lambda_phi), ltp (lambda_theta phi) or norm (the normalization) should be set. All values not found will be set to 0, except for the normalization that will be set to 1 fix_norm (boolean, optional): Fix the normalization to the set value (1 by default) fix_lambdas (boolean, optional): Fix the lambdas to the set values (0 by default) """ val_to_idx = {'norm': 0, 'lth': 1, 'lph': 2, 'ltp': 3} func_str = ('[0] * 3 / (4 * pi * (3 + [1])) * (' '1 + [1] * x[0]*x[0] + ' '[2] * (1 - x[0]*x[0]) * cos(2 * x[1] * pi / 180) + ' '[3] * 2 * x[0] * sqrt(1 - x[0]*x[0]) * cos(x[1] * pi / 180)' ')') func = r.TF2(create_random_str(), func_str, -1, 1, -180, 180) for name, idx in val_to_idx.iteritems(): func.SetParName(idx, name) default_vals = {'norm': 1, 'lth': 0, 'lph': 0, 'ltp': 0} if set_vals is not None: default_vals.update(set_vals) for name, val in default_vals.iteritems(): func.SetParameter(val_to_idx[name], val) fix_params = [] if fix_norm: fix_params.append(0) if fix_lambdas: fix_params.extend([1, 2, 3]) for ipar in fix_params: func.FixParameter(ipar, func.GetParameter(ipar)) return func
def plot(self, wsp, pdfname, snapname='', add_cut='', **kwargs): """ Make a plot of the model and the fitted data and save as pdf. Args: wsp (ROOT.RooWorkspace): workspace where all the information (including data and model) is stored pdfname (str): Name of the created pdf under which the plot will be stored. snapname (str, optional): Name of snapshot that will be loaded before plotting add_cut (str, optional): Additional cut to apply to the dataset before plotting. NOTE: all used variables have to be present in the dataset Keyword Args: logy (bool): Set log on y scale of distribution plot weighted_fit (bool): Assume that the fit has done using weights and calculate the data uncertainties using these weights. verbose (bool): Put information on the fit status and the covariance matrix quality onto the fit """ fitresname = None if snapname: wsp.loadSnapshot(snapname) fitresname = snapname.replace('snap', 'fit_res') # Starting from ROOT v6.12 it is possible to set this on individual axis r.TGaxis.SetMaxDigits(3) mvar = get_var(wsp, self.mname) plot_data = wsp.data('full_data') if add_cut: plot_data = plot_data.reduce(add_cut) nbins = int((mvar.getMax() - mvar.getMin()) / BIN_WIDTH) # TODO: find a better heuristic to do this here nevents_data = plot_data.numEntries() logging.debug( 'Number of events in reduced sample: {}'.format(nevents_data)) while nevents_data / nbins < 50: nbins /= 2 frame = mvar.frame(rf.Bins(nbins)) frame.SetTitle("") frame.GetYaxis().SetTitleOffset(1.3) frame.GetYaxis().SetTitle('Events / {:.1f} MeV'.format( (mvar.getMax() - mvar.getMin()) / nbins * 1000)) full_pdf = wsp.pdf(self.full_model) leg = self._setup_legend() if kwargs.pop('weighted_fit', True): plot_data.plotOn(frame, rf.MarkerSize(0.8), rf.Name('data_hist'), rf.DataError(r.RooAbsData.SumW2)) else: plot_data.plotOn(frame, rf.MarkerSize(0.8), rf.Name('data_hist')) full_pdf.plotOn( frame, rf.LineWidth(2), # rf.ProjWData(plot_data), rf.Name('full_pdf_curve')) leg.AddEntry(frame.getCurve('full_pdf_curve'), 'sum', 'l') for name, lst, lcol, legentry in self.components: full_pdf.plotOn(frame, rf.Components(name), rf.LineStyle(lst), rf.LineColor(lcol), rf.LineWidth(2), rf.Name(name)) leg.AddEntry(frame.getCurve(name), legentry, 'l') info_text = [] if fitresname is not None: fit_res = wsp.genobj(fitresname) chi2, ndf = get_chi2_ndf(frame, 'full_pdf_curve', 'data_hist', fit_res) info_text.append( (0.15, 0.875, '#chi^{{2}}/ndf = {:.1f} / {}'.format(chi2, ndf))) if kwargs.pop('verbose', False): status = get_var(wsp, '__fit_status__').getVal() cov_qual = get_var(wsp, '__cov_qual__').getVal() info_text.append( (0.15, 0.825, 'status = {}, covQual = {}'.format(int(status), int(cov_qual)))) if fitresname is not None: aic = calc_info_crit(fit_res) bic = calc_info_crit(fit_res, nevents_data) info_text.append( (0.15, 0.775, 'AIC = {:.1f}, BIC = {:.1f}'.format(aic, bic))) pull_frame = mvar.frame(rf.Bins(nbins)) hpull = frame.pullHist('data_hist', 'full_pdf_curve', True) hpull.SetMarkerSize(0.8) pull_frame.addPlotable(hpull, 'P') pull_frame.SetTitle("") pull_frame.GetYaxis().SetTitle("pull") pull_frame.GetXaxis().SetTitleSize(0.08) pull_frame.GetYaxis().SetTitleSize(0.08) pull_frame.GetXaxis().SetLabelSize(0.08) pull_frame.GetYaxis().SetLabelSize(0.08) pull_frame.GetYaxis().SetTitleOffset(0.4) pull_frame.GetYaxis().SetRangeUser(-5.99, 5.99) latex = setup_latex() can = r.TCanvas(create_random_str(16), '', 50, 50, 600, 600) can.cd() pad = r.TPad('mass_pad', 'mass_pad', 0, 0.3, 1, 1) r.SetOwnership(pad, False) pad.Draw() pad.cd() if kwargs.pop('logy', False): pad.SetLogy() pdfname = pdfname.replace('mass_fit', 'mass_fit_log') frame.Draw() leg.Draw() put_on_latex(latex, info_text, ndc=True) can.cd() pull_pad = r.TPad('pull_pad', 'pull_pad', 0, 0, 1, 0.3) r.SetOwnership(pull_pad, False) pull_pad.Draw() pull_pad.SetGridy() pull_pad.SetBottomMargin(0.2) pull_pad.cd() pull_frame.Draw() can.SaveAs(pdfname)
def baseline_plot(baseline, compplots, **kwargs): """ Make a plot and compare the compplots with the baseline plots. Divides the plot into an absolute value plot in the top and a ratio plot using the baseline as denominator in the bottom. Args: baseline (plotable ROOT object): The baseline plotable that will be used as comparison for all the compplots compplots (plotable ROOT objects): The plotables that should be compared to the baseline Keyword Args: basename (str): Legend entry to be used for the baseline legEntries (list of str): legend entries to be used for the compplots yRangeRatio (tuple of floats): The minimum and maximum y-value for the ratio pad compname (str): Name describing the whole of the compplots that will be used in the y-axis of the ratio pad putline (float or list of floats): Put horizontal lines into the ratio pad at the given values Other Keyword Args are forwarded to mkplot. Returns: TCanvasWrapper: Transparent wrapper holding the plot and all its objects See Also: mkplot, plot_on_canvas """ comp_attr = kwargs.pop('attr', None) if comp_attr is None: comp_attr = default_attributes(open_markers=False, size=1.0) # the baseline will always be black. Try to match the size of the markers to # the one that were requested by the user base_attr = {'color': 1, 'marker': 20, 'size': 1.5} sizes = np.unique([a['size'] for a in comp_attr if 'size' in a]) if len(sizes) == 1: base_attr['size'] = sizes[0] attr = [base_attr] + comp_attr # use the xLabel only for the lower plot xLabel = kwargs.pop('xLabel', None) # add the baseline name to the legend entries (if any) legEntries = kwargs.pop('legEntries', None) base_name = kwargs.pop('basename', 'baseline') if legEntries is not None: legEntries = [base_name] + legEntries # setup canvas can = kwargs.pop('can', None) if can is None: can_name = create_random_str() can = r.TCanvas(can_name, '', 50, 50, 600, 600) can.cd() ppad = r.TPad('_'.join([can.GetName(), 'plotpad']), 'plotpad', 0, 0.3, 1, 1) r.SetOwnership(ppad, False) ppad.Draw() # plot the comparison plot ppad = mkplot([baseline] + make_iterable(compplots), attr=attr, can=ppad, legEntries=legEntries, **kwargs) can.cd() rpad = r.TPad('_'.join([can.GetName(), 'ratiopad']), 'rpad', 0, 0, 1, 0.33) rpad.SetBottomMargin(0.2) r.SetOwnership(rpad, False) rpad.Draw() # remove some kwargs for kwarg in ['yLabel', 'legPos', 'leg', 'legEntries', 'yRange', 'logy']: kwargs.pop(kwarg, None) ratios = [divide(p, baseline) for p in make_iterable(compplots)] # determine the ratios and plot them rpad = mkplot(ratios, attr=comp_attr, can=rpad, xLabel=xLabel, yLabel=' / '.join( [kwargs.pop('compname', 'distribution(s)'), base_name]), yRange=kwargs.pop('yRangeRatio', [None, None]), **kwargs) for hist in rpad.pltables: _set_ratio_properties(hist) putlines = kwargs.pop('putline', None) if putlines is not None: # Simply assume that all others follow the same as the first one lines = _get_ratio_lines(rpad.pltables[0], putlines) for line in lines: line.Draw() rpad.add_pltables(lines) # attach all plots to the returned canvas if not isinstance(can, TCanvasWrapper): can = TCanvasWrapper(can) # can.add_pltables(ppad.pltables + rpad.pltables) # for obj in ppad.attached_tobjects: # can.add_tobject(obj) can.add_tobject(ppad) can.add_tobject(rpad) return can
def from_array(array, binning, **kwargs): """ Create a ROOT histogram from the passed array Overflow is directly handled by the array2hist function used internally Args: array (np.array): Array holding the bin contents for each bin binning (np.array): Array holding the binning information for each direction Keyword Args: errors (np.array, optional): Array of the same shape as array holding the bin uncertainties Returns: ROOT.TH1: Up to TH3D depending on the shape of the passed array See also: root_numpy.array2hist """ n_dim = len(array.shape) if n_dim > 3 or n_dim < 0: raise ValueError('Invalid number of dimensions of array: {}' .format(n_dim)) raise_err = False if binning.shape[0] != n_dim: if n_dim == 1: shape = [len(binning) + 1] binning = np.array([binning]) else: raise_err = True # Necessary to have shape defined for all possible (in)valid inputs shape = binning.shape else: shape = [len(b) + 1 for b in binning] # include under- and overflow if not raise_err: for axis, bins in enumerate(shape): if array.shape[axis] != bins and array.shape[axis] != bins - 2: raise_err = True if raise_err: raise ValueError('Invalid binning for passed array: {} vs {}' .format(array.shape, shape)) uncer = kwargs.pop('errors', None) if uncer is not None: if uncer.shape != array.shape: raise ValueError('array and errors have to have the same shape') hist_type = 'TH{}D'.format(n_dim) hist_sett = list(chain.from_iterable((len(b) - 1, b.astype(np.float64)) for b in binning)) try: hist = getattr(r, hist_type)(create_random_str(8), '', *hist_sett) except TypeError as exc: logging.error ('Could not construct {} with passed binning: {}'. format(hist_type, binning)) raise exc set_hist_opts(hist) if rnp_version == '4.7.2': array2hist(array, hist) if uncer is not None: # Need to take care of the under- and overflow here since SetError # treats the first as underflow (See implementation of array2hist) # No need to do all the checks since array2hist will already fail if # array has the wrong shape and uncer has the same shape as array slices = [] for axis, bins in enumerate(shape): if uncer.shape[axis] == bins - 2: slices.append(slice(1, -1)) else: slices.append(slice(None)) uncer_of = np.zeros(shape, dtype=np.float64) uncer_of[tuple(slices)] = uncer uncer = uncer_of hist.SetError(np.ravel(uncer.T)) if rnp_version == '4.7.3': logging.debug('Using root_numpys capabilities of setting the errors') array2hist(array, hist, errors=uncer) return hist