def phsp_variables(df):
"""Returns m12, m34, cos1, cos2, phi1"""
mode = gcm()
# implementation using pybind11::array requires some special treatment
# here, otherwise the passed arrays are of non-matching type.
if not is_dummy_run(df):
vals = vec_phsp_variables(
df[vars.dtf_pt(mode.Pi_OS1)], df[vars.dtf_eta(mode.Pi_OS1)],
df[vars.dtf_phi(mode.Pi_OS1)], config.PDG_MASSES['Pi'],
df[vars.dtf_pt(mode.Pi_SS)], df[vars.dtf_eta(mode.Pi_SS)],
df[vars.dtf_phi(mode.Pi_SS)], config.PDG_MASSES['Pi'],
df[vars.dtf_pt(mode.K)], df[vars.dtf_eta(mode.K)],
df[vars.dtf_phi(mode.K)], config.PDG_MASSES['K'],
df[vars.dtf_pt(mode.Pi_OS2)], df[vars.dtf_eta(mode.Pi_OS2)],
df[vars.dtf_phi(mode.Pi_OS2)], config.PDG_MASSES['Pi'])
return pd.DataFrame(
{
'm12': vals[0],
'm34': vals[1],
'cos1': vals[2],
'cos2': vals[3],
'phi1': vals[4]
},
index=df.index)
else:
vals = (df[vars.dtf_pt(mode.K)], df[vars.dtf_eta(mode.K)],
df[vars.dtf_phi(mode.K)], config.PDG_MASSES['K'],
df[vars.dtf_pt(mode.Pi_OS1)], df[vars.dtf_eta(mode.Pi_OS1)],
df[vars.dtf_phi(mode.Pi_OS1)], config.PDG_MASSES['Pi'],
df[vars.dtf_pt(mode.Pi_SS)], df[vars.dtf_eta(mode.Pi_SS)],
df[vars.dtf_phi(mode.Pi_SS)], config.PDG_MASSES['Pi'],
df[vars.dtf_pt(mode.Pi_OS2)], df[vars.dtf_eta(mode.Pi_OS2)],
df[vars.dtf_phi(mode.Pi_OS2)], config.PDG_MASSES['Pi'])
return 1.
def _ltime_ratio(df):
mode = gcm()
ret = df[vars.ltime(mode.D0)] / config.Dz_ltime
if is_dummy_run(df):
return 1
return pd.Series(ret, name='ltime_ratio', index=df.index)
def _dstp_slowpi_angle(df):
mode = gcm()
ret = compute_delta_angle(
df[vars.pt(mode.D0)],
df[vars.eta(mode.D0)],
df[vars.phi(mode.D0)],
df[vars.m(mode.D0)],
df[vars.pt(mode.Pislow)],
df[vars.eta(mode.Pislow)],
df[vars.phi(mode.Pislow)],
config.PDG_MASSES[config.pion],
)
if is_dummy_run(df):
return 1
return pd.Series(ret, name='dstp_slowpi_angle', index=df.index)
def double_misid_d0(df):
"""Returns d0 mass with changed kaon and ss pion mass hypthesis"""
mode = gcm()
val = double_misid_d0_mass(
df[vars.dtf_pt(mode.K)], df[vars.dtf_eta(mode.K)],
df[vars.dtf_phi(mode.K)], config.PDG_MASSES['Pi'],
df[vars.dtf_pt(mode.Pi_SS)], df[vars.dtf_eta(mode.Pi_SS)],
df[vars.dtf_phi(mode.Pi_SS)], config.PDG_MASSES['K'],
df[vars.dtf_pt(mode.Pi_OS1)], df[vars.dtf_eta(mode.Pi_OS1)],
df[vars.dtf_phi(mode.Pi_OS1)], config.PDG_MASSES['Pi'],
df[vars.dtf_pt(mode.Pi_OS2)], df[vars.dtf_eta(mode.Pi_OS2)],
df[vars.dtf_phi(mode.Pi_OS2)], config.PDG_MASSES['Pi'])
if not is_dummy_run(df):
return pd.Series(val, name=vars.m(gcm().D0), index=df.index)
return 1
def get_named_bdt_discriminant(df, name='KnnFlatness', comb_bkg=False):
# Trigger the loading of the needed objects
if selective_load.is_dummy_run(df):
[
df[f.functor(f.particle)] for f in gcm().bdt_vars
if f.functor != vars.angle
] # NOQA
return 1
log.info('Getting discriminant {} for {}'.format(
name, 'comb. bkg' if comb_bkg else 'rand. pion bkg.'))
if comb_bkg:
features = [f.functor(f.particle) for f in gcm().comb_bkg_bdt_vars]
bdt_vars = gcm().comb_bkg_bdt_vars
else:
features = [f.functor(f.particle) for f in gcm().rand_spi_bdt_vars]
bdt_vars = gcm().rand_spi_bdt_vars
log.info('Features: {}'.format(' '.join(features)))
if vars.angle() in features:
log.info('Adding angle.')
add_variables.append_angle(df)
for f in bdt_vars:
if f.convert is not None:
log.info('Converting {}'.format(f.var))
df[f.var] = f.convert(df[f.var])
df = df[features]
classifiers = bdt_utils.load_classifiers(comb_bkg=comb_bkg)
assert False not in (features == df.columns), 'Mismatching feature order'
bdt = classifiers[name]
probs = bdt.predict_proba(df).transpose()[1]
log.info('Returning probability.')
return pd.Series(probs, index=df.index)
def double_misid(df):
"""Constructs dataframe of all double misid masses to play about"""
mode = gcm()
if is_dummy_run(df):
valSS = (df[vars.pt(mode.K)], df[vars.eta(mode.K)], df[vars.phi(
mode.K)], config.PDG_MASSES['Pi'], df[vars.pt(mode.Pi_SS)],
df[vars.eta(mode.Pi_SS)], df[vars.phi(mode.Pi_SS)],
config.PDG_MASSES['K'], df[vars.pt(mode.Pi_OS1)],
df[vars.eta(mode.Pi_OS1)], df[vars.phi(mode.Pi_OS1)],
config.PDG_MASSES['Pi'], df[vars.pt(mode.Pi_OS2)],
df[vars.eta(mode.Pi_OS2)], df[vars.phi(mode.Pi_OS2)],
config.PDG_MASSES['Pi'], df[vars.pt(mode.Pislow)],
df[vars.eta(mode.Pislow)], df[vars.phi(
mode.Pislow)], config.PDG_MASSES['Pi'])
return 1.
# First some sortingbdtdata
os1bigger = df[vars.pt(mode.Pi_OS1)] > df[vars.pt(mode.Pi_OS2)]
os2bigger = df[vars.pt(mode.Pi_OS1)] <= df[vars.pt(mode.Pi_OS2)]
df.loc[os1bigger, 'H_PT'] = df[vars.pt(mode.Pi_OS1)]
df.loc[os1bigger, 'L_PT'] = df[vars.pt(mode.Pi_OS2)]
df.loc[os1bigger, 'H_ETA'] = df[vars.eta(mode.Pi_OS1)]
df.loc[os1bigger, 'L_ETA'] = df[vars.eta(mode.Pi_OS2)]
df.loc[os1bigger, 'H_PHI'] = df[vars.phi(mode.Pi_OS1)]
df.loc[os1bigger, 'L_PHI'] = df[vars.phi(mode.Pi_OS2)]
df.loc[os2bigger, 'H_PT'] = df[vars.pt(mode.Pi_OS2)]
df.loc[os2bigger, 'L_PT'] = df[vars.pt(mode.Pi_OS1)]
df.loc[os2bigger, 'H_ETA'] = df[vars.eta(mode.Pi_OS2)]
df.loc[os2bigger, 'L_ETA'] = df[vars.eta(mode.Pi_OS1)]
df.loc[os2bigger, 'H_PHI'] = df[vars.phi(mode.Pi_OS2)]
df.loc[os2bigger, 'L_PHI'] = df[vars.phi(mode.Pi_OS1)]
# Correct assignment as cross-check
valC = vec_compute_four_delta_mass(
df[vars.pt(mode.K)], df[vars.eta(mode.K)], df[vars.phi(mode.K)],
config.PDG_MASSES['K'], df[vars.pt(mode.Pi_SS)], df[vars.eta(
mode.Pi_SS)], df[vars.phi(mode.Pi_SS)], config.PDG_MASSES['Pi'],
df['H_PT'], df['H_ETA'], df['H_PHI'], config.PDG_MASSES['Pi'],
df['L_PT'], df['L_ETA'], df['L_PHI'], config.PDG_MASSES['Pi'],
df[vars.pt(mode.Pislow)], df[vars.eta(mode.Pislow)],
df[vars.phi(mode.Pislow)], config.PDG_MASSES['Pi'])
# Exchange K <-> SS Pion
valSS = vec_compute_four_delta_mass(
df[vars.pt(mode.K)], df[vars.eta(mode.K)], df[vars.phi(mode.K)],
config.PDG_MASSES['Pi'], df[vars.pt(mode.Pi_SS)], df[vars.eta(
mode.Pi_SS)], df[vars.phi(mode.Pi_SS)], config.PDG_MASSES['K'],
df['H_PT'], df['H_ETA'], df['H_PHI'], config.PDG_MASSES['Pi'],
df['L_PT'], df['L_ETA'], df['L_PHI'], config.PDG_MASSES['Pi'],
df[vars.pt(mode.Pislow)], df[vars.eta(mode.Pislow)],
df[vars.phi(mode.Pislow)], config.PDG_MASSES['Pi'])
# Exchange K <-> OS1 Pion
valOS1 = vec_compute_four_delta_mass(
df[vars.pt(mode.K)], df[vars.eta(mode.K)], df[vars.phi(mode.K)],
config.PDG_MASSES['Pi'], df[vars.pt(mode.Pi_SS)], df[vars.eta(
mode.Pi_SS)], df[vars.phi(mode.Pi_SS)], config.PDG_MASSES['Pi'],
df['H_PT'], df['H_ETA'], df['H_PHI'], config.PDG_MASSES['K'],
df['L_PT'], df['L_ETA'], df['L_PHI'], config.PDG_MASSES['Pi'],
df[vars.pt(mode.Pislow)], df[vars.eta(mode.Pislow)],
df[vars.phi(mode.Pislow)], config.PDG_MASSES['Pi'])
# Exchange K <-> OS2 Pion
valOS2 = vec_compute_four_delta_mass(
df[vars.pt(mode.K)], df[vars.eta(mode.K)], df[vars.phi(mode.K)],
config.PDG_MASSES['Pi'], df[vars.pt(mode.Pi_SS)], df[vars.eta(
mode.Pi_SS)], df[vars.phi(mode.Pi_SS)], config.PDG_MASSES['Pi'],
df['H_PT'], df['H_ETA'], df['H_PHI'], config.PDG_MASSES['Pi'],
df['L_PT'], df['L_ETA'], df['L_PHI'], config.PDG_MASSES['K'],
df[vars.pt(mode.Pislow)], df[vars.eta(mode.Pislow)],
df[vars.phi(mode.Pislow)], config.PDG_MASSES['Pi'])
if not is_dummy_run(df):
return pd.DataFrame(
{
'm_SS': valSS[0],
'dm_SS': valSS[1],
'm_OSH': valOS1[0],
'dm_OSH': valOS1[1],
'm_C': valC[0],
'dm_C': valC[1],
'm_OSL': valOS2[0],
'dm_OSL': valOS2[1]
},
index=df.index)
return 1