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 create_workspace(model, datafile, binvar, binning, massrange, fitfile, weights=None): """ Create the workspace with the data already imported and the model defined, also in charge of writing the bin info json file """ wsp = r.RooWorkspace('ws_mass_fit') massrange = [float(v) for v in massrange.split(',')] # load the data and apply the mass selection of the fitting range immediately bin_var = parse_func_var(binvar) # necessary for loading variables = [model.mname, bin_var[0]] if weights is not None: variables.append(weights) data = apply_selections(get_dataframe(datafile, columns=variables), select_bin(model.mname, *massrange)) costh_bins, costh_means = get_costh_bins(binning, bin_var, data) create_bin_info_json(fitfile.replace('.root', '_bin_sel_info.json'), costh_bins, costh_means, bin_var[0], datafile) # Create the variables in the workspace try_factory(wsp, '{}[{}, {}]'.format(model.mname, *massrange)) if 'abs' in bin_var[1].__name__: try_factory( wsp, '{}[{}, {}]'.format(bin_var[0], -np.max(costh_bins), np.max(costh_bins))) else: try_factory( wsp, '{}[{}, {}]'.format(bin_var[0], np.min(costh_bins), np.max(costh_bins))) dset_vars = r.RooArgSet(get_var(wsp, model.mname), get_var(wsp, bin_var[0])) tree = array2tree(data.to_records(index=False)) if weights is not None: try_factory(wsp, '{}[0, 1e5]'.format(weights)) dataset = r.RooDataSet('full_data', 'full data sample', tree, dset_vars, '', weights) else: dataset = r.RooDataSet('full_data', 'full data sample', tree, dset_vars) ws_import(wsp, dataset) return wsp, costh_bins
def _pull_plot(self, wsp, frame, publication): """ Make the frame containing the pulls """ pulls = frame.pullHist('data_hist', 'full_pdf_curve', True) pulls.SetMarkerSize(0.8) pull_frame = get_var(wsp, self.fit_var).frame() if publication: pull_frame.addPlotable(pulls, 'PEX0') pulls.SetMarkerSize(0.7) pull_frame.GetXaxis().SetTitle(YLABELS.get('chicMass')) else: pull_frame.addPlotable(pulls, 'P') pull_frame.SetTitle("") pull_frame.GetYaxis().SetTitle("pull") if publication: pull_frame.GetXaxis().SetTitleSize(0.12) pull_frame.GetYaxis().SetTitleSize(0.12) pull_frame.GetXaxis().SetLabelSize(0.12) pull_frame.GetYaxis().SetLabelSize(0.12) pull_frame.GetYaxis().SetTitleOffset(0.625) pull_frame.GetYaxis().SetRangeUser(-3.99, 3.99) pull_frame.GetYaxis().SetNdivisions(507) else: 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) return pull_frame
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 main(args): """Main""" with open(args.configfile, 'r') as conff: config = json.load(conff) model = BinnedFitModel(config) wsp = r.RooWorkspace('test_ws') test_vars = {} dset_vars = r.RooArgSet() for var, bounds in model.get_load_vars(): try_factory(wsp, '{}[{}]'.format(var, bounds)) low, high = [float(v) for v in bounds.split(',')] test_vars[var] = np.random.uniform(low, high, 10000) dset_vars.add(get_var(wsp, var)) data = pd.DataFrame(test_vars) tree = array2tree(data.to_records(index=False)) dataset = r.RooDataSet('full_data', 'test data', tree, dset_vars) ws_import (wsp, dataset) if model.define_model(wsp): print('configfile is a valid model definition')
def _plot_bin(self, wsp, full_data, bin_name, bin_borders, n_bins, publication=False): """Make the distribution plot for a given bin""" leg = setup_legend(*self.plot_config['legpos']) # In order to remove horizontal error bars from the plot the option # XErrorSize(0) has to be used, but that breaks the calculation of the # chi2 wrt a given curve since the x-errors are used to determine the # bin and the corresponding mean value of the data points in the bin # which is necessary to calculate the point at which the curve should be # evaluated. So here I plot two curves, the one labeled 'data_hist_plot' # is the visible one, whereas the one labeled 'data_hist' will be the # with associated x uncertainties that can be used to calculate the chi2 # This will only be done in "publication mode" if publication: data_args = (rf.MarkerSize(0.5), rf.Name('data_hist_plot'), rf.XErrorSize(0)) leg.SetTextSize(0.04) else: data_args = (rf.MarkerSize(0.7), rf.Name('data_hist')) fit_var = get_var(wsp, self.fit_var) frame = fit_var.frame(rf.Bins(n_bins)) cut = get_bin_cut(self.bin_cut_vars, bin_borders) full_data.reduce(cut).plotOn(frame, *data_args) if publication: full_data.reduce(cut).plotOn(frame, rf.MarkerSize(0), rf.Name('data_hist'), rf.LineColor(0), rf.LineWidth(0)) leg.AddEntry(frame.getHist('data_hist_plot'), 'Data', 'PE') full_pdf = wsp.pdf(self.full_model + '_' + bin_name) full_pdf.plotOn(frame, rf.LineWidth(2), rf.Name('full_pdf_curve')) leg.AddEntry(frame.getCurve('full_pdf_curve'), 'Fit result', 'l') for name, settings in self.components: full_pdf.plotOn(frame, rf.Components(name + '_' + bin_name), rf.LineStyle(settings['line']), rf.LineColor(settings['color']), rf.LineWidth(2), rf.Name(name)) leg.AddEntry(frame.getCurve(name), settings['label'], 'l') # calculate the bin width in MeV bw_mev = lambda v, n: (v.getMax() - v.getMin()) / n * 1000 frame.GetYaxis().SetTitleOffset(1.3) frame.GetYaxis().SetTitle('Events / {:.2f} MeV'.format( bw_mev(fit_var, n_bins))) frame.SetTitle("") return frame, leg
def get_yields(wsp, binname, yield_names): """ Get the sum of all the yields for a given bin """ yields = [] for name in yield_names: var_name = '_'.join([name, binname]) yields.append(get_var(wsp, var_name).getVal()) return np.sum(yields)
def bin_mean(self, wsp, var, bin_name): """ Get the mean value of the passed variable in the passed bin """ bin_data = wsp.data('full_data').reduce( get_bin_cut(self.bin_cut_vars, self.bins[bin_name])) # Check if the variable can be taken as it is or if it has to be treated # specially if var in self.bin_cut_vars: meanval = bin_data.mean(get_var(wsp, var)) else: # Check to which bin cut variable it belongs form = [v for v in self.bin_cut_vars if var in v] if len(form) > 1: logging.warning('Cannot uniquely identify to which bin ' 'definition \'{}\' belongs'.format(var)) # Create a formula var and add it to the temporary data set to # easily obtain the mean value binvarf = r.RooFormulaVar('binvar', 'temporary bin variable', form[0], r.RooArgList(get_var(wsp, var))) bin_var = bin_data.addColumn(binvarf) meanval = bin_data.mean(bin_var) return meanval
def get_state_fractions(bin_data, wsp, model, binname): """ Get the probabilities for each event to be a chi1 or a chi2 event for a given bin """ full_pdf = wsp.pdf('_'.join([model.full_model, binname])) chi1_pdf = wsp.pdf('_'.join([CHI1_MODEL, binname])) chi2_pdf = wsp.pdf('_'.join([CHI2_MODEL, binname])) full_yields = get_yields(wsp, binname, model.nevent_yields) chi1_yields = get_yields(wsp, binname, ['Nchic1']) chi2_yields = get_yields(wsp, binname, ['Nchic2']) mname = model.fit_var mvar = get_var(wsp, mname) full_pdf_vs = eval_pdf(full_pdf, mvar, bin_data.loc[:, mname]) * full_yields chi1_pdf_vs = eval_pdf(chi1_pdf, mvar, bin_data.loc[:, mname]) * chi1_yields chi2_pdf_vs = eval_pdf(chi2_pdf, mvar, bin_data.loc[:, mname]) * chi2_yields return chi1_pdf_vs / full_pdf_vs, chi2_pdf_vs / full_pdf_vs
def create_workspace(workspace_name, datafile, model): """ Create the workspace and already load the data into it """ wsp = r.RooWorkspace(workspace_name) dset_vars = r.RooArgSet() variables = [] for var, bounds in model.get_load_vars(): try_factory(wsp, '{}[{}]'.format(var, bounds)) dset_vars.add(get_var(wsp, var)) variables.append(var) data = get_dataframe(datafile, columns=variables) tree = array2tree(data.to_records(index=False)) dataset = r.RooDataSet('full_data', 'full data sample', tree, dset_vars) ws_import(wsp, dataset) wsp.Print() return wsp
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 define_model(self, wsp): """ Define the jpsi mass model in the passed workspace Args: wsp (ROOT.RooWorkspace): workspace into which the fit model is imported """ # rapidity dependent sigma of J/psi wsp.factory('CBsigma_p0_jpsi[0.02, 0.015, 0.025]') wsp.factory('CBsigma_p1_jpsi[0, -0.01, 0.01]') wsp.factory('CBsigma_p2_jpsi[0.01, 0.005, 0.025]') CBsigma_jpsi = r.RooFormulaVar('CBsigma_jpsi', '@0 + @1 * abs(@3) + @2 * abs(@3) * abs(@3)', r.RooArgList(get_var(wsp, 'CBsigma_p0_jpsi'), get_var(wsp, 'CBsigma_p1_jpsi'), get_var(wsp, 'CBsigma_p2_jpsi'), get_var(wsp, 'JpsiRap'))) wsp.factory('CBalpha_p0_jpsi[1.729, 1.2, 2.5]') wsp.factory('CBalpha_p1_jpsi[0.191, 0, 0.5]') CBalpha_jpsi = r.RooFormulaVar('CBalpha_jpsi', '@0 + @1 * abs(@2)', r.RooArgList(get_var(wsp, 'CBalpha_p0_jpsi'), get_var(wsp, 'CBalpha_p1_jpsi'), get_var(wsp, 'JpsiRap'))) wsp.factory('CBmass_p0_jpsi[3.094, 3.086, 3.098]') wsp.factory('CBmass_p1_jpsi[0.001, -0.002, 0.002]') wsp.factory('CBmass_p2_jpsi[-0.003, -0.005, 0.001]') CBmass_jpsi = r.RooFormulaVar('CBmass_jpsi', '@0 + @1 * abs(@3) + @2 * abs(@3) * abs(@3)', r.RooArgList(get_var(wsp, 'CBmass_p0_jpsi'), get_var(wsp, 'CBmass_p1_jpsi'), get_var(wsp, 'CBmass_p2_jpsi'), get_var(wsp, 'JpsiRap'))) # Fraction to ensure that sigma 2 is smaller than sigma 1 wsp.factory('r_sigma2_jpsi[0.5, 0, 0.9]') CBsigma2_jpsi = r.RooFormulaVar('CBsigma2_jpsi', '@0 * @1', r.RooArgList(get_var(wsp, 'r_sigma2_jpsi'), CBsigma_jpsi)) ws_import(wsp, r.RooArgSet(CBalpha_jpsi, CBmass_jpsi, CBsigma2_jpsi)) wsp.factory('RooCBShape::{}({}, CBmass_jpsi, CBsigma_jpsi, CBalpha_jpsi,' 'CBn_jpsi[2.5, 1.8, 6])'.format('CB_shape1_jpsi', self.mname)) # wsp.factory('RooCBShape::{}({}, CBmass_jpsi, CBsigma2_jpsi, CBalpha_jpsi,' # ' CBn_jpsi)'.format('CB_shape2_jpsi', self.mname)) wsp.factory('RooGaussian::{}({}, CBmass_jpsi, CBsigma2_g_jpsi[0.01, 0.0, 0.025])' .format('CB_shape2_jpsi', self.mname)) wsp.factory('SUM::{}(frac_CB1[0.5, 0, 1] * CB_shape1_jpsi, CB_shape2_jpsi)' .format(self.signal)) wsp.factory('RooExponential::{}({}, lambda_jpsi[0, -10, 10])' .format(self.bkg_model, self.mname)) # fix some parameters (as was done previously, values also from there) par_fix_vals = [ # ('CBsigma_p1_jpsi', 0), ('CBmass_p1_jpsi', 0), ('CBmass_p2_jpsi', 0), # ('CBalpha_p1_jpsi', 0), # ('CBsigma_p2_jpsi', 0.0125), ('CBn_jpsi', 2.5) ] fix_params(wsp, par_fix_vals) wsp.factory('{}[1e6, 0, 5e7]'.format(self.nevent_vars[0])) # Njpsi wsp.factory('{}[1e5, 0, 5e6]'.format(self.nevent_vars[-1])) # Nbkg_jpsi wsp.factory('SUM::{}({} * {}, {} * {})' .format(self.full_model, self.nevent_vars[0], self.signal, self.nevent_vars[-1], self.bkg_model))
def plot_simvar_graphs(self, wsp, simvars, sym_uncer=False): """ Create TGraphs for all the "simple" proto parameters Args: wsp (ROOT.RooWorkspace): Workspace holding the fit results from a simultaneous fit. simvars (list of strings): List of proto parameter names that do not have a functional dependency but instead are independent for each bin. sym_uncer (boolean, optional): Return the graphs with symmetric uncertainties instead of asymmetric ones (False) Returns: list of TGraphAsymmErrors: Graphs of the passed proto_parameters as a function of the binning variables. For each bin in either direction a graph is created as a function of the other direction. The name of the returned graph indicates on which variable the graph depends and in which bin of the other variable the values are taken from. E.g.: 'p_v_var1_bin_0' means that parameter 'p' is plotted as a function of 'var1' in bin 0 of 'var2' """ #auxiliary definitions bin_var_X = self.bin_vars[0] bin_var_Y = self.bin_vars[1] #auxiliary dict for getting the parameters in the right order of bins p_name = OrderedDict() p_name[bin_var_X] = ('j', 'i') p_name[bin_var_Y] = ('i', 'j') #all bin means calculated previously bin_means = OrderedDict() for bin_var in self.bin_vars: for bin_name in self.bins: mean_name = '__'.join([bin_var, bin_name]) bin_means[mean_name] = self.bin_mean(wsp, bin_var, bin_name) graphs = [] for bintovar, bin_var in enumerate(self.bin_vars): #get parameter in the right binning order p_getname = '{param}_{x_var}_{' + p_name[bin_var][ 0] + '}_{y_var}_{' + p_name[bin_var][1] + '}' b_getname = p_getname[8:] for param in simvars: vals = [] #get the values of the parameter for each bin of X,Y for i in xrange(len(self.binning[1 - bintovar]) - 1): vals.append([ get_var( wsp, p_getname.format(param=param, x_var=bin_var_X, y_var=bin_var_Y, i=i, j=j)) for j in xrange(len(self.binning[bintovar]) - 1) ]) #plot graphs as a function of binning var graph = self.plot_free_pars(wsp, bin_var, vals, sym_uncer) #setting the correct mean (errors adjust accordingly), setting graph name for i in xrange(len(self.binning[1 - bintovar]) - 1): g_mean = np.array([ bin_means['__'.join([ bin_var, b_getname.format(x_var=bin_var_X, y_var=bin_var_Y, i=i, j=j) ])] for j in xrange(len(self.binning[bintovar]) - 1) ]) graph[i] = assign_x(graph[i], g_mean) # graph name of the form [y_axis]_v_[x_axis]_bin_[bin of other bin_var] sym_add = '_sym_uncer' if sym_uncer else '' graph[i].SetName('{}_v_{}_bin_{}{}'.format( param, bin_var, i, sym_add)) graphs.append(graph[i]) return graphs
def plot(self, wsp, verbose=False, publication=False): """ Plot the fit results from the passed workspace Args: wsp (ROOT.RooWorkspace): The workspace that holds the model and the data as well as the fit results verbose (boolean): Print some more information about the fit status onto the plots (mainly for debugging) publication (boolean, optional): Make the plots in publication quality (negates the verbose argument automatically) """ data = wsp.data('full_data') cans = OrderedDict() g_chi2, g_ndf = 0, 0 fvar = get_var(wsp, self.fit_var) n_bins = int(round(((fvar.getMax() - fvar.getMin()) / BIN_WIDTH))) wsp.loadSnapshot('snap_two_dim') for bin_name, bin_borders in self.bins.iteritems(): frame, leg = self._plot_bin(wsp, data, bin_name, bin_borders, n_bins, publication) b_chi2, b_ndf = get_chi2_ndf(frame, 'full_pdf_curve', 'data_hist') logging.debug('bin: {}: chi2 / number bins = {:.2f} / {}'.format( bin_name, b_chi2, b_ndf)) g_chi2 += b_chi2 g_ndf += b_ndf can = _setup_canvas(None) # returns a TCanvasWrapper can.cd() pad = r.TPad('result_pad', 'result pad', 0, 0.3, 1, 0.98) r.SetOwnership(pad, False) if publication: pad.SetBottomMargin(0) pad.Draw() pad.cd() frame.Draw() can.add_tobject(pad) leg.Draw() can.add_tobject(leg) # pulls pull_frame = self._pull_plot(wsp, frame, publication) can.cd() pull_pad = r.TPad('pull_pad', 'pull pad', 0, 0, 1, 0.3) r.SetOwnership(pull_pad, False) if publication: pull_pad.SetTopMargin(0) pull_pad.SetBottomMargin(0.275) else: pull_pad.SetBottomMargin(0.2) pull_pad.Draw() pull_pad.SetGridy() pull_pad.cd() pull_frame.Draw() can.add_tobject(pull_pad) cans[bin_name] = can fit_res = wsp.genobj('fit_res_two_dim') n_float_pars = fit_res.floatParsFinal().getSize() g_ndf -= n_float_pars logging.debug('Floating parameters in fit: {}'.format(n_float_pars)) logging.info('Global chi2 / ndf = {:.2f} / {}'.format(g_chi2, g_ndf)) add_info = [] # loop again over all canvases and put the global chi2 / ndf there add_info.append( (0.15, 0.875, '(global) #chi^{{2}} / ndf = {:.1f} / {}'.format(g_chi2, g_ndf))) if verbose and not publication: status = get_var(wsp, '__fit_status__').getVal() cov_qual = get_var(wsp, '__cov_qual__').getVal() add_info.append( (0.15, 0.825, 'status = {}, covQual = {}'.format(int(status), int(cov_qual)))) # Check if the simultaneous negative log-likelihood is present in # the workspace or create it if it is not yet present. # Since the snapshot values of the fit result are already loaded it # can be directly evaluated to get the value at the minimum sim_nll = get_var(wsp, 'sim_nll') if sim_nll is None: sim_nll = self._create_nll( wsp, (rf.Extended(True), rf.Offset(True))) min_nll = sim_nll.getVal() aic = calc_info_crit(fit_res, min_nll) bic = calc_info_crit(fit_res, min_nll, data.numEntries()) add_info.append( (0.15, 0.775, 'AIC = {:.1f}, BIC = {:.1f}'.format(aic, bic))) latex = setup_latex() if publication: latex.SetTextSize(0.04) add_info[0] = (0.185, 0.89, add_info[0][2]) for bin_name, can in cans.iteritems(): res_pad = can.attached_tobjects[0] res_pad.cd() put_on_latex(latex, add_info, ndc=True) bin_info = get_bin_info_text(bin_name, self.bins) put_on_latex(latex, bin_info, ndc=True) # for easier debugging at the moment return cans
def define_model(self, ws): """ Build the chic mass model in the passed workspace Args: ws (ROOT.RooWorkspace): Workspace into which the mass model is constructed """ ## build mass model ws.factory('RooCBShape::{}({}, CBmass1[3.5,3.45,3.54],' 'CBsigma1[0.008, 0.003, 0.02], CBalpha1[0.6, 0.2, 1.1], ' 'CBn[2.5, 1.8, 3.2])'.format(self.chic1, self.mname)) pes = r.RooFormulaVar( 'PES', '(@0 - {0}) / {1}'.format(m_psiPDG, m_chic1PDG - m_psiPDG), r.RooArgList(get_var(ws, 'CBmass1'))) ws_import(ws, pes) CBmass0 = r.RooFormulaVar( 'CBmass0', '(@0 * {0}) + {1}'.format(m_chic0PDG - m_psiPDG, m_psiPDG), r.RooArgList(get_var(ws, 'PES'))) CBmass2 = r.RooFormulaVar( 'CBmass2', '(@0 * {0}) + {1}'.format(m_chic2PDG - m_psiPDG, m_psiPDG), r.RooArgList(get_var(ws, 'PES'))) ws_import(ws, r.RooArgSet(CBmass0, CBmass2)) # ws.factory('CBmass2[3.54, 3.49, 3.58]') # ws.factory('CBmass0[3.42, 3.395, 3.45]') CBsigma0 = r.RooFormulaVar( 'CBsigma0', '@0 * {0}'.format( (m_chic0PDG - m_psiPDG) / (m_chic1PDG - m_psiPDG)), r.RooArgList(get_var(ws, 'CBsigma1'))) CBsigma2 = r.RooFormulaVar( 'CBsigma2', '@0 * {0}'.format( (m_chic2PDG - m_psiPDG) / (m_chic1PDG - m_psiPDG)), r.RooArgList(get_var(ws, 'CBsigma1'))) ws_import(ws, r.RooArgSet(CBsigma0, CBsigma2)) # ws.factory('CBsigma2[0.008, 0.003, 0.02]') # ws.factory('CBsigma0[0.008, 0.003, 0.02]') ws.factory( 'RooCBShape::{}({}, CBmass2, CBsigma2, CBalpha2[0.6, 0.2, 1.1], CBn)' .format(self.chic2, self.mname)) # ws.factory('RooCBShape::M_chic2(chicMass, CBmass2, CBsigma2, CBalpha1, CBn)') CBalpha0 = r.RooFormulaVar( 'CBalpha0', '(@0 + @1) / 2.0', r.RooArgList(get_var(ws, 'CBalpha1'), get_var(ws, 'CBalpha2'))) # r.RooArgList(ws.var('CBalpha1'), ws.var('CBalpha1'))) ws_import(ws, r.RooArgSet(CBalpha0)) ws.factory( 'RooVoigtian::{}({}, CBmass0, CBsigma0, CBwitdh[0.0104])'.format( self.chic0, self.mname)) ## background # ws.factory('q01S[3.1, 3.0, 3.2]') # ws.factory('alpha1[0.6, 0, 5.0]') # ws.factory('beta1[-2.5, -10, 0]') # a1 = r.RooFormulaVar('a1', 'TMath::Abs(@0 - @1)', # r.RooArgList(get_var(ws, 'chicMass'), # get_var(ws, 'q01S'))) # b1 = r.RooFormulaVar('b1', '@0 * (@1 - @2)', # r.RooArgList(get_var(ws, 'beta1'), # get_var(ws, 'chicMass'), # get_var(ws, 'q01S'))) # signum1 = r.RooFormulaVar('signum1', # '(TMath::Sign(-1., @0 - @1) + 1) / 2.', # r.RooArgList(get_var(ws, self.mname), # get_var(ws, 'q01S'))) # ws_import(ws, r.RooArgSet(a1, b1, signum1)) # M_background = r.RooGenericPdf(self.bkg_model, 'signum1 * pow (a1, alpha1) * exp(b1)', # r.RooArgList(ws.function('a1'), ws.function('signum1'), ws.function('b1'), # ws.var('alpha1'))) # polynomial background ws.factory('BK_p1[0, -1, 1]') ws.factory('BK_p2[0, -1, 1]') M_background = r.RooPolynomial( self.bkg_model, self.bkg_model, ws.var('chicMass'), r.RooArgList(ws.var('BK_p1'), ws.var('BK_p2'))) # M_background = r.RooExponential(self.bkg_model, self.bkg_model, ws.var('chicMass'), ws.var('BK_p1')) ws_import(ws, M_background) fix_params( ws, [ ('CBn', 2.75), ('BK_p2', 1e-10), # ('q01S', 3.1) ]) ws.var('CBmass1').setVal(3.510) # ws.var('alpha1').setConstant(True) # ws.var('q01S').setConstant(True) # ws.var('beta1').setConstant(True) ## combine model ws.factory('{}[10000, 0, 200000]'.format( self.nevent_vars[1])) # Nchic1 ws.factory('{}[300, 0, 10000]'.format(self.nevent_vars[0])) # Nchic0 ws.factory('{}[30000, 0, 200000]'.format(self.nevent_vars[-1])) # Nbkg ## leave Nchic2 free and make ratio a formula var ws.factory('{}[10000, 0, 200000]'.format( self.nevent_vars[2])) # Nchic2 r_c2_c1 = r.RooFormulaVar( 'r_chic2_chic1', '@0 / @1', r.RooArgList(get_var(ws, self.nevent_vars[2]), get_var(ws, self.nevent_vars[1]))) ws_import(ws, r_c2_c1) ## make ratio free and Nchic2 depend on it # ws.factory('r_chic2_chic1[0.5, 0, 1]') # Nchic2 = r.RooFormulaVar(self.nevent_vars[2], '@0 * @1', # r.RooArgList(get_var(ws, 'r_chic2_chic1'), # get_var(ws, self.nevent_vars[1]))) # ws_import(ws, Nchic2) Nsignal = r.RooFormulaVar( 'Nsignal', '@0 + @1 + @2', r.RooArgList(get_var(ws, self.nevent_vars[1]), get_var(ws, self.nevent_vars[2]), get_var(ws, self.nevent_vars[0]))) ws_import(ws, Nsignal) ws.factory('SUM::{}({} * {}, {} * {}, {} * {})'.format( self.signal, self.nevent_vars[1], self.chic1, self.nevent_vars[2], self.chic2, self.nevent_vars[0], self.chic0)) ws.factory('SUM::{}(Nsignal * {}, {} * {})'.format( self.full_model, self.signal, self.nevent_vars[-1], self.bkg_model))
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)