def main(args):
"""Main"""
vals_free_b = get_vals_uncers('/'.join([args.resultdir, 'fit_results_free_betas.root']),
get_par_to_idx(0))
vals_default = get_vals_uncers('/'.join([args.resultdir, 'fit_results_equal_total_cs.root']),
get_par_to_idx(1))
vals_equal_b = get_vals_uncers('/'.join([args.resultdir, 'fit_results_equal_betas.root']),
get_par_to_idx(2))
data_default = get_dataframe('/'.join([args.resultdir, 'scan_ptm_5_2500k_equal_total_cs.root']))
data_equal_b = get_dataframe('/'.join([args.resultdir, 'scan_ptm_5_2500k_equal_betas.root']))
data_free_b = get_dataframe('/'.join([args.resultdir, 'scan_ptm_5_2500k_free_betas.root']))
jpsi_dir = lambda d: 1 - (d.r_chic2_jpsi + d.r_chic1_jpsi + d.r_psip_jpsi)
data_default.loc[:, 'r_jpsi_direct'] = jpsi_dir(data_default)
data_equal_b.loc[:, 'r_jpsi_direct'] = jpsi_dir(data_equal_b)
data_free_b.loc[:, 'r_jpsi_direct'] = jpsi_dir(data_free_b)
with open_tex_file(args.outfile, PREAMBLE) as texf:
texf.write(
create_fit_par_table(vals_free_b, vals_default, vals_equal_b)
)
texf.write('\n\n\n')
texf.write(
create_deriv_par_table(data_free_b, data_default, data_equal_b)
)
def samples(toydatadir, scenario, sel=None):
chi1, chi2 = scenario
chi1_file = '/'.join([toydatadir, chi1, 'toy_data.root'])
chi2_file = '/'.join([toydatadir, chi2, 'toy_data.root'])
return (apply_selections(get_dataframe(chi1_file), sel),
apply_selections(get_dataframe(chi2_file), sel))
def main(args):
"""Main"""
mcdata = get_dataframe(args.mcfile, columns=MC_VARIABLES)
chi1data = get_dataframe(args.chic1file, 'tr', columns=TOY_VARIABLES)
chi2data = get_dataframe(args.chic2file, 'tr', columns=TOY_VARIABLES)
mc_ratios = OrderedDict()
toy_ratios = OrderedDict()
for name in SELECTIONS_TOY:
toy_ratios[name] = OrderedDict()
mc_ratios[name] = OrderedDict()
toy_sel = SELECTIONS_TOY[name]
mc_sel = SELECTIONS_MC[name]
for var in VARIABLES:
hist_sett = VARIABLES[var]['sett']
toy_ratios[name][var] = get_ratio(chi1data, chi2data,
VARIABLES[var]['var'], toy_sel,
hist_sett, get_weight_func())
mc_ratios[name][var] = get_ratio_mc(mcdata, VARIABLES[var]['var'],
mc_sel, hist_sett)
set_TDR_style()
for name in SELECTIONS_TOY:
for var in VARIABLES:
leg = setup_legend(*LEGPOS[var])
can = mkplot(mc_ratios[name][var],
attr=[PLOT_ATTRIBUTES[name]['mc']],
drawOpt='E2',
legOpt='F',
xLabel=LABELS[var],
yLabel='#chi_{c2} / #chi_{c1}',
leg=leg,
legEntries=['real mc'])
can = mkplot(toy_ratios[name][var],
attr=[PLOT_ATTRIBUTES[name]['toy']],
drawOpt='E1same',
legOpt='PLE',
can=can,
xLabel=LABELS[var],
yLabel='#chi_{c2} / #chi_{c1}',
leg=leg,
legEntries=['toy mc'])
latex = setup_latex()
can.add_tobject(latex)
put_on_latex(latex, [(0.18, 0.96, name)])
savename = '_'.join(['comp_real_toy', name,
var]).replace(' + ', '_')
can.SaveAs(savename + '.pdf')
def main(args):
"""Main"""
treename = get_treename(args.inputfile)
branchname = args.branchname
while branch_already_present(args.inputfile, treename, branchname):
logging.warning(
'\'{}\' is already present in tree \'{}\' in file {}'.format(
branchname, treename, args.inputfile))
branchname = raw_input('Please enter a different branch name: ')
map_file = r.TFile.Open(args.corrmapfile)
corr_map = map_file.Get(args.name)
if not corr_map:
logging.fatal('Cannot find acceptance map \'{}\' in file {}'.format(
args.name, args.corrmapfile))
sys.exit(1)
acc_prov = AcceptanceCorrectionProvider(corr_map)
variables = ['{costh,phi}_PX']
if acc_prov.dim == 3:
variables.append(args.variable)
data = get_dataframe(args.inputfile, treename, columns=variables)
if acc_prov.dim == 2:
corr_weights = acc_prov.eval(data.costh_PX, data.phi_PX)
else:
corr_weights = acc_prov.eval(data.costh_PX, data.phi_PX,
_get_var(data, args.variable))
add_branch(corr_weights, branchname, args.inputfile, treename)
def main(args):
"""Main"""
# In order to not have to load the whole file first determine the names of
# the branches that are necessary
var_names = [('costh', 'phi')]
if args.genlevel:
logging.info('Also adding generator level folding')
var_names.append(('gen_costh', 'gen_phi'))
frames = args.frames.split(',')
load_variables = ['_'.join(p) for p in product(flatten(var_names), frames)]
for infile in args.inputfiles:
logging.info('Processing file {}'.format(infile))
if not args.treename:
treename = get_treename(infile)
else:
treename = args.treename
df = get_dataframe(infile, treename, columns=load_variables)
for var_pair in var_names:
for frame in frames:
costh_f, phi_f = get_folded_frame(df, frame, *var_pair)
add_branch(costh_f, '_'.join([var_pair[0], frame, 'fold']),
infile, treename)
add_branch(phi_f, '_'.join([var_pair[1], frame, 'fold']),
infile, treename)
def load_data(scanfile, var_x, var_y, tree='log_like_scan'):
"""Load the dataframe containing the scan results and do some cleanup"""
# Only load the necessary variables
load_vars = [var_x, var_y, 'llh']
rfile = r.TFile.Open(scanfile)
if has_branch(rfile.Get(tree), ['goodFit']):
load_vars.append('goodFit')
rfile.Close()
data = get_dataframe(scanfile, treename=tree, columns=load_vars)
n_full = data.shape[0]
logging.info('Loaded %d rows of data', n_full)
# Check if the data come from a fit scan
if 'goodFit' in data.columns:
# remove all rows where the fit did not converge and where there are
# nan values for any of the fitted parameters
data = apply_selections(data, lambda d: d.goodFit > 0)
n_bad_fit = n_full - data.shape[0]
logging.info('Removed %d because of no good fit', n_bad_fit)
return data
def main(args):
"""Main"""
lth, lph, frame = args.lambda_theta, args.lambda_phi, args.frame
weight_name = get_storable_name('wPol_{}_lth_{:.2f}'.format(frame, lth))
if lph is not None:
weight_name = get_storable_name('wPol_{}_lth_{:.2f}_lph_{:.2f}'
.format(frame, lth, lph))
mc_frame = get_dataframe(args.inputfile)
if weight_name in mc_frame.columns:
logging.info('weights already present for this polarization scenario')
sys.exit(1)
gen_prefix = 'gen_' if args.genlevel else ''
if lph is None:
weight = calc_weights(mc_frame[gen_prefix + 'costh_' + frame], lth)
else:
weight = calc_weights_2d(mc_frame[gen_prefix + 'costh_' + frame],
mc_frame[gen_prefix + 'phi_' + frame], lth, lph)
logging.debug('Shape of data frame before adding weights: {}'.
format(mc_frame.shape))
mc_frame[weight_name] = weight
logging.debug('Shape of data frame after adding weights: {}'
.format(mc_frame.shape))
outfile = args.outfile if args.outfile else args.inputfile
store_dataframe(mc_frame, outfile, 'chic_mc_tuple')
def main(args):
"""Main"""
data = get_dataframe(args.datafile)
cmfile = r.TFile.Open(args.corrmapfile)
accmap = get_correction_map(cmfile, not args.no_pt, args.acceptance)
cond_mkdir(args.outdir)
plot_args = {'drawOpt': 'colz'}
if args.plot_arguments is not None:
plot_args.update(parse_plot_args(args.plot_arguments.split(';;')))
if isinstance(accmap, r.TH2):
plot = make_overlay_plot(accmap, data, **plot_args)
plot.SaveAs('{}/corrmap_data_overlay_2d.pdf'.format(args.outdir))
else:
pt_binning = get_binning(accmap, 2)
pt_bins = zip(pt_binning[:-1], pt_binning[1:])
for pt_bin in pt_bins:
pdata = apply_selections(data, select_bin('JpsiPt', *pt_bin))
pmap = get_pt_bin(accmap, 0.5 * np.sum(pt_bin))
plot = make_overlay_plot(pmap, pdata, **plot_args)
plot.SaveAs('{}/corrmap_data_overlay_2d_{}_{}.pdf'.format(
args.outdir, int(pt_bin[0]), int(pt_bin[1])))
def main(args):
"""Main"""
frames = args.frames.split(',')
if args.bin_variable is not None:
bin_var = parse_func_var(args.bin_variable)
if args.binning is None:
logging.error('You have to define a binning for \'{}\' if you want '
'to use it as binning variable'.format(bin_var))
sys.exit(1)
binning = parse_binning(args.binning)
else:
binning = None
bin_var = None
logging.info('Processing gen level file')
gen_hists = create_histograms(get_dataframe(args.genlevelfile), frames,
args.ncosth, args.nphi, bin_var, binning)
logging.info('Processing reco level file')
reco_hists = create_histograms(get_dataframe(args.recolevelfile), frames,
args.ncosth, args.nphi, bin_var, binning,
WEIGHT_F)
logging.debug('Scaling gen level hists by '.format(args.scale_gen))
for hist in gen_hists.values():
hist.Scale(args.scale_gen)
logging.info('calculating acceptance maps')
acc_maps = OrderedDict()
for name, hist in reco_hists.iteritems():
gen_map_n = [n for n in gen_hists if n in name]
if len(gen_map_n) > 1:
logging.warning('Found more than 1 gen level map for \'{}\': {}'
.format(name, gen_map_n))
# Still use just the first, since we should always just have 1
acc_maps[name] = divide(hist, gen_hists[gen_map_n[0]])
logging.debug('storing to output file')
outfile = r.TFile(args.outfile, 'recreate' if args.recreate else 'update')
store_hists(outfile, gen_hists, 'gen_costh_phi', bin_var)
store_hists(outfile, reco_hists, 'reco_costh_phi', bin_var)
store_hists(outfile, acc_maps, 'acc_map_costh_phi', bin_var)
def main():
"""Main"""
set_TDR_style()
r.gStyle.SetPadRightMargin(r.gStyle.GetPadRightMargin() + 0.01)
data = apply_selections(get_dataframe(INFILE,
columns=collect_requirements(SELECTIONS)),
SELECTIONS)
can = make_photon_pt_dist_plot(data)
can.SaveAs(os.path.join(OUTDIR, 'photon_pt_dist.pdf'))
def main(args):
"""Main"""
orig_frame = get_dataframe(args.infile, args.treename)
rand_indices = permutation(arange(orig_frame.shape[0]) % args.nsamples)
base_output = args.infile.replace('.root', '__{}.root')
for i in xrange(args.nsamples):
idx_frame = orig_frame.iloc[rand_indices == i]
out_name = base_output.format(i)
store_dataframe(idx_frame, out_name, args.treename)
def main(args):
"""Main"""
data = get_dataframe(args.scanfile, columns=list(flatten(args.variables)))
outfile = r.TFile.Open(args.outfile, 'recreate')
for varx, vary in args.variables:
corr_graphs = make_pair_correlation(data, varx, vary, args.physical_lambdas)
for cg in corr_graphs:
cg.Write()
outfile.Close()
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 main(args):
"""Main"""
data = get_dataframe(args.infile, args.treename)
pserver = PlotServer(args.outfile, 'update')
selections = get_selections((args.ptmin, args.ptmax))
for frame in args.frames.split(','):
for var in ['phi', 'costh']:
for sel in test_selections:
selection = selections[sel]
histset = create_histos(data, selection, args.genlevel, var,
frame)
store_histos(pserver, histset, args.prefix, sel, var, frame,
args.genlevel)
def get_vals_uncers(datafile, par_to_idx):
"""Get the values and uncertainties of all parameters"""
data = get_dataframe(datafile, 'fit_result')
pars = data.parameters[0]
n_pars = np.max(par_to_idx.values()) + 1
if len(pars) != n_pars:
print('Parameters and par_to_idx do not have the same length!')
return
cov = np.reshape(data.cov_matrix[0], (n_pars, n_pars))
vals = {}
for par, idx in par_to_idx.iteritems():
vals[par] = (pars[idx], np.sqrt(cov[idx, idx]))
return vals
def main(args):
"""Main"""
selections = OrderedDict()
for sel in args.selections.split(','):
if sel not in SELECTIONS:
logging.warn('Cannot produce plots for selection \'{}\', '
'because it is not defined'.format(sel))
else:
selections[sel] = SELECTIONS[sel]
varx, vary = args.variablex, args.variabley
load_vars = compile_load_vars(selections, varx, vary)
def histf(data, sel, weight):
"""Wrap the function to conform to interface"""
vrx = varx
vry = vary
if varx in FUNCVARS:
vrx = FUNCVARS[varx]
if vary in FUNCVARS:
vry = FUNCVARS[vary]
return make_2D_hist(data,
vrx,
vry,
selection=sel,
weight=weight,
nbinsx=args.nbinsx,
nbinsy=args.nbinsy,
minx=args.minx,
maxx=args.maxx,
miny=args.miny,
maxy=args.maxy)
data = get_dataframe(args.inputfile, columns=load_vars)
outfile = r.TFile(args.outfile, 'recreate' if args.recreate else 'update')
if args.costh_int:
hists = create_hists(data, histf, selections)
else:
hists = create_costh_binned_hists(data, histf, selections,
np.linspace(0, 1, 11))
store_hists(outfile, hists, args.basename)
outfile.Write('', r.TObject.kWriteDelete)
outfile.Close()
def main():
"""Main"""
set_TDR_style()
cond_mkdir(OUTDIR)
data = apply_selections(get_dataframe(INFILE), SELECTIONS)
r.gStyle.SetPadRightMargin(0.129)
r.gStyle.SetPadLeftMargin(r.gStyle.GetPadLeftMargin() - 0.007)
can = make_costh_phi_plot(data, 'HX')
can.SaveAs(os.path.join(OUTDIR, 'costh_phi_fold_HX_pt_12_18_all.pdf'))
can = make_costh_phi_plot(data, 'CS')
can.SaveAs(os.path.join(OUTDIR, 'costh_phi_fold_CS_pt_12_18_all.pdf'))
set_TDR_style()
can = make_costh_mu_pt_plot(data)
can.SaveAs(os.path.join(OUTDIR, 'costh_HX_comp_mu_pt_jpsipt_12_18_all.pdf'))
def load_data(filen, model):
"""
Load the raw data making sure that only the data specified in the fitting
range as well as in the range defined by the binning to avoid loading
events that are not used in the fit.
NOTE: This still does not beat a proper pre-selection!
"""
mass_sel = select_bin(model.fit_var, *model.fit_range)
selections = [mass_sel]
for var, bounds in model.get_load_vars():
selections.append(
select_bin(var, *[float(v) for v in bounds.split(',')]))
load_vars = ['{costh,phi}_HX_fold'] + collect_requirements(selections)
return apply_selections(get_dataframe(filen, columns=load_vars),
selections)
def main(args):
"""Main"""
data = get_dataframe(args.scanfile)
# Get a "helper" ptm binning, that is only used in the groupby to easily
# identify the different ptM values
ptm_binning, ptm_vals = get_variable_binning(data.ptM, get_vals=True)
ptm_groupby = lambda d: d.groupby(pd.cut(d.ptM, ptm_binning))
outfile = r.TFile.Open(args.outfile, 'recreate')
for var, trans in args.variables.iteritems():
graph = make_graph(data, ptm_groupby, ptm_vals, var, trans,
args.physical_lambdas)
graph.SetName('_'.join([get_var_name(var, trans.__name__), 'v',
'ptm']))
graph.Write()
outfile.Close()
def main(args):
"""Main"""
toy_data = get_dataframe(args.toydata)
outfile = r.TFile.Open(args.outfile, 'recreate')
# combined efficiencies calculated on smeared variables
# NOTE: photon efficiencies are parametrized in percent
comb_effs = lambda df: df.gamma_eff_sm * 0.01 * \
df.lepP_eff_sm * df.lepN_eff_sm
gen_effs = lambda df: df.gamma_eff * 0.01 * \
df.lepP_eff * df.lepN_eff
## store gen level vars
store_all_plots(outfile, toy_data, SELECTION_SETS['gen'], GEN_VARS, 'gen',
True)
## store smeared vars at generation level
store_all_plots(outfile, toy_data, SELECTION_SETS['gen'], SM_VARS, 'gen',
False)
store_all_plots(outfile, toy_data, SELECTION_SETS['gen'], SPECIAL_VARS,
'gen', False)
## store smeared level variables for different selection sets
for sel in SELECTION_SETS:
if sel == 'gen':
continue # 'gen' already stored above
store_base = sel
for var_set in [SM_VARS, SPECIAL_VARS]:
# store once without efficiencies
store_all_plots(outfile, toy_data, SELECTION_SETS[sel], var_set,
store_base, False)
## ... and once with efficiencies
store_all_plots(outfile, toy_data, SELECTION_SETS[sel], var_set,
store_base + '_eff', False, comb_effs)
# store gen level selected and efficiency weighted histograms
store_all_plots(outfile, toy_data, SELECTION_SETS[sel], SM_VARS,
'gen_' + store_base, True)
store_all_plots(outfile, toy_data, SELECTION_SETS[sel], SM_VARS,
'gen_' + store_base + '_eff', True, gen_effs)
outfile.Write()
def main(args):
"""Main"""
# While loading also disregard the events with JpsiRap > 1.2 which have been
# accidentally stored for parts of the generations
data = get_dataframe(args.toymcfile,
columns=['costh_HX', '*eff_sm'],
where='TMath::Abs(JpsiRap) < 1.2')
# per-event efficiency as product of photon x muon efficiencies
eff = lambda d: d.gamma_eff_sm * 0.01 * d.lepP_eff_sm * d.lepN_eff_sm
costh_binning = read_costh_binning(args.binningfile)
hist = create_histogram(data.costh_HX.abs(),
costh_binning,
weights=eff(data),
name=args.name,
x_axis='|cos#vartheta^{HX}|')
outfile = r.TFile.Open(args.output, 'update')
hist.Write()
outfile.Close()
def do_chic_ratio(args):
"""
Make the chic ratio plot
"""
ffile = r.TFile.Open(args.fitfile)
ws = ffile.Get('ws_mass_fit')
outdir = get_outdir(args.outdir, args.fitfile)
bin_sel_info = get_bin_sel_info(args.pklfile, args.fitfile)
costh_bins = bin_sel_info['costh_bins']
costh_means = bin_sel_info['costh_means']
graph = get_ratio_graph(ws, 'Nchic2', ws, 'Nchic1', costh_bins,
costh_means)
ptmin, ptmax = get_pt_range(bin_sel_info['basic_sel'])
mcdf = get_dataframe(args.mcfile, 'chic_mc_tuple')
mc_sel = (np.abs(mcdf.JpsiRap) < 1.2) & ((mcdf.trigger & 2) == 2) & \
(mcdf.vtxProb > 0.01) & \
(get_bin_cut_df(mcdf, 'chicPt', 0, 990)) & \
(get_bin_cut_df(mcdf, 'JpsiPt', ptmin, ptmax))
unpol = get_mc_ratio_graph(mcdf[mc_sel], mcdf[mc_sel], 'wChic2', 'wChic1',
costh_bins, costh_means, None, None)
nrqcd = get_mc_ratio_graph(mcdf[mc_sel], mcdf[mc_sel], 'wChic2', 'wChic1',
costh_bins, costh_means, -0.3, 0.5)
extreme = get_mc_ratio_graph(mcdf[mc_sel], mcdf[mc_sel], 'wChic2',
'wChic1', costh_bins, costh_means, -0.6, 1)
mc_graphs = {'unpol': unpol, 'nrqcd': nrqcd, 'extreme': extreme}
plot_sett = {
'range': [0, 0, 1, 0.75],
'xtitle': '|cos#theta^{HX}|',
'ytitle': '#chi_{c2} / #chi_{c1}'
}
plot_name = '{}/chic2_chic1_pt{}_{}_nbins{}_costh_HX.pdf'.format(
outdir, ptmin, ptmax, len(costh_bins))
make_plot(graph, mc_graphs, plot_name, plot_sett)
def main(args):
"""Main"""
selections = OrderedDict()
selections['trigger'] = sf.trigger_sel_(args.trigger)
selections['muon'] = get_muon_sel(args.muon)
selections['vtx prob'] = sf.vtx_prob_sel_(args.vtxprob)
selections['photon sel'] = sf.photon_sel_(sf.flat_pt(0.4, 1.5))
selections['jpsi kin sel'] = get_jpsi_sel(args.jpsi)
selections['lifetime cut'] = get_lt_selection(args.mc, args.lifetime)
selections['deta cut (MC only)'] = get_deta_sel(args.deta)
# selections['chis mass cut'] = sf.chic_mass_sel # not strictly necessary from a PS point of view
# To ensure rectangular region in costh-phi
# selections['costh'] = lambda d: d.costh_HX_fold.abs() < 0.625
selections['state_sel'] = get_state_sel(args.mc, args.state)
selections['costh_sel'] = get_costh_sel(args.max_costh)
global VARIABLES
VARIABLES.extend(sf.collect_requirements(selections.values()))
if args.deta:
VARIABLES.append('gen_photonEta')
if args.mc:
VARIABLES.append('pdgId')
VARIABLES = list(set(VARIABLES))
data = get_dataframe(args.infile, columns=VARIABLES, where='trigger > 0')
if not args.hist:
sel_data = apply_selections(data, selections.values())
else:
sel_data, n_evts = apply_selection_one_by_one(data,
selections.values())
hist = create_store_sel_hist(n_evts, selections.keys())
# compute the Q-value based mass
sel_data.loc[:,
'mQ'] = sel_data.mumugammaMass - sel_data.JpsiMass + m_psiPDG
store_dataframe(sel_data, args.outfile, 'chic_tuple')
def main(args):
"""Main"""
dfr = get_dataframe(args.mcfile)
costh_bins = get_costh_binning(dfr, args.nbins)
h_int = create_mass_hist(dfr)
h_int.Scale(1.0 / args.nbins)
h_bins = [create_mass_hist(dfr, cbin) for cbin in costh_bins]
attr = [{'color': 1, 'marker': 20, 'size': 1.0}]
attr += default_attributes(size=1.0, open_markers=True) * 3
legEntries = ['costh integrated']
legEntries += [
'{:.2f} < |cos#vartheta^{{HX}}| < {:.2f}'.format(*cbin)
for cbin in costh_bins
]
leg = setup_legend(0.675, 0.6, 0.95, 0.94)
leg.SetTextSize(0.025)
can = r.TCanvas('plot_can', '', 600, 600)
can.SetRightMargin(0.02)
can.SetTopMargin(0.05)
can.SetLeftMargin(0.12)
can = mkplot([h_int] + h_bins,
can=can,
xRange=[3.325, 3.725],
yRange=[0.1, None],
attr=attr,
logy=args.logy,
xLabel='M^{#chi}',
yLabel='# events',
legOpt='PLE',
leg=leg,
legEntries=legEntries)
can.SaveAs(args.output)
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 main(args):
"""Main"""
set_basic_style()
llh_data = get_dataframe(args.scanfile, 'log_like_scan')
cond_mkdir(args.outdir)
can = make_plot_good_fit(llh_data, args.variable_x, args.variable_y,
args.transform_x, args.transform_y)
can.SaveAs('/'.join([args.outdir, 'scan_good_fit.pdf']))
can = make_plot_min_chi2(llh_data, args.variable_x, args.variable_y,
args.transform_x, args.transform_y)
can.SaveAs('/'.join([args.outdir, 'chi2_vals_all_fits.pdf']))
can = make_plot_min_chi2(llh_data,
args.variable_x,
args.variable_y,
args.transform_x,
args.transform_y,
gf_only=True)
can.SaveAs('/'.join([args.outdir, 'chi2_vals_good_fits.pdf']))
def main(args):
"""Main"""
muon_effs = get_effs(args.efficiencies, args.muoneffs, MuonEfficiencies)
phot_effs = get_effs(args.efficiencies, args.photoneffs,
PhotonEfficiencies)
treename = get_treename(
args.datafile) # do this here and fail early in case
# only read in the necessary columns and add newly created branches to file
columns = [n + '{Pt,Eta}' for n in MUON_NAMES + [PHOTON_NAME]]
data = get_dataframe(args.datafile, treename, columns=columns)
if phot_effs is not None:
photon_effs = np.array(calc_effs(data, phot_effs, PHOTON_NAME))
add_branch(photon_effs, '_'.join(['gamma_sm', args.name]),
args.datafile, treename)
if muon_effs is not None:
for muon in MUON_NAMES:
mu_effs = np.array(calc_effs(data, muon_effs, muon))
lep = 'lepP' if 'P' in muon else 'lepN' # for consistency
add_branch(mu_effs, '_'.join([lep, 'sm', args.name]),
args.datafile, treename)
def do_jpsi_ratios(args):
"""
Make the jpsi ratio plots
"""
chic_ff = r.TFile.Open(args.chic_ff)
chic_ws = chic_ff.Get('ws_mass_fit')
jpsi_ff = r.TFile.Open(args.jpsi_ff)
jpsi_ws = jpsi_ff.Get('ws_mass_fit')
outdir = get_outdir(args.outdir, args.chic_ff)
# Only need the chic information here, since the information in the J/psi
# file that is needed here is redundant
bin_info = get_bin_sel_info(args.pklfile, args.chic_ff)
costh_bins = bin_info['costh_bins']
costh_means = bin_info['costh_means']
chic1_graph = get_ratio_graph(chic_ws, 'Nchic1', jpsi_ws, 'Njpsi',
costh_bins, costh_means)
chic2_graph = get_ratio_graph(chic_ws, 'Nchic2', jpsi_ws, 'Njpsi',
costh_bins, costh_means)
ptmin, ptmax = get_pt_range(bin_info['basic_sel'])
chic_df = get_dataframe(args.mc_chic, 'chic_mc_tuple')
chic_sel_mc = (np.abs(chic_df.JpsiRap) < 1.2) & \
((chic_df.trigger & 2) == 2) & (chic_df.vtxProb > 0.01) & \
(get_bin_cut_df(chic_df, 'chicPt', 0, 990)) & \
(get_bin_cut_df(chic_df, 'JpsiPt', ptmin, ptmax))
jpsi_df = get_dataframe(args.mc_jpsi, 'jpsi_mc_tuple')
jpsi_sel_mc = (np.abs(jpsi_df.JpsiRap) < 1.2) & (jpsi_df.trigger == 1) & \
(jpsi_df.vtxProb > 0.01) & \
(get_bin_cut_df(jpsi_df, 'JpsiPt', ptmin, ptmax))
unpol = get_mc_ratio_graph(chic_df[chic_sel_mc], jpsi_df[jpsi_sel_mc],
'wChic1', 'wJpsi', costh_bins, costh_means,
None, None)
nrqcd = get_mc_ratio_graph(chic_df[chic_sel_mc], jpsi_df[jpsi_sel_mc],
'wChic1', 'wJpsi', costh_bins, costh_means, 0.5,
None)
extreme = get_mc_ratio_graph(chic_df[chic_sel_mc], jpsi_df[jpsi_sel_mc],
'wChic1', 'wJpsi', costh_bins, costh_means, 1,
None)
can = r.TCanvas('rcan', 'rcan', 50, 50, 600, 600)
frame = can.DrawFrame(0, 0, 1, 0.1)
mc_graphs = {'unpol': unpol, 'nrqcd': nrqcd, 'extreme': extreme}
plot_sett = {
'range': [0, 0, 1, 0.125],
'xtitle': '|cos#theta^{HX}|',
'ytitle': '#chi_{c1} / J/#psi'
}
plot_name = '{}/chic1_jpsi_pt{}_{}_nbins{}_costh_HX.pdf'.format(
outdir, ptmin, ptmax, len(costh_bins))
make_plot(chic1_graph, mc_graphs, plot_name, plot_sett)
unpol = get_mc_ratio_graph(chic_df[chic_sel_mc], jpsi_df[jpsi_sel_mc],
'wChic2', 'wJpsi', costh_bins, costh_means,
None, None)
nrqcd = get_mc_ratio_graph(chic_df[chic_sel_mc], jpsi_df[jpsi_sel_mc],
'wChic2', 'wJpsi', costh_bins, costh_means,
-0.3, None)
extreme = get_mc_ratio_graph(chic_df[chic_sel_mc], jpsi_df[jpsi_sel_mc],
'wChic2', 'wJpsi', costh_bins, costh_means,
-0.6, None)
mc_graphs = {'unpol': unpol, 'nrqcd': nrqcd, 'extreme': extreme}
plot_sett['range'] = [0, 0, 1, 0.075]
plot_sett['ytitle'] = '#chi_{c2} / J/#psi'
plot_name = plot_name.replace('chic1', 'chic2')
make_plot(chic2_graph, mc_graphs, plot_name, plot_sett)
Get the chic2 / chic1 ratio for a given variable in a given frame
"""
chi1_data = apply_selections(df, lambda df: chic_state_sel(df, 'chic1'))
chi2_data = apply_selections(df, lambda df: chic_state_sel(df, 'chic2'))
chi1_hist = get_var_histo(chi1_data, var, frame, gen)
chi2_hist = get_var_histo(chi2_data, var, frame, gen)
chi2_hist.Divide(chi1_hist)
chi2_hist.SetYTitle('#chi_{c2} / #chi_{c1}')
return chi2_hist
gen_data = get_dataframe(
'/afs/hephy.at/data/tmadlener01/ChicPol/Chic2012/InputFiles/NewMCGen2012/flat_tuples/pGunRunI-chic-gen_flat_new_tupling.root'
)
reco_data = get_dataframe(
'/afs/hephy.at/data/tmadlener01/ChicPol/Chic2012/InputFiles/NewMCGen2012/flat_tuples/pGunRunI-chic-lowPt_merged_new_tupling.root'
)
gen_selection = (lambda df: photon_sel(df, flat_pt(0, 1.5), gen=True),
lambda df: fiducial_muon_sel(df, gen=True),
lambda df: jpsi_kin_sel(df, 8, 20, 1.2, gen=True))
# reco_selection = gen_selection + (trigger_sel,)
reco_selection = gen_selection
print('Number of events after selection: gen: {}, reco: {}'.format(
get_n_events(gen_data, gen_selection),
get_n_events(reco_data, reco_selection)))
# save_1D_histograms(chi2_hists, outdir, 'chic2', drawOpt='E1',
# yRange=[0, None], yLabel='N^{#chi_{c2}}')
# can_chi1 = save_1D_histograms(chi1_hists, outdir, 'jpsipt', drawOpt='E1',
# yRange=[0, None])
# can_chi2 = save_1D_histograms(chi2_hists, outdir, 'jpsipt', drawOpt='E1',
# yRange=[0, None])
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser(description='Script for checking effects '
'of single muon acceptance cuts on costh '
'distribution at gen level')
parser.add_argument('genfile', help='File containing the flat genlevel tuple')
parser.add_argument('--ptmin', type=float, default=8, help='minimum jpsi pt')
parser.add_argument('--ptmax', type=float, default=20, help='maximum jpsi pt')
parser.add_argument('-t', '--treename', default='chic_mc_tuple', type=str,
help='name of the tree in which the original variables are '
'(used for storing output file).')
parser.add_argument('-f', '--frame', help='reference frame', default='HX')
parser.add_argument('-o', '--outdir', help='output directory', default='.')
args = parser.parse_args()
data = get_dataframe(args.genfile, args.treename)
make_comp_plots(data, args.frame, args.outdir, (args.ptmin, args.ptmax))
