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)