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()