def plot_fit(suffix=None, wsp=None):
from . import roofit_to_matplotlib
from . import fit_config
shapes.load_shape_class('RooCruijff')
shapes.load_shape_class('RooJohnsonSU')
shapes.load_shape_class('RooBackground')
mode = gcm()
if wsp is None:
wsp = fit_config.load_workspace(mode)
sel = selection.get_final_selection()
df = mode.get_data([dtf_dm(), m(mode.D0)])
df = df[sel]
data = fit_config.pandas_to_roodataset(df, wsp.set('datavars'))
fit_config.WS_DMASS_NAME = dtf_dm()
fit_config.WS_MASS_NAME = m(mode.D0)
outfile = mode.get_output_path('sweight_fit') + 'fits{}.pdf'.format(
suffix if suffix is not None else '')
with PdfPages(outfile) as pdf:
for func in [m, dtf_dm]:
roofit_to_matplotlib.plot_fit(
mode.D0, wsp, func, data=data, pdf=pdf,
do_comb_bkg=mode.mode in config.twotag_modes)
roofit_to_matplotlib.plot_fit(
mode.D0, wsp, func, data=data, pdf=pdf, do_pulls=False,
do_comb_bkg=mode.mode in config.twotag_modes)
def mass_fiducial_selection(df):
ret = True
ret &= (df[m(gcm().D0)] >= 1810.)
ret &= (df[m(gcm().D0)] < 1920.)
ret &= (df[dtf_dm()] >= 140.5)
ret &= (df[dtf_dm()] < 160.5)
return ret
def rand_spi_sideband_region(df):
"""Selects the signal D0 peak and delta mass sidebands to get a random
slow pion enriched sample"""
ret = True
ret &= np.abs(df[m(gcm().D0)] - config.PDG_MASSES['D0']) < 18.
ret &= np.abs(df[dtf_dm()] - config.PDG_MASSES['delta']) > 2.3
return ret
def mass_signal_region(df):
"""Selects the signal peak in both D0 and delta mass to create a signal
enriched sample."""
ret = True
ret &= np.abs(df[m(gcm().D0)] - config.PDG_MASSES['D0']) < 18.
ret &= np.abs(df[dtf_dm()] - config.PDG_MASSES['delta']) < 0.5
return ret
def fit():
"""Runs the mass fit. Either nominal with making pretty plots or
in spearmint mode which does not save the workspace and returns a
metric."""
# Get the data
# TODO: rewrite selection to use gcm itself
mode = gcm()
sel = selection.get_final_selection()
df = mode.get_data([dtf_dm(), m(mode.D0)])
df = df[sel]
from . import fit_config
from ROOT import RooFit as RF
from .fit_setup import setup_workspace
wsp, _ = setup_workspace()
data = fit_config.pandas_to_roodataset(df, wsp.set('datavars'))
model = wsp.pdf('total')
plot_fit('_start_values', wsp=wsp)
result = model.fitTo(data, RF.NumCPU(4), RF.Save(True), RF.Strategy(2),
RF.Extended(True))
if not helpers.check_fit_result(result, log):
log.error('Bad fit quality')
fit_config.dump_workspace(mode, wsp)
def setup_workspace():
mode = modes.gcm()
wsp = ROOT.RooWorkspace(mode.mode, mode.mode)
fit_config.WS_DMASS_NAME = dtf_dm()
fit_config.WS_MASS_NAME = m(mode.D0)
wsp.factory('{}[{},{}]'.format(m(mode.D0), 1810., 1920.))
wsp.factory('{}[{},{}]'.format(dtf_dm(), 140.5, 160.5))
wsp.var(dtf_dm()).setRange('plotting', 140.5, 152.5)
wsp.var(m(mode.D0)).setRange('plotting', 1820, 1910)
wsp.defineSet('datavars', '{},{}'.format(dtf_dm(), m(mode.D0)))
vs = setup_pdf(wsp)
return wsp, vs
def overlap_plotting():
df = gcm().get_data([vars.dtf_dm()])
sel = extended_selection.get_complete_selection(True)
sel &= misid_selection.misid_cut()
passed = remove_right_sign_candidates()
outfile = gcm().get_output_path('selection') + 'RS_candidates.pdf'
with PdfPages(outfile) as pdf:
nbins = 50
xmin = min(df[sel][vars.dtf_dm()])
xmax = max(df[sel][vars.dtf_dm()])
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist(df[sel & passed][vars.dtf_dm()],
bins=nbins,
range=(xmin, xmax),
color='#006EB6',
edgecolor='#006EB6',
label='Ghost')
ax.set_xlabel(vars.dtf_dm.latex(with_unit=True))
ax.set_xlim((xmin, xmax))
ax.set_ylabel('Arbitrary units')
pdf.savefig(fig)
plt.clf()
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist(df[sel & ~passed][vars.dtf_dm()],
bins=nbins,
range=(xmin, xmax),
color='#006EB6',
edgecolor='#006EB6',
label='Ghost')
ax.set_xlim((xmin, xmax))
ax.set_xlabel(vars.dtf_dm.latex(with_unit=True))
ax.set_ylabel('Arbitrary units')
pdf.savefig(fig)
plt.clf()
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist(df[sel & passed][vars.dtf_dm()],
bins=nbins,
color='#D3EFFB',
range=(xmin, xmax),
label='Kept',
edgecolor='#D3EFFB')
ax.hist(df[sel & ~passed][vars.dtf_dm()],
bins=nbins,
range=(xmin, xmax),
label='Removed',
color='#006EB6',
edgecolor='#006EB6')
ax.set_xlim((xmin, xmax))
ax.set_xlabel(vars.dtf_dm.latex(with_unit=True))
ax.set_ylabel('Candidates')
ax.legend()
pdf.savefig(fig)
plt.clf()
def remove_right_sign_candidates():
"""Remove wrong sign D0 candidates which are combined and end up
in the signal window in the right sign sample"""
# Get the necessary information from the current mode
year = gcm().year
polarity = gcm().polarity
polarity = gcm().polarity
if gcm().mode not in config.twotag_modes:
rs, ws = 'RS', 'WS'
else:
rs, ws = '2tag_RS', '2tag_WS'
with MODE(polarity, year, rs):
RS = gcm().get_data(
[vars.evt_num(),
vars.run_num(),
vars.dtf_dm(),
vars.pt(gcm().D0)])
rs_sel = extended_selection.get_complete_selection(True)
# RS modes should not be selected using this:
if gcm().mode not in config.wrong_sign_modes:
return pd.Series(True, RS.index)
with MODE(polarity, year, ws):
WS = gcm().get_data([
vars.evt_num(),
vars.dtf_dm(),
vars.pt(gcm().D0),
vars.dtf_chi2(gcm().head)
])
OL = RS[rs_sel].merge(WS,
on=['eventNumber'],
left_index=True,
suffixes=['_RS', '_WS'])
dm_ref = config.PDG_MASSES['delta']
OLS = OL.query('(abs(delta_m_dtf_RS-{})<1.) &'
'(abs(D0_PT_RS-D0_PT_WS)<1.)'.format(dm_ref))
return pd.Series(~WS.index.isin(OLS.index), index=WS.index)
def run_spearmint_fit(spearmint_selection=None, metric='punzi'):
"""Runs the mass fit. Either nominal with making pretty plots or
in spearmint mode which does not save the workspace and returns a
metric."""
from . import fit_config
from ROOT import RooFit as RF
shapes.load_shape_class('RooCruijff')
shapes.load_shape_class('RooJohnsonSU')
shapes.load_shape_class('RooBackground')
mode = gcm()
wsp = fit_config.load_workspace(mode)
sel = selection.get_final_selection()
# Get the data
df = mode.get_data([dtf_dm(), m(mode.D0)])
if spearmint_selection is not None:
sel = sel & spearmint_selection
df = df[sel]
data = fit_config.pandas_to_roodataset(df, wsp.set('datavars'))
model = wsp.pdf('total')
metric = get_metric(metric)(wsp)
if spearmint_selection is not None:
result = model.fitTo(data, RF.NumCPU(4), RF.Save(True), RF.Strategy(2),
RF.Extended(True))
if not helpers.check_fit_result(result, log):
result = model.fitTo(data, RF.NumCPU(4), RF.Save(True),
RF.Strategy(1), RF.Extended(True))
if not helpers.check_fit_result(result, log):
result = model.fitTo(data, RF.NumCPU(4), RF.Save(True),
RF.Strategy(0), RF.Extended(True))
if not helpers.check_fit_result(result, log):
log.warn('Bad fit quality')
return 0.0
return metric()
def get_sweights(do_comb_bkg=False):
helpers.allow_root()
df = gcm().get_data([m(gcm().D0), dtf_dm()])
from . import fit_config
from hep_ml import splot
shapes.load_shape_class('RooCruijff')
shapes.load_shape_class('RooJohnsonSU')
shapes.load_shape_class('RooBackground')
wsp = fit_config.load_workspace(gcm())
sel = selection.get_final_selection()
do_comb_bkg = gcm().mode in config.twotag_modes
df = df[sel]
sig_pdf = wsp.pdf('signal')
rnd_pdf = wsp.pdf('random')
comb_pdf = wsp.pdf('combinatorial')
sig_prob = call_after_set(sig_pdf, wsp, **df)
rnd_prob = call_after_set(rnd_pdf, wsp, **df)
if do_comb_bkg:
comb_prob = call_after_set(comb_pdf, wsp, **df)
if do_comb_bkg:
probs = pd.DataFrame(dict(sig=sig_prob*wsp.var('NSig').getVal(),
rnd=rnd_prob*wsp.var('NSPi').getVal(),
comb=comb_prob*wsp.var('NBkg').getVal()),
index=df.index)
else:
probs = pd.DataFrame(dict(sig=sig_prob*wsp.var('NSig').getVal(),
rnd=rnd_prob*wsp.var('NSPi').getVal()),
index=df.index)
probs = probs.div(probs.sum(axis=1), axis=0)
sweights = splot.compute_sweights(probs)
sweights.index = probs.index
if not do_comb_bkg:
sweights['comb'] = 0.0
return sweights
def misid_plots():
"""Remove wrong sign D0 candidates which are combined and end up
in the signal window in the right sign sample"""
# Get the necessary information from the current mode
if gcm().mode in config.wrong_sign_modes:
wrong_spi = add_variables.other_slowpi_ws()
else:
wrong_spi = add_variables.other_slowpi()
dst_mass = gcm().get_data([vars.m(gcm().head)])[vars.m(gcm().head)]
sel = final_selection.get_final_selection()
bins, xmin, xmax = gcm().mass_var.binning
ybins, ymin, ymax = gcm().dmass_var.binning
bins = 30
df_sel = final_selection.get_final_selection()
misid = add_variables.double_misid()
data = gcm().get_data([vars.dtf_dm(), vars.m(gcm().D0)])
outfile = gcm().get_output_path('misid') + 'overview.pdf'
with PdfPages(outfile) as pdf:
for i, pc in enumerate(double_misid_pc):
fig, ax = plt.subplots(figsize=(10, 10))
nbins, xmin, xmax = pc.binning
ax.hist(misid[df_sel][pc.var], bins=nbins, range=(xmin, xmax))
ax.set_xlabel(pc.xlabel)
ax.set_ylabel('Candidates')
ax.set_xlim((xmin, xmax))
pdf.savefig(fig)
plt.close()
if i % 2 == 0:
fig, ax = plt.subplots(figsize=(10, 10))
nbins, xmin, xmax = pc.binning
cutvar = double_misid_pc[i+1].var
narrow = misid[cutvar] < 147.5
ax.hist(misid[df_sel&narrow][pc.var], bins=nbins, range=(xmin, xmax)) # NOQA
ax.set_xlabel(pc.xlabel)
ax.set_ylabel(r'Candidates with $\Delta m <147.5$')
ax.set_xlim((xmin, xmax))
pdf.savefig(fig)
plt.close()
cut = misid_selection.misid_cut()
dm = gcm().dmass_var
nbins, xmin, xmax = dm.binning
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist(data[dm.var][sel & cut], bins=nbins, color='#D3EFFB', # NOQA
range=(xmin, xmax), label='Kept', edgecolor='#D3EFFB')
ax.hist(data[dm.var][sel & ~cut], bins=nbins,
range=(xmin, xmax), label='Removed', color='#006EB6', edgecolor='#006EB6') # NOQA
ax.set_xlim((xmin, xmax))
ax.set_xlabel(dm.xlabel)
ax.set_ylabel('Candidates')
ax.legend()
pdf.savefig(fig)
plt.clf()
outfile = gcm().get_output_path('misid') + 'wrong_spi.pdf'
pdf = PdfPages(outfile)
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist(wrong_spi[sel], bins=bins, range=(xmin, xmax), normed=True, color='#006EB6', edgecolor='#006EB6') # NOQA
ax.set_xlabel(gcm().mass_var.xlabel)
ax.set_xlim((xmin, xmax))
ax.set_ylabel('Arbitrary units')
pdf.savefig(fig)
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist((dst_mass - wrong_spi)[sel], bins=ybins, range=(ymin, ymax), color='#006EB6', edgecolor='#006EB6') # NOQA
ax.set_xlabel(gcm().dmass_var.xlabel)
ax.set_xlim((xmin, xmax))
pdf.savefig(fig)
plt.clf()
plt.clf()
pdf.close()
def plot_mass_regions():
sel = get_final_selection()
df = gcm().get_data([vars.m(gcm().D0), vars.dtf_dm()])
selected = df[sel]
nbins = 100
name = 'mass_regions'
if config.optimised_selection:
name += '_opt'
if config.candidates_selection:
name += '_cand'
outfile = gcm().get_output_path('selection') + name + '.pdf'
with PdfPages(outfile) as pdf:
fig, ax = plt.subplots(figsize=(10, 10))
# Doing D0 mass first
xmin, xmax = 1810, 1920
# Signal window boundaries
sw_lo = config.PDG_MASSES['D0'] - 18.
sw_hi = config.PDG_MASSES['D0'] + 18.
# Lower sideband boundaries
sb_lo_lo = xmin
sb_lo_hi = config.PDG_MASSES['D0'] - 30.
# Upper sideband boundaries
sb_hi_lo = config.PDG_MASSES['D0'] + 30.
sb_hi_hi = xmax
bkg = np.array([(sb_lo_hi + sb_lo_lo) / 2.,
(sb_hi_hi + sb_hi_lo) / 2.])
bkgw = np.array([(sb_lo_hi - sb_lo_lo), (sb_hi_hi - sb_hi_lo)])
sig = np.array([(sw_lo + sw_hi) / 2.])
sigw = np.array([(sw_hi - sw_lo)])
h_vals, edges = np.histogram(selected[vars.m(gcm().D0)],
bins=nbins,
range=(xmin, xmax))
h_errorbars = np.sqrt(h_vals)
x_ctr = (edges[1:] + edges[:-1]) / 2.
width = (edges[1:] - edges[:-1])
x_err = width / 2.
dt_options = dict(fmt='o',
markersize=5,
capthick=1,
capsize=0,
elinewidth=2,
color='#000000',
markeredgecolor='#000000')
ax.errorbar(x_ctr, h_vals, xerr=x_err, yerr=h_errorbars, **dt_options)
hmax = np.max(ax.lines[0].get_ydata())
ax.bar(sig,
1.10 * np.array(hmax),
sigw,
color='#D3EFFB',
edgecolor='#D3EFFB',
label='Signal',
alpha=0.5)
ax.bar(bkg,
1.10 * np.ones(len(bkg)) * hmax,
bkgw,
label='Background',
color='#006EB6',
edgecolor='#006EB6',
alpha=0.5)
ax.set_xlabel(vars.m.latex((gcm().D0), with_unit=True))
unit = r'{} {}'.format((xmax - xmin) / nbins, vars.m.unit)
ylabel = r'Candidates / ({0})'.format(unit)
ax.set_ylabel(ylabel)
ax.legend()
ax.set_xlim(xmin, 0.9999 * xmax)
plot_utils.y_margin_scaler(ax, lf=0, la=True)
pdf.savefig(fig)
plt.clf()
# Now delta mass
fig, ax = plt.subplots(figsize=(10, 10))
xmin, xmax = 140.5, 152.5
# Signal window boundaries
sw_lo = config.PDG_MASSES['delta'] - 0.5
sw_hi = config.PDG_MASSES['delta'] + 0.5
# Lower sideband boundaries
sb_lo_lo = xmin
sb_lo_hi = config.PDG_MASSES['delta'] - 2.3
# Upper sideband boundaries
sb_hi_lo = config.PDG_MASSES['delta'] + 2.3
sb_hi_hi = xmax
bkg = np.array([(sb_lo_hi + sb_lo_lo) / 2.,
(sb_hi_hi + sb_hi_lo) / 2.])
bkgw = np.array([(sb_lo_hi - sb_lo_lo), (sb_hi_hi - sb_hi_lo)])
sig = np.array([(sw_lo + sw_hi) / 2.])
sigw = np.array([(sw_hi - sw_lo)])
h_vals, edges = np.histogram(selected[vars.dtf_dm()],
bins=nbins,
range=(xmin, xmax))
h_errorbars = np.sqrt(h_vals)
x_ctr = (edges[1:] + edges[:-1]) / 2.
width = (edges[1:] - edges[:-1])
x_err = width / 2.
ax.errorbar(x_ctr, h_vals, xerr=x_err, yerr=h_errorbars, **dt_options)
hmax = np.max(ax.lines[0].get_ydata())
ax.bar(sig,
1.10 * np.array(hmax),
sigw,
color='#D3EFFB',
edgecolor='#D3EFFB',
label='Signal',
alpha=0.5)
ax.bar(bkg,
1.10 * np.ones(len(bkg)) * hmax,
bkgw,
label='Background',
color='#006EB6',
edgecolor='#006EB6',
alpha=0.5)
ax.set_xlabel(vars.dtf_dm.latex(with_unit=True))
unit = r'{} {}'.format((xmax - xmin) / nbins, vars.dtf_dm.unit)
ylabel = r'Candidates / ({0})'.format(unit)
ax.set_ylabel(ylabel)
ax.legend()
ax.set_xlim(xmin, 0.9999 * xmax)
plot_utils.y_margin_scaler(ax, lf=0, la=True)
pdf.savefig(fig)
plt.clf()
def comb_bkg_sideband_region(df):
"""Selects the D0 mass sidebands to create a comb background
enriched sample"""
ret = np.abs(df[m(gcm().D0)] - config.PDG_MASSES['D0']) > 30.
ret &= np.abs(df[dtf_dm()] - config.PDG_MASSES['delta']) > 2.3
return ret
def mass_sideband_region(df):
"""Selects the delta mass sidebands to create a background sample"""
ret = True
ret &= np.abs(df[dtf_dm()] - config.PDG_MASSES['delta']) > 2.3
return ret
def delta_mass_wide_signal_region(df):
"""Selects the signal peak in both D0 and delta mass to create a signal
enriched sample."""
ret = True
ret &= np.abs(df[dtf_dm()] - config.PDG_MASSES['delta']) < 1.0
return ret
def remove_clones():
"""Remove events with cloned tracks. Different treatment for RS and WS
due to naming convention."""
# Get the necessary information from the current mode
ret = None
sel = extended_selection.get_complete_selection(True)
passed = remove_right_sign_candidates()
masses = config.PDG_MASSES
outfile = gcm().get_output_path('selection') + 'clones.pdf'
pdf = PdfPages(outfile)
for A, B in combinations(gcm().head.all_daughters(), 2):
log.info('Checking angle between {} and {}'.format(A.name, B.name))
df = gcm().get_data([
vars.phi(A),
vars.pt(A),
vars.eta(A),
vars.phi(B),
vars.pt(B),
vars.eta(B)
])
angle = compute_delta_angle(df[vars.pt(A)], df[vars.eta(A)],
df[vars.phi(A)], masses[A.pid],
df[vars.pt(B)], df[vars.eta(B)],
df[vars.phi(B)], masses[B.pid])
angle = pd.Series(angle, index=df.index)
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist(angle[sel & passed],
bins=50,
range=(0, 0.01),
color='#006EB6',
edgecolor='#006EB6')
ax.set_xlabel(r'$\angle({},{})$ [rad]'.format(A.title.replace('$', ''),
B.title.replace('$',
'')))
ax.set_ylabel('Candidates')
ax.set_xlim((0, 0.01))
pdf.savefig(fig)
plt.clf()
if ret is None:
ret = (angle > 0.0005)
else:
ret &= (angle > 0.0005)
df = gcm().get_data([vars.dtf_dm()])
xmin = min(df[sel & passed & ~ret][vars.dtf_dm()])
xmax = max(df[sel & passed & ~ret][vars.dtf_dm()])
fig, ax = plt.subplots(figsize=(10, 10))
ax.hist(df[sel & passed & ret][vars.dtf_dm()],
bins=50,
range=(xmin, xmax),
label='Kept')
ax.hist(df[sel & passed & ~ret][vars.dtf_dm()],
bins=50,
range=(xmin, xmax),
label='Removed')
ax.set_xlabel(vars.dtf_dm.latex(with_unit=True))
ax.yaxis.set_visible(False)
ax.legend()
pdf.savefig(fig)
plt.clf()
pdf.close()
return ret