def __init__(self, **kwargs):
self.configure_from_args(file_list='PbPb_file_list.txt',
n_events=500,
output='output',
run_jetfinder=False)
super(Data2MP4, self).__init__(**kwargs)
self.data_io = DataIO(file_list=self.file_list)
#self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -1, 1, 51, 0., 2.*ROOT.TMath.Pi())
#self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -1, 1, 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
#self.hetaphi_jet = ROOT.TH2F('_hetaphi_jet', '_hetaphi_jet;#eta;#varphi', 51, -1, 1, 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
#self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 101, -ROOT.TMath.Pi(), ROOT.TMath.Pi(), 101, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -1,
1, 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.hetaphi_jet = ROOT.TH2F('_hetaphi_jet',
'_hetaphi_jet;#eta;#varphi', 51,
-ROOT.TMath.Pi(), ROOT.TMath.Pi(), 51,
-ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.mean_eta = ROOT.TH1F('_hmean_eta', '_hmean_eta;<#eta>', 101, -0.3,
0.3)
# self.mean_phi = ROOT.TH1F('_hmean_phi', '_hmean_phi;<#varphi>', 101, ROOT.TMath.Pi()-0.3, ROOT.TMath.Pi()+0.3)
self.mean_phi = ROOT.TH1F('_hmean_phi', '_hmean_phi;<#varphi>', 101,
-0.3, +0.3)
self.mean_e = ROOT.TH1F('_hmean_e', '_hmean_e;<E (GeV)>', 100, 0, 2)
self.tc = None
if self.run_jetfinder:
fj.ClusterSequence.print_banner()
def __init__(self, **kwargs):
self.configure_from_args(tree_name='tree_Particle',
tree_name_gen='tree_Particle_gen',
args=None)
super(Embedding, self).__init__(**kwargs)
self.copy_attributes(self.args)
self.jet_def = fj.JetDefinition(fj.antikt_algorithm, self.jetR)
if self.benchmark:
self.jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax(
100.0) & fj.SelectorAbsEtaMax(self.max_eta - 1.05 * self.jetR)
# jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * self.jetR)
else:
self.jet_selector = fj.SelectorAbsEtaMax(self.max_eta -
1.05 * self.jetR)
self.parts_selector = fj.SelectorAbsEtaMax(self.max_eta)
self.output = EmbeddingOutput(args=self.args)
# self.output.copy_attributes(self)
self.sd = fjcontrib.SoftDrop(0, self.sd_zcut, self.jetR)
self.ja_part = JetAnalysis(jet_R=self.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=5.,
particle_eta_max=self.max_eta)
self.ja_det = JetAnalysis(jet_R=self.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=self.jetptcut,
particle_eta_max=self.max_eta)
self.ja_hybrid = JetAnalysis(jet_R=self.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=5.,
particle_eta_max=self.max_eta)
self.dataPbPb = DataBackgroundIO(name='Data PbPb',
file_list=self.datalistAA)
self.det_sim = DataIO(name='Sim Pythia Detector level',
file_list=self.simulationpp,
random_file_order=False)
self.part_sim = DataIO(name='Sim Pythia Particle level',
file_list=self.simulationpp,
random_file_order=False,
tree_name='tree_Particle_gen')
self.cs = None
if self.dRmax > 0:
self.cs = CEventSubtractor(alpha=self.alpha,
max_distance=self.dRmax,
max_eta=self.max_eta,
bge_rho_grid_size=0.25,
max_pt_correct=100)
def main_make(args):
if os.path.isfile(args.output):
if not args.overwrite:
print('[i] output', args.output, 'exists - use --overwrite to do just that...')
return
print(args)
# alice specific
max_eta = 0.9
outf = ROOT.TFile(args.output, 'recreate')
outf.cd()
bgestim = BackgroundEstimator(name='bg_estim', grid_size=0.1, particle_eta_max=max_eta)
print(bgestim)
be = None
data = None
if len(args.data) > 0:
data = DataIO(file_list=args.data)
print(data)
else:
be = BoltzmannEvent(mean_pt=0.6, multiplicity=2000 * max_eta * 2, max_eta=max_eta, max_pt=100)
print(be)
parts_selector = fj.SelectorAbsEtaMax(max_eta)
if args.nev < 1:
args.nev = 1
### EVENT LOOP STARTS HERE
for iev in tqdm.tqdm(range(args.nev)):
if data:
bg_parts = data.load_event(offset=10000)
else:
if be:
bg_parts = be.generate(offset=10000)
bgestim.fill_grid(bg_parts)
outf.cd()
if be:
be.write()
bgestim.write()
outf.Write()
outf.Close()
print('[i] written', outf.GetName())
def __init__(self, **kwargs):
self.configure_from_args(file_list='PbPb_file_list.txt',
n_events=500,
output='output')
super(Data2MP4Morph, self).__init__(**kwargs)
self.data_io = DataIO(file_list=self.file_list)
# self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -ROOT.TMath.Pi(), ROOT.TMath.Pi(), 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
# self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 101, -1, 1, 101, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#varphi;#eta', 101,
-ROOT.TMath.Pi(), +ROOT.TMath.Pi(), 101, -1,
1)
self.tc = None
self.iframe = 0
self.pythia = PythiaJetty(pthatmin=100,
eta_max=1,
jet_pt_min=20,
jet_R0=0.4)
class Embedding(MPBase):
def add_arguments_to_parser(parser):
parser.add_argument('-o',
'--output-filename',
default="output.root",
type=str)
parser.add_argument('datalistAA',
help='run through a file list',
default='',
type=str)
parser.add_argument('simulationpp',
help='run through a file list',
default='',
type=str)
parser.add_argument('--jetR', default=0.4, type=float)
parser.add_argument('--alpha', default=0, type=float)
parser.add_argument('--dRmax', default=0.25, type=float)
parser.add_argument('--sd-zcut', default=0.1, type=float)
parser.add_argument('--overwrite',
help="overwrite output",
default=False,
action='store_true')
parser.add_argument('--benchmark',
help='benchmark pthat setting - 80 GeV',
default=False,
action='store_true')
parser.add_argument('--jetptcut',
help='remove jets below the cut',
default=1.e-3,
type=float)
parser.add_argument('--nev',
help='number of events to run',
default=0,
type=int)
parser.add_argument('--max-eta',
help='max eta for particles',
default=0.9)
def __init__(self, **kwargs):
self.configure_from_args(tree_name='tree_Particle',
tree_name_gen='tree_Particle_gen',
args=None)
super(Embedding, self).__init__(**kwargs)
self.copy_attributes(self.args)
self.jet_def = fj.JetDefinition(fj.antikt_algorithm, self.jetR)
if self.benchmark:
self.jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax(
100.0) & fj.SelectorAbsEtaMax(self.max_eta - 1.05 * self.jetR)
# jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * self.jetR)
else:
self.jet_selector = fj.SelectorAbsEtaMax(self.max_eta -
1.05 * self.jetR)
self.parts_selector = fj.SelectorAbsEtaMax(self.max_eta)
self.output = EmbeddingOutput(args=self.args)
# self.output.copy_attributes(self)
self.sd = fjcontrib.SoftDrop(0, self.sd_zcut, self.jetR)
self.ja_part = JetAnalysis(jet_R=self.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=5.,
particle_eta_max=self.max_eta)
self.ja_det = JetAnalysis(jet_R=self.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=self.jetptcut,
particle_eta_max=self.max_eta)
self.ja_hybrid = JetAnalysis(jet_R=self.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=5.,
particle_eta_max=self.max_eta)
self.dataPbPb = DataBackgroundIO(name='Data PbPb',
file_list=self.datalistAA)
self.det_sim = DataIO(name='Sim Pythia Detector level',
file_list=self.simulationpp,
random_file_order=False)
self.part_sim = DataIO(name='Sim Pythia Particle level',
file_list=self.simulationpp,
random_file_order=False,
tree_name='tree_Particle_gen')
self.cs = None
if self.dRmax > 0:
self.cs = CEventSubtractor(alpha=self.alpha,
max_distance=self.dRmax,
max_eta=self.max_eta,
bge_rho_grid_size=0.25,
max_pt_correct=100)
def run(self):
# need to change this for data to drive...
delta_t = 0
start_t = time.time()
iev = 1
# while self.det_sim.load_event() and self.part_sim.load_event():
while self.det_sim.load_event():
iev = iev + 1
if self.nev > 0:
if iev > self.nev:
iev = iev - 1
break
if iev % 1000 == 0:
delta_t = time.time() - start_t
pinfo('processing event', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
# find jets on detector level
if len(self.det_sim.particles) < 1:
pwarning(iev, 'event skipped N detector parts',
len(self.det_sim.particles))
continue
self.ja_det.analyze_event(self.det_sim.particles)
_jets_det = self.ja_det.jets
# _x = [pdebug(' -d ', j) for j in _jets_det]
if len(_jets_det) < 1:
continue
_too_high_pt = [
p.pt() for j in _jets_det for p in j.constituents()
if p.pt() > 100.
]
if len(_too_high_pt) > 0:
pwarning(iev, 'a likely fake high pT particle(s)',
_too_high_pt, '- skipping whole event')
continue
_output_fname = os.path.expanduser(
os.path.expandvars(self.det_sim.file_io.file_input))
_output_fname = _output_fname.replace("/", "_")
self.output.initialize_output(_output_fname)
self.output.fill_det_level(iev, _jets_det)
# load the corresponding event on particle level
self.part_sim.open_afile(afile=self.det_sim.file_io.file_input)
if not self.part_sim.load_event_with_loc(
self.det_sim.event.run_number, self.det_sim.event.ev_id,
0):
perror('unable to load partL event run#:',
self.det_sim.event.run_number, 'ev_id:',
self.det_sim.event.ev_id)
continue
if self.det_sim.event.run_number != self.part_sim.event.run_number:
perror('run# missmatch detL:', self.det_sim.event.run_number,
'partL:', self.part_sim.event.run_number)
continue
if self.det_sim.event.ev_id != self.part_sim.event.ev_id:
perror('ev_id# missmatch detL:', self.det_sim.event.ev_id,
'partL:', self.part_sim.event.ev_id)
continue
# find jets on particle level
if len(self.part_sim.particles) < 1:
pwarning(iev, 'event skipped N particle parts',
len(self.part_sim.particles))
continue
self.ja_part.analyze_event(self.part_sim.particles)
_jets_part = self.ja_part.jets
# _x = [pdebug(' -p ', j) for j in _jets_part]
if len(_jets_part) < 1:
continue
# match in pp simulations
_det_part_matches = []
_n_matches = 0
_part_psjv = self.ja_part.jets_as_psj_vector()
for j_det in _jets_det:
_mactches_pp = fjtools.matched_Reta(j_det, _part_psjv,
0.6 * self.jetR)
#_mactches_pp = fjtools.matched_Ry(j_det, _part_psjv, 0.6 * self.jetR)
_n_matches = _n_matches + len(_mactches_pp)
if len(_mactches_pp) > 1:
pwarning('event:', iev, 'jet pt=', j_det.pt(),
'more than one match in pp jets',
[i for i in _mactches_pp])
if len(_mactches_pp) == 1:
j_part = _part_psjv[_mactches_pp[0]]
# pinfo('j_det', j_det, 'j_part', j_part)
_det_part_matches.append([j_det, j_part])
self.output.fill_pp_pairs(iev, [j_det, j_part])
if _n_matches < 1:
if _n_matches < 1:
pwarning('event:', iev,
'- no matched jets in simulation!?',
len(_det_part_matches))
# here embedding to PbPb data
_offset = 10000
while _offset < len(self.det_sim.particles):
_offset = _offset + 1000
pwarning('increasing bg index offset to', _offset)
_PbPb_loaded = 0
while _PbPb_loaded == 0:
if not self.dataPbPb.load_event(offset=_offset):
perror('unable to load next PbPb event')
_PbPb_loaded = -1
else:
_hybrid_event = self.dataPbPb.particles
_nparts_hybrid_no_emb = len(_hybrid_event)
if _nparts_hybrid_no_emb < 1:
pwarning(
'hybrid event with no particles! trying another one'
)
_PbPb_loaded = 0
else:
_PbPb_loaded = 1
if _PbPb_loaded < 0:
perror(
'unable to load PbPb event - permanent - bailing out here.'
)
break
_tmp = [_hybrid_event.push_back(p) for p in self.det_sim.particles]
if self.cs:
cs_parts = self.cs.process_event(_hybrid_event)
rho = self.cs.bge_rho.rho()
self.ja_hybrid.analyze_event(cs_parts)
else:
self.ja_hybrid.analyze_event(_hybrid_event)
_hybrid_matches = []
_hybrid_psjv = self.ja_hybrid.jets_as_psj_vector()
for m in _det_part_matches:
j_det = m[0]
j_part = m[1]
_mactches_hybrid = fjtools.matched_Reta(
j_det, _hybrid_psjv, 0.6 * self.jetR)
if len(_mactches_hybrid) > 1:
pwarning('event:', iev, 'jet pt=', j_det.pt(),
'more than one match in hybrid jets',
[i for i in _mactches_hybrid])
if len(_mactches_hybrid) == 1:
# m.append(_hybrid_psjv[_mactches_hybrid[0]])
j_hybr = _hybrid_psjv[_mactches_hybrid[0]]
# pdebug('L302', 'j_det', j_det, 'j_part', j_part, 'j_hybr', j_hybr)
_hybrid_matches.append([j_det, j_part, j_hybr])
self.output.fill_emb_3(iev, [j_det, j_part, j_hybr])
_n_matches_hybrid = len(_hybrid_matches)
if _n_matches_hybrid < 1:
if _n_matches_hybrid < 1:
pwarning('event:', iev, '- no matched jets in embedding!?',
_n_matches_hybrid)
delta_t = time.time() - start_t
pinfo('processed events', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
self.output.close()
def main():
parser = argparse.ArgumentParser(description='test groomers',
prog=os.path.basename(__file__))
parser.add_argument('-o',
'--output-filename',
default="output.root",
type=str)
parser.add_argument('datalistpp',
help='run through a file list',
default='',
type=str)
parser.add_argument('--datalistAA',
help='run through a file list - embedding mode',
default='',
type=str)
parser.add_argument('--jetR', default=0.4, type=float)
parser.add_argument('--alpha', default=0, type=float)
parser.add_argument('--dRmax', default=0.25, type=float)
parser.add_argument('--overwrite',
help="overwrite output",
default=False,
action='store_true')
parser.add_argument('--jetptcut',
help='remove jets below the cut',
default=50.,
type=float)
parser.add_argument('--nev',
help='number of events to run',
default=0,
type=int)
parser.add_argument('--max-eta',
help='max eta for particles',
default=0.9,
type=float)
parser.add_argument('--npart-cut',
help='npart cut on centrality low,high hint:' +
npart_cents,
default='325,450',
type=str)
args = parser.parse_args()
try:
npart_min = int(args.npart_cut.split(',')[0])
npart_max = int(args.npart_cut.split(',')[1])
except:
perror(
'unable to parse npart centrality selection - two integer numbers with a coma in-between needed - specified:',
args.npart_cut)
return 1
# initialize constituent subtractor
cs = None
if args.dRmax > 0:
cs = CEventSubtractor(alpha=args.alpha,
max_distance=args.dRmax,
max_eta=args.max_eta,
bge_rho_grid_size=0.25,
max_pt_correct=100)
pp_data = DataIO(name='Sim Pythia Detector level',
file_list=args.datalistpp,
random_file_order=False,
tree_name='tree_Particle_gen')
ja_pp = JetAnalysis(jet_R=args.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=50.,
particle_eta_max=args.max_eta)
if args.datalistAA:
aa_data = DataBackgroundIO(name='PbPb',
file_list=args.datalistAA,
tree_name='tree_Particle_gen')
ja_emb = JetAnalysis(jet_R=args.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=50.,
particle_eta_max=args.max_eta)
ja_aa = JetAnalysis(jet_R=args.jetR,
jet_algorithm=fj.antikt_algorithm,
jet_pt_min=50.,
particle_eta_max=args.max_eta)
dndeta_selector = fj.SelectorAbsEtaMax(1.)
# tg = thg.ThermalGenerator()
print(cs)
# print the banner first
fj.ClusterSequence.print_banner()
print()
gout = GroomerOutput(args.output_filename,
enable_aa_trees=bool(args.datalistAA))
delta_t = 0
start_t = time.time()
iev = 1
while pp_data.load_event(offset=0):
iev = iev + 1
if args.nev > 0:
if iev > args.nev:
iev = iev - 1
break
if iev % 1000 == 0:
delta_t = time.time() - start_t
pinfo('processing event', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
# find jets on detector level
if len(pp_data.particles) < 1:
pwarning(iev, 'pp event skipped N parts', len(pp_data.particles))
continue
ja_pp.analyze_event(pp_data.particles)
if len(ja_pp.jets) < 1:
continue
# pinfo('n particles', len(pp_data.particles))
dndeta0 = dndeta_selector(pp_data.particles)
[
gout.fill_branches(j, syst=0, dndeta=len(dndeta0) / 2.)
for j in ja_pp.jets
]
# pinfo('n jets', len(ja_pp.jets))
if args.datalistAA:
while True:
aa_loaded = aa_data.load_event(offset=10000)
if aa_data.event.npart < npart_min or aa_data.event.npart >= npart_max:
continue
else:
if len(aa_data.particles) < 1:
pwarning(iev, 'AA event skipped N parts',
len(aa_data.particles))
continue
else:
break
if aa_loaded:
ja_aa.analyze_event(aa_data.particles)
dndeta1 = dndeta_selector(aa_data.particles)
if len(ja_aa.jets) > 0:
[
gout.fill_branches(j, syst=1, dndeta=len(dndeta1) / 2.)
for j in ja_aa.jets
]
else:
# pwarning('no jets in AA event?', len(ja_aa.jets), 'while dndeta=', len(dndeta1)/2.)
pass
emb_event = fj.vectorPJ()
[emb_event.push_back(p) for p in pp_data.particles]
[emb_event.push_back(p) for p in aa_data.particles]
rho = 0
if cs:
cs_parts = cs.process_event(emb_event)
rho = cs.bge_rho.rho()
ja_emb.analyze_event(cs_parts)
else:
ja_emb.analyze_event(emb_event)
# matches = [[jpp, jemb] for jpp in ja_pp.jets for jemb in ja_emb.jets if fjtools.matched_pt(jemb, jpp) > 0.5]
# for mj in matches:
# gout.fill_branches(mj[0], syst=2, dndeta=len(dndeta1)/2., rho=rho)
# gout.fill_branches(mj[1], syst=3)
[
gout.fill_branches_prong_matching(j_pp,
j_emb,
dndeta=len(dndeta1) / 2.,
rho=rho)
for j_pp in ja_pp.jets for j_emb in ja_emb.jets
]
delta_t = time.time() - start_t
pinfo('processed events', iev, ' - ev/sec =', iev / delta_t, 'elapsed =',
delta_t)
gout.write()
def main():
parser = argparse.ArgumentParser(description='test groomers',
prog=os.path.basename(__file__))
parser.add_argument('-o',
'--output-filename',
default="centrality_output.root",
type=str)
parser.add_argument('datalist',
help='run through a file list',
default='',
type=str)
parser.add_argument('--overwrite',
help="overwrite output",
default=False,
action='store_true')
parser.add_argument('--nev',
help='number of events to run',
default=0,
type=int)
parser.add_argument('--max-eta', help='max eta for particles', default=0.9)
parser.add_argument('--thermal',
help='enable thermal generator',
action='store_true',
default=False)
parser.add_argument('--thermal-default',
help='enable thermal generator',
action='store_true',
default=False)
parser.add_argument('--particles',
help='stream particles',
action='store_true',
default=False)
parser.add_argument('--npart-cut',
help='npart cut on centrality low,high hint:' +
npart_cents,
default='325,450',
type=str)
parser.add_argument('--nch-cut',
help='nch cut on centrality low,high hint:' +
nch_cents,
default='18467,50000',
type=str)
args = parser.parse_args()
try:
npart_min = int(args.npart_cut.split(',')[0])
npart_max = int(args.npart_cut.split(',')[1])
except:
perror(
'unable to parse npart centrality selection - two integer numbers with a coma in-between needed - specified:',
args.npart_cut)
return 1
try:
nch_min = int(args.nch_cut.split(',')[0])
nch_max = int(args.nch_cut.split(',')[1])
except:
perror(
'unable to parse nch centrality selection - two integer numbers with a coma in-between needed - specified:',
args.nch_cut)
return 1
outf = ROOT.TFile(args.output_filename, 'recreate')
outf.cd()
t = ROOT.TTree('t', 't')
tw = RTreeWriter(tree=t)
hpt_antyr = ROOT.TH1F('hpt_antyr', 'hpt_antyr', 100, 0, 100)
hpt_antyr_c = ROOT.TH1F('hpt_antyr_c', 'hpt_antyr_c', 100, 0, 100)
hpt_therm = ROOT.TH1F('hpt_therm', 'hpt_therm', 100, 0, 100)
hpt_therm_c = ROOT.TH1F('hpt_therm_c', 'hpt_therm_c', 100, 0, 100)
data = DataIO(name='Sim Pythia Detector level',
file_list=args.datalist,
random_file_order=False,
tree_name='tree_Particle_gen')
dndeta_selector = fj.SelectorAbsEtaMax(abs(
args.max_eta)) & fj.SelectorPtMin(0.15)
tg_default = None
if args.thermal_default:
tg_default = thg.ThermalGenerator()
print(tg_default)
tg_central = None
if args.thermal:
tg_central = thg.ThermalGenerator(beta=0.5, N_avg=3000, sigma_N=500)
print(tg_central)
delta_t = 0
start_t = time.time()
iev = 1
while data.load_event(offset=0):
iev = iev + 1
if args.nev > 0:
if iev > args.nev:
iev = iev - 1
break
if iev % 1000 == 0:
delta_t = time.time() - start_t
pinfo('processing event', iev, ' - ev/sec =', iev / delta_t,
'elapsed =', delta_t)
# find jets on detector level
if len(data.particles) < 1:
pwarning(iev, 'pp event skipped N parts', len(data.particles))
continue
# print(data.event)
dndeta0_parts = dndeta_selector(data.particles)
dndeta0 = len(dndeta0_parts) / (abs(args.max_eta * 2.))
[hpt_antyr.Fill(p.perp()) for p in dndeta0_parts]
if args.particles:
tw.fill_branches(dndeta=dndeta0, p=data.particles)
else:
tw.fill_branches(dndeta=dndeta0)
tw.fill_branches_attribs(
data.event, ['sigma', 'npart', 'nch', 'nchfwd', 'nchselect'],
prefix='antyr_')
if data.event.npart < npart_min or data.event.npart >= npart_max:
tw.fill_branches(cent10npart=0)
else:
tw.fill_branches(cent10npart=1)
[hpt_antyr_c.Fill(p.perp()) for p in dndeta0_parts]
if data.event.nch < nch_min or data.event.nch >= nch_max:
tw.fill_branches(cent10nch=0)
else:
tw.fill_branches(cent10nch=1)
if tg_default:
thg_particles = tg_default.load_event()
dndetathg_default = dndeta_selector(thg_particles)
if args.particles:
tw.fill_branches(dndeta_thg_0=len(dndetathg_default) /
(abs(args.max_eta * 2.)),
p_thg_0=thg_particles)
else:
tw.fill_branches(dndeta_thg_0=len(dndetathg_default) /
(abs(args.max_eta * 2.)))
if tg_central:
thg_parts_central = tg_central.load_event()
dndetathg_central = dndeta_selector(thg_parts_central)
[hpt_therm_c.Fill(p.perp()) for p in dndetathg_central]
if args.particles:
tw.fill_branches(dndeta_thg_c=len(dndetathg_central) /
(abs(args.max_eta * 2.)),
p_thg_c=thg_parts_central)
else:
tw.fill_branches(dndeta_thg_c=len(dndetathg_central) /
(abs(args.max_eta * 2.)))
tw.fill_tree()
delta_t = time.time() - start_t
pinfo('processed events', iev, ' - ev/sec =', iev / delta_t, 'elapsed =',
delta_t)
outf.Write()
outf.Close()
def main_test(args):
if os.path.isfile(args.output):
if not args.overwrite:
print('[i] output', args.output, 'exists - use --overwrite to do just that...')
return
print(args)
# alice specific
max_eta = 0.9
jet_R0 = 0.4
mycfg = []
if args.benchmark:
mycfg = ['PhaseSpace:pThatMin = 80', 'PhaseSpace:pThatMax = -1']
jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax(100.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
# jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
else:
args.py_biaspow = 4
args.py_biasref = 10
jet_selector = fj.SelectorPtMin(20) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
# jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
if args.ignore_mycfg:
mycfg = []
pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
if not pythia:
print("[e] pythia initialization failed.")
return
sd_zcut = 0.1
sd = fjcontrib.SoftDrop(0, sd_zcut, jet_R0)
parts_selector = fj.SelectorAbsEtaMax(max_eta)
fout = ROOT.TFile(args.output, 'recreate')
fout.cd()
twj = RTreeWriter(tree_name='jets', fout=fout)
twjm = RTreeWriter(tree_name='jetsm', fout=fout)
twp = RTreeWriter(tree_name='parts', fout=fout)
bgestim = BackgroundEstimator(input_file=args.input, name='bg_estim')
print(bgestim)
boltzmann_subtractor = BoltzmannSubtractor(max_pt_subtract=10.)
print(boltzmann_subtractor)
be = None
data = None
if len(args.data) > 0:
data = DataIO(file_list=args.data)
print(data)
else:
# be = BoltzmannEvent(mean_pt=0.6, multiplicity=2000 * max_eta * 2, max_eta=max_eta, max_pt=100)
pass
be = fjtools.BoltzmannBackground(0.7, 0.01, 1.5)
print(be.description())
parts_selector = fj.SelectorAbsEtaMax(max_eta)
if args.nev < 1:
args.nev = 1
jas = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta)
ja = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta)
jabg = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta)
### EVENT LOOP STARTS HERE
pbar = tqdm.tqdm(total=args.nev)
iev = 0
while pbar.n < args.nev:
iev = iev + 1
if not pythia.next():
continue
parts_pythia = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged])
parts_gen = parts_selector(parts_pythia)
jas.analyze_event(parts_gen)
signal_jets = fj.sorted_by_pt(jet_selector(jas.jets))
if len(signal_jets) < 1:
continue
pbar.set_description('tried %i' % iev)
pbar.update(1)
if data:
bg_parts = data.load_event(offset=10000)
else:
if be:
# bg_parts = be.generate(offset=10000)
bg_parts = be.generate(int(2000 * max_eta * 2.), max_eta, 10000)
# len_bg_parts = len(bg_parts)
# subtr_parts_b = be.subtract(bg_parts, 0.7, len(bg_parts))
# print('x nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
# subtr_parts_b = be.subtract(bg_parts)
# print(' - nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
_pythia_inserts = []
_pythia_inserts_indexes = []
for sjet in signal_jets:
_pythia_inserts = [bg_parts.push_back(psj) for psj in sjet.constituents() if psj.user_index() >= 0]
_pythia_inserts_indexes.extend([psj.user_index() for psj in sjet.constituents() if psj.user_index() >= 0])
print ('pythia indexes:', len(_pythia_inserts_indexes), sorted(_pythia_inserts_indexes))
# subtr_parts_b = be.subtract(bg_parts)
# print(' - nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
# subtr_parts_b = be.subtract(bg_parts, 0.7, len(_pythia_inserts))
# print('. nparts_bg=', len_bg_parts, 'nparts_pass='******' ', be.get_formula())
# _mpt = mean_pt(bg_parts)
# subtr_parts_b = be.subtract(bg_parts, _mpt, len(_pythia_inserts))
# print('. nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
# subtr_parts = bgestim.subtracted_particles(bg_parts)
# ja.analyze_event(subtr_parts)
# for j in ja.jets:
# if matched_pt(j, sjet) > 0.5:
# tw.fill_branch('sjet_bg', j)
# tw.fill_branch('delta_pt', sjet.pt() - j.pt())
#ja.analyze_event(bg_parts)
#bg_signal_jets_b = fj.sorted_by_pt(ja.jets)
#for j in bg_signal_jets_b:
# if matched_pt(j, sjet) > 0.5:
# tw.fill_branch('sjet_hybrid', j)
# tw.fill_branch('delta_pt_hybrid', sjet.pt() - j.pt())
# subtr_parts_b = boltzmann_subtractor.subtracted_particles(bg_parts)
# _it1 = [p for p in be.subtract(bg_parts)]
# subtr_parts_b = be.subtract(_it1)
jabg.analyze_event(bg_parts)
bg_parts_noleads = remove_jets(bg_parts, jabg.jets[:2])
# print(bg_parts_noleads.size())
#for j in jabg.jets:
# pass
# continue
# if i > 2:
# break
# _tmp = [p.pt() for p in j.constituents()]
# reduce_total_pt = reduce_total_pt + sum(_tmp)
#print ('reduce total pt by', reduce_total_pt)
# subtr_parts_b = be.subtract(bg_parts)
# print(be.get_formula())
# print('nparts_bg=', len(bg_parts), 'nparts_subtr=', len(subtr_parts_b), 'nparts_pythia=', len(tmp))
_mpt = mean_pt(bg_parts_noleads)
# subtr_parts_b = be.subtract(bg_parts, _mpt, len(_pythia_inserts))
# subtr_parts_b = be.subtract(bg_parts, _mpt, len(bg_parts) - len(bg_parts_noleads))
subtr_parts_b = be.subtract(bg_parts, _mpt, len(bg_parts_noleads))
print('. nparts_bg=', len(bg_parts), 'nparts_pass='******' ', be.get_formula())
_indexes_subtr_parts_b = [p.user_index() for p in subtr_parts_b]
print ('remaining indexes:', len(_indexes_subtr_parts_b), sorted(_indexes_subtr_parts_b))
ja.analyze_event(subtr_parts_b)
bg_signal_jets_b = fj.sorted_by_pt(ja.jets)
w = pythia.info.weight()
s = pythia.info.sigmaGen()
twj.fill_branch('sjet', signal_jets)
twj.fill_tree()
print( 'py jets=', len([j for j in signal_jets if j.pt() > 0.1 ]),
'sub jets=', len([j for j in bg_signal_jets_b if j.pt() > 0.1]))
_tmp = [fill_tree_matched(sjet, j, twjm, sd, 0, iev, weight=w, sigma=s)
for sjet in signal_jets
for j in bg_signal_jets_b]
twp.fill_branch('p', bg_parts)
#twp.fill_branch('sp', subtr_parts)
twp.fill_branch('sp_b', subtr_parts_b)
twp.fill_tree()
pbar.close()
fout.cd()
fout.Write()
fout.Close()
print('[i] written', fout.GetName())
class Data2MP4(MPBase):
def __init__(self, **kwargs):
self.configure_from_args(file_list='PbPb_file_list.txt',
n_events=500,
output='output',
run_jetfinder=False)
super(Data2MP4, self).__init__(**kwargs)
self.data_io = DataIO(file_list=self.file_list)
#self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -1, 1, 51, 0., 2.*ROOT.TMath.Pi())
#self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -1, 1, 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
#self.hetaphi_jet = ROOT.TH2F('_hetaphi_jet', '_hetaphi_jet;#eta;#varphi', 51, -1, 1, 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
#self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 101, -ROOT.TMath.Pi(), ROOT.TMath.Pi(), 101, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -1,
1, 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.hetaphi_jet = ROOT.TH2F('_hetaphi_jet',
'_hetaphi_jet;#eta;#varphi', 51,
-ROOT.TMath.Pi(), ROOT.TMath.Pi(), 51,
-ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.mean_eta = ROOT.TH1F('_hmean_eta', '_hmean_eta;<#eta>', 101, -0.3,
0.3)
# self.mean_phi = ROOT.TH1F('_hmean_phi', '_hmean_phi;<#varphi>', 101, ROOT.TMath.Pi()-0.3, ROOT.TMath.Pi()+0.3)
self.mean_phi = ROOT.TH1F('_hmean_phi', '_hmean_phi;<#varphi>', 101,
-0.3, +0.3)
self.mean_e = ROOT.TH1F('_hmean_e', '_hmean_e;<E (GeV)>', 100, 0, 2)
self.tc = None
if self.run_jetfinder:
fj.ClusterSequence.print_banner()
def run(self):
_err_level_tmp = ROOT.gErrorIgnoreLevel
if self.tc is None:
self.tc = ROOT.TCanvas('_tc', '_tc', 800, 800)
self.tc.Divide(2, 2)
ROOT.gErrorIgnoreLevel = ROOT.kWarning
for iev in tqdm.tqdm(range(self.n_events)):
parts = self.data_io.load_event(offset=10000)
if self.run_jetfinder:
jet_R0 = 0.4
jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
jets = fj.sorted_by_pt(jet_def(parts))
if len(jets) > 0:
jet = jets[0]
else:
jet = fj.PseudoJet()
else:
jet = fj.PseudoJet()
self.hetaphi.Reset()
self.hetaphi_jet.Reset()
#_tmp = [self.hetaphi.Fill(p.eta(), jet.delta_phi_to(p), p.e()) for p in parts]
_tmp = [
self.hetaphi.Fill(p.eta(),
p.phi() - ROOT.TMath.Pi(), p.e())
for p in parts
]
if self.run_jetfinder:
# _tmp = [self.hetaphi_jet.Fill(p.eta(), jet.delta_phi_to(p), p.e()) for p in jet.constituents()]
_tmp = [
self.hetaphi_jet.Fill(p.eta(),
p.phi() - ROOT.TMath.Pi(), p.e())
for p in jet.constituents()
]
_e = [p.e() for p in parts]
_phi = [p.phi() - ROOT.TMath.Pi() for p in parts]
# _phi = [jet.delta_phi_to(p) for p in parts]
_eta = [p.eta() for p in parts]
# self.mean_eta.Fill(self.hetaphi.GetMean(1))
self.mean_eta.Fill(np.mean(_eta))
# self.mean_phi.Fill(self.hetaphi.GetMean(2))
self.mean_phi.Fill(np.mean(_phi))
self.mean_e.Fill(np.mean(_e))
self.tc.cd()
# self.hetaphi.Scale(1./self.hetaphi.Integral())
self.hetaphi.SetMaximum(10.)
self.hetaphi.SetMinimum(0.)
self.tc.cd(1)
self.hetaphi.Draw('colz')
# self.hetaphi.Draw('lego2')
self.hetaphi_jet.SetLineColor(ROOT.kRed)
self.hetaphi_jet.SetFillColor(ROOT.kRed)
self.hetaphi_jet.SetFillStyle(1001)
self.hetaphi_jet.Draw('cont3 same')
self.tc.cd(2)
self.mean_e.Draw()
self.tc.cd(3)
self.mean_phi.Draw()
self.tc.cd(4)
self.mean_eta.Draw()
self.tc.SaveAs('_{}_{}.png'.format(self.output, iev), '.png')
ROOT.gErrorIgnoreLevel = _err_level_tmp
def save_mp4(self):
# os.system('ffmpeg -i _{}_%01d.png -vcodec mpeg4 -y {}.mp4'.format(self.output, self.output))
os.system(
'ffmpeg -r 4 -f image2 -s 1920x1080 -i _{}_%d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p {}.mp4'
.format(self.output, self.output))
class Data2MP4Morph(MPBase):
def __init__(self, **kwargs):
self.configure_from_args(file_list='PbPb_file_list.txt',
n_events=500,
output='output')
super(Data2MP4Morph, self).__init__(**kwargs)
self.data_io = DataIO(file_list=self.file_list)
# self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 51, -ROOT.TMath.Pi(), ROOT.TMath.Pi(), 51, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
# self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#eta;#varphi', 101, -1, 1, 101, -ROOT.TMath.Pi(), +ROOT.TMath.Pi())
self.hetaphi = ROOT.TH2F('_hetaphi', '_hetaphi;#varphi;#eta', 101,
-ROOT.TMath.Pi(), +ROOT.TMath.Pi(), 101, -1,
1)
self.tc = None
self.iframe = 0
self.pythia = PythiaJetty(pthatmin=100,
eta_max=1,
jet_pt_min=20,
jet_R0=0.4)
def make_a_frame(self, h):
self.tc.cd()
for ibx in range(1, h.GetNbinsX() + 1):
for iby in range(1, h.GetNbinsY() + 1):
if h.GetBinContent(ibx, iby) < 1e-1:
h.SetBinContent(ibx, iby, 1e-1)
h.SetMaximum(100.)
h.SetMinimum(1e-1)
# h.Draw('colz')
h.SetName('{}_f{}'.format(h.GetName().split('_f')[0], self.iframe))
h.Draw('LEGO2 FB BB')
# h.Draw('surf2 FB BB')
# h.Draw('CONT3 COLZ')
# h.Draw('CYL LEGO2 COLZ')
h.SetLineColor(ROOT.kWhite)
h.SetLineColorAlpha(ROOT.kWhite, 1.)
# h.Draw('CYL SURF2')
# h.Draw('SURF2')
# h.Draw('PSR LEGO2')
# self.tc.SetLogz()
self.tc.SaveAs('_{}_{}.png'.format(self.output, self.iframe), '.png')
self.iframe = self.iframe + 1
def morph(self, hfrom, hto, nframes):
_hmorph = hfrom.Clone("_hmorph")
_hmorph.SetDirectory(0)
_hdeltas = hto.Clone('_delta')
_hdeltas.Add(hfrom, -1.)
for ifr in tqdm.tqdm(range(1, nframes + 1)):
fr = 1.0 * ifr / nframes
for ibx in range(1, hfrom.GetNbinsX() + 1):
for iby in range(1, hfrom.GetNbinsY() + 1):
v = hfrom.GetBinContent(
ibx, iby) + fr * _hdeltas.GetBinContent(ibx, iby)
_hmorph.SetBinContent(ibx, iby, v)
self.make_a_frame(_hmorph)
def run(self):
_err_level_tmp = ROOT.gErrorIgnoreLevel
if self.tc is None:
self.tc = ROOT.TCanvas('_tc', '_tc', 800, 800)
ROOT.gErrorIgnoreLevel = ROOT.kWarning
_hetaphi_next = self.hetaphi.Clone("_next")
for iev in tqdm.tqdm(range(self.n_events)):
parts = self.data_io.load_event(offset=10000)
_hetaphi_next.Reset()
#_tmp = [_hetaphi_next.Fill(p.eta(), p.phi()-ROOT.TMath.Pi(), p.e()) for p in parts]
#_tmp = [_hetaphi_next.Fill(p.phi()-ROOT.TMath.Pi(), p.eta(), p.e()) for p in parts]
pyev = None
if self.pythia:
pyev = self.pythia.get_event()
if pyev:
_tmp = [
_hetaphi_next.Fill(p.phi() - ROOT.TMath.Pi(), p.eta(),
p.e()) for p in pyev
]
if iev > 0:
self.morph(self.hetaphi, _hetaphi_next, 10)
self.hetaphi.Reset()
#_tmp = [self.hetaphi.Fill(p.phi()-ROOT.TMath.Pi(), p.eta(), p.e()) for p in parts]
if pyev:
_tmp = [
self.hetaphi.Fill(p.phi() - ROOT.TMath.Pi(), p.eta(),
p.e()) for p in pyev
]
self.make_a_frame(self.hetaphi)
ROOT.gErrorIgnoreLevel = _err_level_tmp
def save_mp4(self):
# os.system('ffmpeg -r 10 -i _{}_%d.png -vcodec mpeg4 -y {}.mp4'.format(self.output, self.output))
# good page: https://hamelot.io/visualization/using-ffmpeg-to-convert-a-set-of-images-into-a-video/
os.system(
'ffmpeg -r 60 -f image2 -s 1920x1080 -i _{}_%d.png -vcodec libx264 -crf 25 -pix_fmt yuv420p {}.mp4'
.format(self.output, self.output))