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)