def fj_example_02_area(event):
# cluster the event
jet_def = fj.JetDefinition(fj.antikt_algorithm, 0.4)
area_def = fj.AreaDefinition(fj.active_area, fj.GhostedAreaSpec(5.0))
cs = fj.ClusterSequenceArea(event, jet_def, area_def)
jets = fj.SelectorPtMin(5.0)(fj.sorted_by_pt(cs.inclusive_jets()))
print("jet def:", jet_def)
print("area def:", area_def)
print("#-------------------- initial jets --------------------")
print_jets(jets)
#----------------------------------------------------------------------
# estimate the background
maxrap = 4.0
grid_spacing = 0.55
gmbge = fj.GridMedianBackgroundEstimator(maxrap, grid_spacing)
gmbge.set_particles(event)
print("#-------------------- background properties --------------------")
print("rho = ", gmbge.rho())
print("sigma = ", gmbge.sigma())
print()
#----------------------------------------------------------------------
# subtract the jets
subtractor = fj.Subtractor(gmbge)
subtracted_jets = subtractor(jets)
print("#-------------------- subtracted jets --------------------")
print_jets(subtracted_jets)
def analyze_event(self, parts):
self.particles = self.particle_selector(parts)
if len(self.particles) < 1:
self.cs = None
self.jets = []
else:
self.cs = fj.ClusterSequenceArea(self.particles, self.jet_def,
self.jet_area_def)
self.jets = fj.sorted_by_pt(
self.jet_selector(self.cs.inclusive_jets()))
def analyze_event(self, parts): if len(parts) < 1: self.rho = 0.0 self.cs = None self.jets = [] self.corr_jet_pt = [] else: self.bg_estimator.set_particles(parts) self.rho = self.bg_estimator.rho() self.cs = fj.ClusterSequenceArea(parts, self.jet_def, self.jet_area_def) self.jets = fj.sorted_by_pt(self.jet_selector(self.cs.inclusive_jets())) self.corr_jet_pt = [j.pt() - j.area() * self.rho for j in self.jets]
def do_cs_jet_by_jet(full_event, ptmin=100.):
# clustering with ghosts and get the jets
jet_def = fj.JetDefinition(fj.antikt_algorithm, 0.4)
ghost_RapMax = 3.
ghost_spec = fj.GhostedAreaSpec(ghost_RapMax, 1, ghost_area)
area_def = fj.AreaDefinition(fj.active_area_explicit_ghosts, ghost_spec)
clust_seq_full = fj.ClusterSequenceArea(full_event, jet_def, area_def)
full_jets = clust_seq_full.inclusive_jets(ptmin)
# background estimation
jet_def_bge = fj.JetDefinition(fj.kt_algorithm, 0.4)
area_def_bge = fj.AreaDefinition(fj.active_area_explicit_ghosts,
ghost_spec)
bge_range = fj.SelectorAbsRapMax(3.0)
bge = fj.JetMedianBackgroundEstimator(bge_range, jet_def_bge, area_def_bge)
bge.set_particles(full_event)
# subtractor
subtractor = cs.ConstituentSubtractor()
subtractor.set_distance_type(0)
subtractor.set_max_distance(max_distance)
subtractor.set_alpha(alpha)
subtractor.set_max_eta(3.0)
subtractor.set_background_estimator(bge)
subtractor.set_common_bge_for_rho_and_rhom()
#sel_max_pt = fj.SelectorPtMax(10)
#subtractor.set_particle_selector(sel_max_pt)
# do subtraction
corrected_jets = cs.PseudoJetVec()
for jet in full_jets:
subtracted_jet = subtractor.result(jet)
corrected_jets.push_back(subtracted_jet)
# pt cut
selector = fj.SelectorPtMin(ptmin)
return selector(corrected_jets), clust_seq_full
def process_event(self, event):
# assume this is done already (when creating the ntuples):
# self.fj_particles_selected = self.particle_selector(fj_particles)
if self.constit_subtractor:
self.fj_particles_selected = self.constit_subtractor.process_event(
event.particles)
else:
self.fj_particles_selected = event.particles
self.cs = fj.ClusterSequenceArea(self.fj_particles_selected,
self.jet_def, self.jet_area_def)
# self.cs = fj.ClusterSequence(self.fj_particles_selected, self.jet_def)
self.jets = fj.sorted_by_pt(self.cs.inclusive_jets())
self.rho = 0
if self.bg_estimator:
self.bg_estimator.set_particles(self.fj_particles_selected)
self.rho = self.bg_estimator.rho()
if self.jets[0].pt(
) - self.rho * self.jets[0].area() < self.jet_pt_min:
return False
self.jets_selected = self.jet_selector(self.jets)
# remove jets containing bad (high pT) particles
self.jets_detok = list(
filter(
lambda j: max([t.pt() for t in j.constituents()]) < self.
particle_pt_max, self.jets_selected))
# print(len(self.jets_selected), len(self.jets_detok))
self.sd_jets = []
self.sd_jets_info = []
for isd, sd in enumerate(self.sds):
self.sd_jets.append([])
self.sd_jets_info.append([])
for j in self.jets_detok:
for isd, sd in enumerate(self.sds):
jet_sd = sd.result(j)
self.sd_jets[isd].append(jet_sd)
self.sd_jets_info[isd].append(
fjcontrib.get_SD_jet_info(jet_sd))
if self.event_output:
self.event_output.fill_branch('ev_id', event.ev_id)
self.event_output.fill_branch('run_number', event.run_number)
self.event_output.fill_branch('rho', self.rho)
self.event_output.fill_tree()
if self.jet_output:
self.jet_output.fill_branch('ev_id', event.ev_id)
self.jet_output.fill_branch('run_number', event.run_number)
self.jet_output.fill_branch('rho', self.rho)
self.jet_output.fill_branch('jet', [j for j in self.jets_detok])
self.jet_output.fill_branch(
'jet_ptsub',
[j.pt() - (self.rho * j.area()) for j in self.jets_detok])
for isd, sd in enumerate(self.sds):
bname = 'jet_sd{}pt'.format(self.sd_betas[isd])
self.jet_output.fill_branch(
bname, [j.pt() for j in self.sd_jets[isd]])
bname = 'jet_sd{}zg'.format(self.sd_betas[isd])
self.jet_output.fill_branch(
bname, [j.z for j in self.sd_jets_info[isd]])
bname = 'jet_sd{}Rg'.format(self.sd_betas[isd])
self.jet_output.fill_branch(
bname, [j.dR for j in self.sd_jets_info[isd]])
bname = 'jet_sd{}thetag'.format(self.sd_betas[isd])
self.jet_output.fill_branch(
bname, [j.dR / self.jet_R for j in self.sd_jets_info[isd]])
self.jet_output.fill_tree()
return True
def run(data, n_events = 1000000):
out = []
njets = []
met = []
rho = []
jets_corr = []
area = []
# run jet clustering with AntiKt, R=1.0
R = 1.0
jet_def = fj.JetDefinition(fj.antikt_algorithm, R)
# Area Definition
ghost_maxrap = 4.7
area_def = fj.AreaDefinition(fj.active_area, fj.GhostedAreaSpec(ghost_maxrap))
# Background estimator
select_rapidity = fj.SelectorAbsRapMax(ghost_maxrap)
bge = fj.JetMedianBackgroundEstimator(select_rapidity, jet_def, area_def)
# Loop over events
for ievt in range(n_events):
# Build a list of all particles and MET Information
pjs = []
pjmet = fj.PseudoJet()
pjmets = fj.PseudoJet()
for i in range(int(data.shape[1]/3)):
pj = fj.PseudoJet()
pj.reset_PtYPhiM(data.at[ievt, 3*i+0], data.at[ievt, 3*i+1], data.at[ievt, 3*i+2], 0)
if pj.pt() > 1.0:
pjs.append(pj)
pjmet.reset_momentum(-pj.px(), -pj.py(), 0, pj.E())
pjmets = pjmets + pjmet
met.append(pjmets)
# Cluster sequence
clust_seq = fj.ClusterSequenceArea(pjs, jet_def, area_def)
jets = fj.sorted_by_pt(clust_seq.inclusive_jets())
jets = [j for j in jets if j.pt() > 30. and j.eta() < 4.7]
area.append([jets[0].area(), jets[1].area()])
# Save the two leading jets and njets
jets_sel = (fj.SelectorNHardest(2))(jets)
out.append(jets_sel)
njets.append(len(jets))
# Background estimator
bge.set_particles(pjs)
rho.append(bge.rho())
# Correct Jets
sub = fj.Subtractor(bge)
sub_jets = sub(jets_sel)
jets_corr.append(sub_jets)
return out, area, met, njets, rho, jets_corr
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(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.SelectorPtMax(
1000.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'])
_ts = pds_trks.loc[pds_trks['ev_id'] == iev_id]
start = time.time()
_tpsj = fj_parts_from_tracks_numpy(_ts)
end = time.time()
dt_swig = end - start
start = time.time()
_tpsj_for = fj_parts_from_tracks(_ts)
end = time.time()
dt_for = end - start
# print ('len {} =?= {}'.format(len(_tpsj_for), len(_tpsj)))
print(
'[i] timing (ntracks={}): dt_for: {} dt_swig: {} ratio: {}'.format(
len(_tpsj), dt_for, dt_swig, dt_for / dt_swig))
# 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())
# print("sigma = ", gmbge.sigma())
# _jets = jet_selector(jet_def(_tpsj))
# _jets_a = [[iev_id, j.perp(), j.eta(), j.phi()] for j in _jets]
# _jets_a = pd.DataFrame(np.array([[iev_id, j.perp(), j.eta(), j.phi()] for j in _jets]), columns=['evid', 'pt', 'eta', 'phi'])
_jets_a = 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)
# , columns=['evid, pt, eta, phi']
e_jets = e_jets.append(_jets_a, ignore_index=True)
# print('event', i, 'number of parts', len(_tpsj), 'number of jets', len(_jets))
# print(_jets_a.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()