def __init__(self, **kwargs):
self.configure_from_args(jet_R=0.4,
jet_algorithm=fj.antikt_algorithm,
particle_eta_max=0.9,
jet_pt_min=0.0,
particles=None,
explicit_ghosts=False)
super(JetAnalysis, self).__init__(**kwargs)
self.particle_selector = fj.SelectorAbsEtaMax(self.particle_eta_max)
self.jet_eta_max = self.particle_eta_max - self.jet_R * 1.05
# self.jet_eta_max = self.particle_eta_max - 0.05
self.jet_def = fj.JetDefinition(self.jet_algorithm, self.jet_R)
self.jet_selector = fj.SelectorPtMin(
self.jet_pt_min) & fj.SelectorAbsEtaMax(self.jet_eta_max)
if self.explicit_ghosts:
self.jet_area_def = fj.AreaDefinition(
fj.active_area_explicit_ghosts,
fj.GhostedAreaSpec(self.particle_eta_max))
else:
self.jet_area_def = fj.AreaDefinition(
fj.active_area, fj.GhostedAreaSpec(self.particle_eta_max))
if self.particles:
self.analyze_event(self.particles)
def __init__(self, **kwargs):
self.configure_from_args(jet_R=0.4,
jet_algorithm=fj.antikt_algorithm,
particle_eta_max=0.9,
particles=None)
super(JetAnalysis, self).__init__(**kwargs)
self.bg_rho_range = fj.SelectorAbsEtaMax(self.particle_eta_max)
self.bg_jet_def = fj.JetDefinition(fj.kt_algorithm, self.jet_R)
self.bg_area_def = fj.AreaDefinition(
fj.active_area_explicit_ghosts,
fj.GhostedAreaSpec(self.particle_eta_max))
self.bg_estimator = fj.JetMedianBackgroundEstimator(
self.bg_rho_range, self.bg_jet_def, self.bg_area_def)
self.rho = 0
self.jet_eta_max = self.particle_eta_max - self.jet_R * 1.05
self.jet_def = fj.JetDefinition(self.jet_algorithm, self.jet_R)
self.jet_selector = fj.SelectorPtMin(0.0) & fj.SelectorAbsEtaMax(
self.jet_eta_max)
self.jet_area_def = fj.AreaDefinition(
fj.active_area_explicit_ghosts,
fj.GhostedAreaSpec(self.particle_eta_max))
if self.particles:
self.analyze_event(self.particles)
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_file_list(file_inputs=[],
output_prefix='rg',
tree_name='tree_Particle'):
anl = []
bg_rho_range = fj.SelectorAbsRapMax(0.9)
bg_jet_def = fj.JetDefinition(fj.kt_algorithm, 0.4)
bg_area_def = fj.AreaDefinition(fj.active_area_explicit_ghosts,
fj.GhostedAreaSpec(0.9))
bg_estimator = fj.JetMedianBackgroundEstimator(bg_rho_range, bg_jet_def,
bg_area_def)
print('[i]', bg_estimator.description())
an_std = MPJetAnalysis(output_prefix=output_prefix,
name='std',
bg_estimator=bg_estimator)
anl.append(an_std)
g_bg_estimator = fj.GridMedianBackgroundEstimator(0.9, 0.4)
print('[i]', g_bg_estimator.description())
an_gbg = MPJetAnalysis(output_prefix=output_prefix,
name='bgb',
bg_estimator=g_bg_estimator)
anl.append(an_gbg)
cs004 = csubtractor.CEventSubtractor(alpha=0,
max_distance=0.4,
max_eta=0.9,
bge_rho_grid_size=0.25,
max_pt_correct=100)
an_cs_004 = MPJetAnalysis(output_prefix=output_prefix,
name='cs004',
bg_estimator=cs004.bge_rho,
constit_subtractor=cs004)
anl.append(an_cs_004)
cs404 = csubtractor.CEventSubtractor(alpha=4,
max_distance=0.4,
max_eta=0.9,
bge_rho_grid_size=0.25,
max_pt_correct=100)
an_cs_404 = MPJetAnalysis(output_prefix=output_prefix,
name='cs404',
bg_estimator=cs404.bge_rho,
constit_subtractor=cs404)
anl.append(an_cs_404)
ad = MPAnalysisDriver(file_list=file_inputs,
analyses_list=anl,
tree_name=tree_name)
ad.run()
print()
print(' ---')
print()
print(ad)
print('[i] done.')
def __init__(self, **kwargs):
self.configure_constants(
particle_eta_max=0.9,
particle_pt_max=
100, # reject jets with constituent with pT > max_particle_pt
jet_R=0.4,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=-200,
bge_rho_grid_size=-1, # no background subtraction for -1
sd_z_cuts=[0.1],
sd_betas=[0, 2],
show_progress=False,
name="MPJetAnalysis",
output_prefix='./')
self.jet_eta_max = self.particle_eta_max - self.jet_R * 1.05
self.jet_def = fj.JetDefinition(self.jet_algorithm, self.jet_R)
self.jet_selector = fj.SelectorPtMin(5.0) & fj.SelectorAbsRapMax(
self.jet_eta_max)
self.particle_selector = fj.SelectorPtMin(0.15) & fj.SelectorAbsRapMax(
self.particle_eta_max)
self.jet_area_def = fj.AreaDefinition(
fj.active_area_explicit_ghosts,
fj.GhostedAreaSpec(self.particle_eta_max))
self.bg_estimator = None
self.constit_subtractor = None
self.jet_output = None
self.event_output = None
self.fout = None
super(MPJetAnalysis, self).__init__(**kwargs)
self.filename_output = '{}{}.root'.format(self.output_prefix,
self.name)
if self.fout is None:
self.fout = ROOT.TFile(self.filename_output, 'recreate')
if self.jet_output is None:
self.jet_output = treewriter.RTreeWriter(tree_name='tjet',
fout=self.fout)
if self.event_output is None:
self.event_output = treewriter.RTreeWriter(tree_name='tev',
fout=self.fout)
fj.ClusterSequence.print_banner()
print()
print(self)
print(' . particle selector', self.particle_selector)
print(' . jet selector', self.jet_selector)
print(' . jet definition:', self.jet_def)
# define softdrops
self.sds = []
for beta in self.sd_betas:
for zcut in self.sd_z_cuts:
_sd = fjcontrib.SoftDrop(beta, zcut, self.jet_R)
self.sds.append(_sd)
for _sd in self.sds:
print(' . sd:', _sd.description())
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 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()
