def test_get_var_func(self):
"""Test if using a function works correctly"""
var = mh._get_var(self.test_df, lambda d: d.col2)
pdt.assert_series_equal(var, self.test_df.col2)
var = mh._get_var(self.test_df, lambda d: d.col1**2)
pdt.assert_series_equal(var, self.test_df.col1**2)
def fill_3d(data, frame, weights):
"""3D histograms"""
return hist3d(_get_var(data, 'costh_' + frame),
_get_var(data, 'phi_' + frame),
_get_var(data, *binvar),
hist_sett=hist_sett,
weights=weights)
def test_get_var_np_func(self):
"""Test that passing an additional numpy function works"""
var = mh._get_var(self.test_df, 'col2', np.sqrt)
pdt.assert_series_equal(var, np.sqrt(self.test_df.col2))
var = mh._get_var(self.test_df, 'col1', 'abs')
pdt.assert_series_equal(var, self.test_df.col1.abs())
def test_get_var_string(self):
"""Test that a string or a list of string gets handled correctly"""
var = mh._get_var(self.test_df, 'col2')
pdt.assert_series_equal(var, self.test_df.col2)
var = mh._get_var(self.test_df, ['col1', 'col2'])
pdt.assert_frame_equal(var, self.test_df)
def make_2D_hist(dfr, varx, vary, selection=None, weight=None, **kwargs):
"""
Make the 2D histogram and return it
"""
logging.debug('Making histograms for {}:{}'.format(varx, vary))
dfr, weights = get_df_and_weights(dfr, selection, weight)
return hist2d(_get_var(dfr, varx),
_get_var(dfr, vary),
weights=weights,
**kwargs)
def make_costh_phi_plot(data, frame):
"""Make costh-phi histogram"""
hist = hist2d(_get_var(data, 'costh_{}_fold'.format(frame), np.abs),
_get_var(data, 'phi_{}_fold'.format(frame)),
nbinsx=24, minx=0, maxx=1,
nbinsy=32, miny=0, maxy=90)
cth_lab = '|cos#vartheta^{{{}}}|'.format(frame)
phi_lab = '#varphi^{{{}}}_{{fold}}'.format(frame)
can = mkplot(hist, drawOpt='colz', xRange=[0, 1], xLabel=cth_lab,
yRange=[0, 90], yLabel=phi_lab)
can.pltables[0].SetNdivisions(505, 'X')
can.Update()
return can
def get_var_shifted(dfr, func, name, kwargs):
"""
Hacky way of removing the shift for priors but letting it there for
posteriors.
TODO: Remove this once the analysis is unblinded
"""
for var in ['lth2', 'dlth', 'ltilde2', 'dltilde']:
if var in name and 'prior' in name:
low_bound = kwargs['min'] - RAND_DLTH_SHIFT
high_bound = kwargs['max'] - RAND_DLTH_SHIFT
return _get_var(dfr, func) - RAND_DLTH_SHIFT, low_bound, high_bound
return _get_var(dfr, func), kwargs['min'], kwargs['max']
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 get_n_events(data, selections=None, weight=None):
"""
Get the number of events in the dataframe surviving the passed selection
If weight is not None the corresponding column will be used as weights,
unless weight is a function taking the dataframe as only input, than the sum
of the array returned by that function call will be returned
"""
sel_data = apply_selections(data, selections)
if weight is None:
return sel_data.shape[0]
weights = _get_var(sel_data, weight)
return weights.sum()
def fill_2d(data, frame, weights):
"""2D histograms"""
return hist2d(_get_var(data, 'costh_' + frame),
_get_var(data, 'phi_' + frame),
hist_sett=(n_costh, -1, 1, n_phi, -180, 180),
weights=weights)
def corr_w(data):
corr = corrmap.eval(_get_var(data, 'costh_' + frame),
_get_var(data, 'phi_' + frame),
_get_var(data, *parse_func_var(var)))
print('Percentage in acceptance: ', np.sum(corr > 0) / float(len(corr)))
return (corr > 0) * corr
def test_get_var_var(self):
"""Test that when already passed a variable the same is returned"""
var = mh._get_var(self.test_df, self.test_df.col1)
pdt.assert_series_equal(var, self.test_df.col1)
def __call__(self, dfr):
phot_name = get_gen_name('photon', self.gen)
return pt_eta_sel(_get_var(dfr, phot_name + 'Pt'),
_get_var(dfr, phot_name + 'Eta').abs(),
self.cuts)
def get_df_and_weights(dfr, selection=None, weight=None):
"""
Get the dataframe after selections and the according weights
"""
sdfr = apply_selections(dfr, selection)
return sdfr, _get_var(sdfr, weight)
def jpsi_pt_M_sel(dfr, min_ptM=2.3, max_ptM=5.2, gen=False):
"""Kinematic J/psi selection using a Jpsi pT / chicMass selection"""
jpsi_pt = sf.get_gen_name('JpsiPt', gen)
chic_mass = sf.get_gen_name('chicMass', gen)
pt_M = lambda d: _get_var(d, jpsi_pt) / _get_var(d, chic_mass)
return get_bin_cut_df(dfr, pt_M, min_ptM, max_ptM)
