def get_fjparticles(self, df_particles_grouped): user_index_offset = 0 return fjext.vectorize_pt_eta_phi(df_particles_grouped['pt'].values, df_particles_grouped['y'].values, df_particles_grouped['phi'].values, user_index_offset)
def process_event(self, df): self.pbar.update(1) if 'ev_id_ext' in list(self.event_df): _ev_query = "run_number == {} & ev_id == {} & ev_id_ext == {}".format(df['run_number'], df['ev_id'], df['ev_id_ext']) else: _ev_query = "run_number == {} & ev_id == {}".format(df['run_number'], df['ev_id']) _df_d0 = self.d0_df.query(_ev_query) _df_d0.reset_index(drop=True) _n_d0s = len(_df_d0.index) if _n_d0s < 1: return _ev_query = "" _df_tracks = self.track_df.query(_ev_query) _df_tracks.reset_index(drop=True) _parts = fjext.vectorize_pt_eta_phi(_df_tracks['ParticlePt'].values, _df_tracks['ParticleEta'].values, _df_tracks['ParticlePhi'].values) self._user_index_offset = 10000 # _d0s = fjext.vectorize_pt_eta_phi(_df_d0['pt_cand'].values, _df_d0['eta_cand'].values, _df_d0['phi_cand'].values, self._user_index_offset) _d0s = fjext.vectorize_pt_eta_phi_m(_df_d0['pt_cand'].values, _df_d0['eta_cand'].values, _df_d0['phi_cand'].values, _df_d0['inv_mass'].values, self._user_index_offset) _d0s_gh = [p * 1.e-6 for p in _d0s] _parts_and_ds = _parts _tmp = [_parts_and_ds.push_back(p) for p in _d0s_gh] # pinfo('n parts = ', len(_parts_and_ds)) ja = jet_analysis.JetAnalysis(jet_R = 0.6, particle_eta_max=0.9, jet_pt_min=5.0) ja.analyze_event(_parts_and_ds) _d0_imass_list = _df_d0['inv_mass'].values.tolist() self.tw.fill_branches(dpsj = _d0s) self.tw.fill_branches(dpsjgh = _d0s_gh) self.tw.fill_branches(minv = _d0_imass_list) self.tw.fill_branches(jets = ja.jets_as_psj_vector()) self.tw.fill_tree() self.d0_jet_correl(ja.jets, _d0s, _d0_imass_list)
def fj_parts_from_tracks_numpy(tracks): fjparts = [] _pts = tracks['ParticlePt'] _etas = tracks['ParticleEta'] _phis = tracks['ParticlePhi'] fjparts = fjext.vectorize_pt_eta_phi(_pts.values, _etas.values, _phis.values) return fjparts
def get_event(self, df_tracks): # Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi event = DataEvent([], -1, -1) event.particles = fjext.vectorize_pt_eta_phi( df_tracks['ParticlePt'].values, df_tracks['ParticleEta'].values, df_tracks['ParticlePhi'].values) if len(event.particles) > 0: event.run_number = float(df_tracks['run_number'].values[0]) event.ev_id = float(df_tracks['ev_id'].values[0]) else: event.run_number = -1 event.ev_id = -1 return event
def get_fjparticles(self, df_tracks, offset_indices=False): # If offset_indices is true, then offset the user_index by a large negative value user_index_offset = 0 if offset_indices: user_index_offset = int(-1e6) # Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi fj_particles = fjext.vectorize_pt_eta_phi( df_tracks['ParticlePt'].values, df_tracks['ParticleEta'].values, df_tracks['ParticlePhi'].values, user_index_offset) return fj_particles
def load_event(self): # Decide how many tracks to generate N_tracks = int(np.random.normal(self.N_avg, self.sigma_N)) # Generate tracks, and populate a dataframe pt_array = np.random.gamma(self.alpha, self.beta, N_tracks) eta_array = np.random.uniform(-self.eta_max, self.eta_max, N_tracks) phi_array = np.random.uniform(0., 2 * np.pi, N_tracks) # Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi user_index_offset = int(-1e6) fj_particles = fjext.vectorize_pt_eta_phi(pt_array, eta_array, phi_array, user_index_offset) return fj_particles
def get_event(self, df_tracks): # event_df = self.event_df.loc[(self.event_df['run_number'] == self.run_number) & (self.event_df['ev_id'] == self.ev_id)] # event = DataEvent([], -1, -1, self.event_df) # event = DataEvent(particles=[], run_number=-1, ev_id=-1) # Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi _parts = fjext.vectorize_pt_eta_phi(df_tracks['ParticlePt'].values, df_tracks['ParticleEta'].values, df_tracks['ParticlePhi'].values) event = DataEvent(particles=_parts) for c in df_tracks.columns: if 'Particle' not in c: event.__setattr__(c, df_tracks[c].values[0]) if len(event.particles) <= 0: # backward compat event.run_number = -1 event.ev_id = -1 return event
def analyze(self, df): self.compile_selection(df) _df = df[self.df_selection] self.fj_parts = fjext.vectorize_pt_eta_phi(_df['ParticlePt'].values, _df['ParticleEta'].values, _df['ParticlePhi'].values) _df = _df.drop(columns=['ParticlePt', 'ParticleEta', 'ParticlePhi']) self.fj_Dcands = fjext.vectorize_pt_eta_phi_m(_df['pt_cand'].values, _df['eta_cand'].values, _df['phi_cand'].values, _df['inv_mass'].values) _Dgh_pt = [1e-5 for v in _df['pt_cand'].values] self.fj_DcandsGhosts = fjext.vectorize_pt_eta_phi_m( _Dgh_pt, _df['eta_cand'].values, _df['phi_cand'].values, _df['inv_mass'].values) _df.reset_index(drop=True) _df.drop_duplicates(inplace=True) self.analysis(_df) if self.callback is not None: self.callback(df['ev_id'].values[0])
def get_event(self, df): # Use swig'd function to create a vector of fastjet::PseudoJets from numpy arrays of pt,eta,phi _parts = fjext.vectorize_pt_eta_phi(df['ParticlePt'].values, df['ParticleEta'].values, df['ParticlePhi'].values) #event = DataEvent(particles=_parts) #for c in df_tracks.columns: # if 'Particle' not in c: # event.__setattr__(c, df_tracks[c].values[0]) #if len(event.particles) <= 0: # # backward compat # event.run_number = -1 # event.ev_id = -1 _event = pd.DataFrame({ "run_number": df["run_number"], "ev_id": df["ev_id"], "ParticleFJ": _parts }) event = pd.merge(df, _event, on=['run_number', 'ev_id']) return event
def main(args): fname = args.fname file = uproot.open(fname) all_ttrees = dict( file.allitems( filterclass=lambda cls: issubclass(cls, uproot.tree.TTreeMethods))) tracks = all_ttrees[b'PWGHF_TreeCreator/tree_Particle;1'] pds_trks = tracks.pandas.df() # entrystop=10) events = all_ttrees[b'PWGHF_TreeCreator/tree_event_char;1'] pds_evs = events.pandas.df() # print the banner first fj.ClusterSequence.print_banner() # signal jet definition maxrap = 0.9 jet_R0 = args.jetR jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) jet_selector = fj.SelectorPtMin(0.0) & fj.SelectorAbsEtaMax(1) jet_area_def = fj.AreaDefinition(fj.active_area, fj.GhostedAreaSpec(maxrap)) print(jet_def) # background estimation grid_spacing = maxrap / 10. gmbge = fj.GridMedianBackgroundEstimator(maxrap, grid_spacing) print() output_columns = ['evid', 'pt', 'eta', 'phi', 'area', 'ptsub'] e_jets = pd.DataFrame(columns=output_columns) for i, e in pds_evs.iterrows(): iev_id = int(e['ev_id']) run_number = int(e['run_number']) _ts = pds_trks.loc[pds_trks['ev_id'] == iev_id].loc[ pds_trks['run_number'] == run_number] _tpsj = fjext.vectorize_pt_eta_phi(_ts['ParticlePt'].values, _ts['ParticleEta'].values, _ts['ParticlePhi'].values) # print('maximum particle rapidity:', max([psj.rap() for psj in _tpsj])) _cs = fj.ClusterSequenceArea(_tpsj, jet_def, jet_area_def) _jets = jet_selector(fj.sorted_by_pt(_cs.inclusive_jets())) gmbge.set_particles(_tpsj) # print("rho = ", gmbge.rho(), "sigma = ", gmbge.sigma()) _jets_df = pd.DataFrame([[ iev_id, j.perp(), j.eta(), j.phi(), j.area(), j.perp() - gmbge.rho() * j.area() ] for j in _jets], columns=output_columns) e_jets = e_jets.append(_jets_df, ignore_index=True) # print('event', i, 'number of parts', len(_tpsj), 'number of jets', len(_jets)) # print(_jets_df.describe()) if args.fjsubtract: fj_example_02_area(_tpsj) # print(e_jets.describe()) joblib.dump(e_jets, args.output)
def main(): parser = argparse.ArgumentParser(description='test the TennGen', prog=os.path.basename(__file__)) pyconf.add_standard_pythia_args(parser) parser.add_argument('--ignore-mycfg', help="ignore some settings hardcoded here", default=False, action='store_true') parser.add_argument('--cent-bin', help="centraility bin 0 is the 0-5 % most central bin", type=int, default=0) parser.add_argument('--seed', help="pr gen seed", type=int, default=1111) parser.add_argument('--harmonics', help="set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only)", type=int, default=5) parser.add_argument('--eta', help="set eta range must be uniform (e.g. |eta| < 0.9, which is ALICE TPC fiducial acceptance)", type=float, default=0.9) parser.add_argument('--qa', help="PrintOutQAHistos", default=False, action='store_true') args = parser.parse_args() args.py_pthatmin = 100 mycfg = ['PhaseSpace:pThatMin = {}'.format(args.py_pthatmin)] if args.ignore_mycfg: mycfg = [] pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if not pythia: perror("pythia initialization failed.") return tgbkg = ROOT.TennGen() # //constructor tgbkg.SetCentralityBin(args.cent_bin) # //centraility bin 0 is the 0-5 % most central bin tgbkg.SetRandomSeed(args.seed) # //setting the seed tgbkg.SetHarmonics(args.harmonics) # // set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only) tgbkg.SetEtaRange(args.eta) # //set eta range must be uniform (e.g. |eta| < 0.9, which is ALICE TPC fiducial acceptance) tgbkg.PrintOutQAHistos(args.qa) # tgbkg.InitializeBackground() # jet_R0 = 0.4 jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) jet_selector_pythia = fj.SelectorPtMin(args.py_pthatmin) & fj.SelectorPtMax(1000.0) &fj.SelectorAbsEtaMax(args.eta - jet_R0) jet_selector_hybrid = fj.SelectorPtMin(10) & fj.SelectorPtMax(1000.0) &fj.SelectorAbsEtaMax(args.eta - jet_R0) # jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(200.0) &fj.SelectorAbsEtaMax(1) parts_selector = fj.SelectorAbsEtaMax(args.eta) print(jet_def) tw = treewriter.RTreeWriter(name = 'tparts', file_name = 'test_TennGen.root') if args.nev < 100: args.nev = 100 pbar = tqdm.tqdm(total = args.nev) while pbar.n < args.nev: if not pythia.next(): continue # get pythia particles # parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False) _py_fj_parts = parts_selector(pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)) # get jets w/o area determination # pythia_jets = jet_selector_pythia(jet_def(_py_fj_parts)) # with area determination jet_area_def = fj.AreaDefinition(fj.active_area, fj.GhostedAreaSpec(args.eta)) cs = fj.ClusterSequenceArea(_py_fj_parts, jet_def, jet_area_def) pythia_jets = jet_selector_pythia(cs.inclusive_jets()) if len(pythia_jets) < 1: continue pbar.update(1) tw.fill_branches(pyj = pythia_jets) # now generate bg bg_tclones = tgbkg.GetBackground() # tgbkg.GetRandomSeed() nParticles = bg_tclones.GetEntries(); # pinfo('event', pbar.n, 'number of parts', nParticles) # _parts = { 'pt' : [], 'eta' : [], 'phi' : [], 'kf' : []} _parts = [[], [], [], []] _ = [[_parts[0].append(p[0].Pt()), _parts[1].append(p[0].Eta()), _parts[2].append(p[0].Phi()), _parts[3].append(p[1])] for p in [[tlv_from_tmcparticle(_p), _p.GetKF()] for _p in bg_tclones if _p.GetEnergy()>0]] _bg_fj_parts = fjext.vectorize_pt_eta_phi(_parts[0], _parts[1], _parts[2], 1000) #bg particles with index > 1000 # add background and pythia _fj_parts = [] _ = [_fj_parts.append(_p) for _p in _py_fj_parts] _ = [_fj_parts.append(_p) for _p in _bg_fj_parts] # stream all particles _ = [tw.fill_branches(part_pt = _pfj.perp(), part_eta = _pfj.eta(), part_phi = _pfj.phi(), part_idx=_pfj.user_index()) for _pfj in _fj_parts] # find jets in the hybrid event # w/o area # jets = jet_selector_hybrid(jet_def(_fj_parts)) # w/area cs_hybrid = fj.ClusterSequenceArea(_fj_parts, jet_def, jet_area_def) jets = jet_selector_hybrid(cs_hybrid.inclusive_jets()) # stream jets from the hybrid event tw.fill_branches(j = jets) # estimate the background bg_rho_range = fj.SelectorAbsEtaMax(args.eta * 1.1) bg_jet_def = fj.JetDefinition(fj.kt_algorithm, jet_R0) bg_area_def = fj.AreaDefinition(fj.active_area_explicit_ghosts, fj.GhostedAreaSpec(args.eta)) # bg_area_def = fj.AreaDefinition(fj.active_area, fj.GhostedAreaSpec(args.eta)) #active area defunct for bg estim bg_estimator = fj.JetMedianBackgroundEstimator(bg_rho_range, bg_jet_def, bg_area_def) bg_estimator.set_particles(_fj_parts) if len(_fj_parts) < 0: perror('no particles in the hybrid event?') continue rho = bg_estimator.rho() sigma = bg_estimator.sigma() corr_jet_pt = [j.pt() - j.area() * rho for j in jets] # matches = [j.perp(), matched_jet(j, pythia_jets) for j in jets] delta_pt = [delta_pt_matched(j, pythia_jets, rho) for j in jets] tw.fill_branches(j_corr_pt = corr_jet_pt, dpt = delta_pt) tw.fill_branches(rho = rho, rho_sigma = sigma) tw.fill_tree() bg_tclones.Clear() pbar.close() tgbkg.CloserFunction() tw.write_and_close()