def match_dR(j, partons, drmatch=0.1): mps = [p for p in partons if j.delta_R(p) < drmatch] # for p in fj.sorted_by_pt(mps)[0]: p = fj.sorted_by_pt(mps)[0] pyp = pythiafjext.getPythia8Particle(p) # print(p, pyp.id(), pyp.isQuark(), pyp.isGluon()) return pyp.id(), pyp.isQuark(), pyp.isGluon()
def parton_splittings_parton_only(pythia, selected_jet, maxR=0.4): spinfo = SplitInfo(pythia) for p in fj.sorted_by_pt(selected_jet.constituents()): pyp = pythiafjext.getPythia8Particle(p) if pyp.isParton(): if p.delta_R(selected_jet) < maxR: spinfo.find_split(p.user_index()) print (spinfo) else: continue
def main(): parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__)) pyconf.add_standard_pythia_args(parser) _default_output_filename = os.path.basename(__file__).replace( ".py", "") + "_output.root" parser.add_argument('--output', default=_default_output_filename, type=str) parser.add_argument('--debug', default=0, type=int) args = parser.parse_args() # print the banner first fj.ClusterSequence.print_banner() print() # set up our jet definition and a jet selector jet_R0 = 0.6 jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) jet_selector = fj.SelectorPtMin(2.0) & fj.SelectorAbsEtaMax(2) print(jet_def) jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, jet_R0) lund_gen = fjcontrib.LundGenerator(jet_def_lund) print(jet_def_lund) print(lund_gen) outf = ROOT.TFile(args.output, 'recreate') outf.cd() t = ROOT.TTree('t', 't') tw = RTreeWriter(tree=t) # mycfg = ['PhaseSpace:pThatMin = 100'] mycfg = [] pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if args.nev < 100: args.nev = 100 for i in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue if args.debug: pwarning('-- event', i) # parts = pythiafjext.vectorize(pythia, True, -1, 1, False) parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, True) if args.debug > 5: parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kHadron], 0, True) if args.debug > 10: parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kAny], 0, True) if args.debug > 0: for p in parts: pypart = pythiafjext.getPythia8Particle(p) if pypart.name()[:2] == 'D0': pinfo(pypart.name(), pypart.id(), pypart.status(), 'final =?', pypart.isFinal()) jets = jet_selector(jet_def(parts)) for j in jets: isD0_lead = False lead_part = fj.sorted_by_E(j.constituents())[0] pypart = pythiafjext.getPythia8Particle(lead_part) if args.debug: pinfo('leading id is', pypart.id(), pypart.name(), 'jet', j) if abs(pypart.id()) == 421: # pinfo('leading D0') isD0_lead = True l = lund_gen.result(j) if len(l) > 0: tw.fill_branch('Epair', [s.pair().e() for s in l]) tw.fill_branch('z', [s.z() for s in l]) tw.fill_branch('kt', [s.kt() for s in l]) tw.fill_branch('delta', [s.Delta() for s in l]) tw.fill_branch('D0lead', isD0_lead) tw.fill_branch('lead_id', pypart.id()) tw.fill_tree() else: if args.debug: pwarning("len of a lund is less than 1?", len(l), l) pythia.stat() outf.Write() outf.Close() print('[i] written', outf.GetName())
def parton_splittings(pythia, selected_jet, maxR=0.4): spinfo = SplitInfo(pythia) for p in fj.sorted_by_pt(selected_jet.constituents()): pyp = pythiafjext.getPythia8Particle(p) spinfo.find_split(p.user_index()) print (spinfo)
def main(): parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__)) pyconf.add_standard_pythia_args(parser) # could use --py-seed parser.add_argument('--fj-R', help='jet finder R', default=0.8, type=float) parser.add_argument('--user-seed', help='pythia seed', default=-1, type=int) parser.add_argument('--output', default='{}.root'.format(os.path.basename(__file__)), type=str) parser.add_argument('--min-jet-pt', help='jet pt selection', default=50., type=float) parser.add_argument('--max-jet-pt', help='jet pt selection', default=1000., type=float) parser.add_argument('--npart-min', help='minimum npart in Argantyr', default=2, type=int) parser.add_argument('--as-data', help='write as data - tree naming convention', action='store_true', default=False) parser.add_argument('--run-number', help='set run number', default=1, type=int) parser.add_argument('--part-min-pt', help='minimum pt of a particle', default=0.15, type=float) args = parser.parse_args() if args.user_seed < 0: args.user_seed = -1 mycfg = [] else: pinfo('user seed for pythia', args.user_seed) # mycfg = ['PhaseSpace:pThatMin = 100'] mycfg = ['Random:setSeed=on', 'Random:seed={}'.format(args.user_seed)] pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if args.nev < 100: args.nev = 100 # print the banner first fj.ClusterSequence.print_banner() print() # set up our jet definition and a jet selector jet_R0 = args.fj_R jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) print(jet_def) # hadron level - ALICE max_eta_hadron = 3. pwarning('max eta for particles after hadronization set to', max_eta_hadron) parts_selector_h = fj.SelectorAbsEtaMax(max_eta_hadron) & fj.SelectorPtMin( args.part_min_pt) parts_selector_cent = fj.SelectorAbsEtaMax(5.) & fj.SelectorAbsEtaMin(3.) outf = ROOT.TFile(args.output, 'recreate') outf.cd() # t = ROOT.TTree('t', 't') # tw = RTreeWriter(tree=t) # tch = ROOT.TTree('tch', 'tch') # twch = RTreeWriter(tree=tch) te = ROOT.TTree('te', 'te') twe = RTreeWriter(tree=te) tdf = ROOT.TDirectoryFile('PWGHF_TreeCreator', 'PWGHF_TreeCreator') tdf.cd() if args.as_data: t_p = ROOT.TNtuple( 'tree_Particle', 'tree_Particle', 'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticlePID') else: t_p = ROOT.TNtuple( 'tree_Particle_gen', 'tree_Particle_gen', 'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticlePID') t_e = ROOT.TNtuple( 'tree_event_char', 'tree_event_char', 'run_number:ev_id:z_vtx_reco:is_ev_rej:weight:sigma:npart:nch:nchfwd:nchselect' ) run_number = args.run_number ev_id = 0 # event loop for iev in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue twe.clear() # tw.clear() # twch.clear() weight = pythia.info.weight() if args.py_PbPb: # from main113.cc # Also fill the number of (absorptively and diffractively) # wounded nucleaons. nw = pythia.info.hiinfo.nAbsTarg() + pythia.info.hiinfo.nDiffTarg( ) + pythia.info.hiinfo.nAbsProj() + pythia.info.hiinfo.nDiffProj() fill_ncoll_branches(pythia, twe) else: nw = 2 twe.fill_branch('nw', nw) twe.fill_branch('w', weight) sigma = pythia.info.sigmaGen() twe.fill_branch('sigma', sigma) # parts_pythia_h = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False) # parts_pythia_h_selected = parts_selector_h(parts_pythia_h) parts_pythia_ch = pythiafjext.vectorize_select( pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, True) parts_pythia_ch_selected = parts_selector_h(parts_pythia_ch) nch_total = len(parts_pythia_ch) twe.fill_branch('nch', nch_total) ncharged_fwd = len(parts_selector_cent(parts_pythia_ch)) twe.fill_branch('nchfwd', ncharged_fwd) ncharged_selected = len(parts_pythia_ch_selected) twe.fill_branch('nchselect', ncharged_selected) twe.fill_branch('iev', iev) ev_id = ev_id + 1 if args.py_PbPb and args.npart_min > nw: twe.fill_tree() continue for p in parts_pythia_ch_selected: pyp = pythiafjext.getPythia8Particle(p) t_p.Fill(float(run_number), float(ev_id), p.perp(), p.eta(), p.phi(), pyp.id()) t_e.Fill(float(run_number), float(ev_id), 0, 0, weight, sigma, nw, nch_total, ncharged_fwd, ncharged_selected) twe.fill_tree() pythia.stat() outf.Write() outf.Close()
def main(): parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__)) pyconf.add_standard_pythia_args(parser) _default_output_filename = os.path.basename(__file__).replace( ".py", "") + "_output.root" parser.add_argument('--output', default=_default_output_filename, type=str) parser.add_argument('--user-seed', help='pythia seed', default=1111, type=int) parser.add_argument('--debug', default=0, type=int) args = parser.parse_args() # jets # print the banner first # fj.ClusterSequence.print_banner() # print() # # set up our jet definition and a jet selector jet_R0 = 1.0 # jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) # print(jet_def) # acceptance # hadron level acceptamce max_eta_hadron = 10. from pyjetty.mputils import pwarning pwarning('max eta for particles after hadronization set to', max_eta_hadron) parts_selector_h = fj.SelectorAbsEtaMax(max_eta_hadron) jet_selector = fj.SelectorPtMin(1.0) & fj.SelectorPtMax( 100.0) & fj.SelectorAbsEtaMax(max_eta_hadron - 1.05 * jet_R0) # parton level acceptamce max_eta_parton = max_eta_hadron + 3. * jet_R0 pwarning('max eta for partons set to', max_eta_parton) parts_selector_p = fj.SelectorAbsEtaMax(max_eta_parton) # initialize pythia # mZ = pythia8.Pythia().particleData.m0(23) mZ = 91.188 beams_eCM = "Beams:eCM={}".format(mZ) mycfg = [ "PDF:lepton = off", # Allow no substructure in e+- beams: normal for corrected LEP data. "WeakSingleBoson:ffbar2gmZ = on", # Process selection. "23:onMode = off", # Switch off all Z0 decays and then switch back on those to quarks. "23:onIfAny = 1 2 3 4 5", "Beams:idA = 11", "Beams:idB = -11", beams_eCM, # LEP1 initialization at Z0 mass. "HadronLevel:all=off", # parton level first "PhaseSpace:bias2Selection=off" ] # this is ON by default in pyconf - not OK for these settings pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if not pythia: return # initialize ROOT output outf = ROOT.TFile(args.output, 'recreate') outf.cd() tdf = ROOT.TDirectoryFile('PWGHF_TreeCreator', 'PWGHF_TreeCreator') tdf.cd() t_p = ROOT.TNtuple( 'tree_Particle_P', 'tree_Particle_P', 'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticleID:ParticleIDabs:ParticleCharge:isGluon:isQuark' ) # t_p = ROOT.TNtuple('tree_Particle_gen', 'tree_Particle_gen', 'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticleID:ParticleIDabs:ParticleCharge') t_h = ROOT.TNtuple( 'tree_Particle_H', 'tree_Particle_H', 'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticleID:ParticleIDabs:ParticleCharge:isHadron:isLepton:isVisible' ) t_e = ROOT.TNtuple('tree_event_char', 'tree_event_char', 'run_number:ev_id:z_vtx_reco:is_ev_rej') if args.nev < 100: args.nev = 100 run_number = args.user_seed # main loop for i in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue if args.debug: pwarning('-- event', i) #select particles parts_pythia_p = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, True) parts_pythia_p_selected = parts_selector_p(parts_pythia_p) # hadronize hstatus = pythia.forceHadronLevel() if not hstatus: pwarning('forceHadronLevel false event', iev) continue parts_pythia_h = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, True) parts_pythia_h_selected = parts_selector_h(parts_pythia_h) # charged hadrons/particles only # parts_pythia_h = pythiafjext.vectorize_select(pythia, [pythiafjext.kHadron, pythiafjext.kCharged]) # parts_pythia_hch = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, True) # parts_pythia_hch_selected = parts_selector_h(parts_pythia_hch) # stream to trees ev_id = i t_e.Fill(run_number, ev_id, 0, 0) _tmp = [ t_p.Fill( run_number, ev_id, p.perp(), p.eta(), p.phi(), pythiafjext.getPythia8Particle(p).id(), pythiafjext.getPythia8Particle(p).idAbs(), pythiafjext.getPythia8Particle(p).charge(), pythiafjext.getPythia8Particle(p).isGluon(), pythiafjext.getPythia8Particle(p).isQuark(), ) for p in parts_pythia_p ] _tmp = [ t_h.Fill(run_number, ev_id, p.perp(), p.eta(), p.phi(), pythiafjext.getPythia8Particle(p).id(), pythiafjext.getPythia8Particle(p).idAbs(), pythiafjext.getPythia8Particle(p).charge(), pythiafjext.getPythia8Particle(p).isHadron(), pythiafjext.getPythia8Particle(p).isLepton(), pythiafjext.getPythia8Particle(p).isVisible()) for p in parts_pythia_h ] pythia.stat() outf.Write() outf.Close()
def get_pid(p): py8p = pythiafjext.getPythia8Particle(p) pid = 0 if py8p: pid = py8p.id() return pid
def main(): parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__)) pyconf.add_standard_pythia_args(parser) parser.add_argument('--charged', default=False, action='store_true') parser.add_argument('--neutral', default=False, action='store_true') parser.add_argument('--partons', default=False, action='store_true') parser.add_argument('--quarks', default=False, action='store_true') parser.add_argument('--diquarks', default=False, action='store_true') parser.add_argument('--gluons', default=False, action='store_true') parser.add_argument('--leptons', default=False, action='store_true') parser.add_argument('--Ncharged', default=False, action='store_true') parser.add_argument('--Nneutral', default=False, action='store_true') parser.add_argument('--Npartons', default=False, action='store_true') parser.add_argument('--Nquarks', default=False, action='store_true') parser.add_argument('--Ndiquarks', default=False, action='store_true') parser.add_argument('--Ngluons', default=False, action='store_true') parser.add_argument('--Nleptons', default=False, action='store_true') args = parser.parse_args() mycfg = [] pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if args.nev < 1: args.nev = 1 part_selection = [pythiafjext.kFinal] if args.charged: part_selection.append(pythiafjext.kCharged) if args.neutral: part_selection.append(pythiafjext.kNeutral) if args.partons: part_selection.append(pythiafjext.kParton) if args.quarks: part_selection.append(pythiafjext.kQuark) if args.diquarks: part_selection.append(pythiafjext.kDiquark) if args.gluons: part_selection.append(pythiafjext.kGluon) if args.leptons: part_selection.append(pythiafjext.kLepton) if args.Ncharged: part_selection.append(-pythiafjext.kCharged) if args.Nneutral: part_selection.append(-pythiafjext.kNeutral) if args.Npartons: part_selection.append(-pythiafjext.kParton) if args.Nquarks: part_selection.append(-pythiafjext.kQuark) if args.Ndiquarks: part_selection.append(-pythiafjext.kDiquark) if args.Ngluons: part_selection.append(-pythiafjext.kGluon) if args.Nleptons: part_selection.append(-pythiafjext.kLepton) for iev in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue parts = [] parts = pythiafjext.vectorize_select(pythia, part_selection, 0, True) print('[i] number of particles selected:', len(parts)) for p in parts: pypart = pythiafjext.getPythia8Particle(p) print_pythia_particle(pypart)
def main(): parser = argparse.ArgumentParser(description='pythia8 in python', prog=os.path.basename(__file__)) parser.add_argument('--hepmc', help='write the hepmc files', default=False, action='store_true') parser.add_argument('--ml', help='write the ml format root files', default=False, action='store_true') parser.add_argument('-c', '--nperfile', help='nevents per output file', default=1000, type=int) pyconf.add_standard_pythia_args(parser) args = parser.parse_args() if args.hepmc: hepmc2_output_base = "pythia_gen_qorg.hepmc" hepmc2_output_name = None hepmc2_writer = None hepmc2_fileno = 0 if args.ml: if root_avail is None: print('[error] unable to load ROOT so --ml is defunct') args.ml = False ml_root_output_base = "pythia_gen_qorg.root" ml_root_output_name = None ml_root_file = None ml_root_ntuple_parts = None ml_root_ntuple_ev = None ml_root_fileno = 0 if args.hepmc is False and args.ml is False: print('[error] one of the --hepmc or --ml is required') print(' --help for more options') return mycfg = [] pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if args.nev < 10: args.nev = 10 run_number = args.py_seed event_number = 0 pbar = tqdm.tqdm(range(args.nev)) for i in pbar: if not pythia.next(): pbar.update(-1) continue event_number = event_number + 1 if i == 0 or i % args.nperfile == 0: if args.hepmc: if hepmc2_writer: del hepmc2_writer hepmc2_output_name = format_output_file( hepmc2_output_base, hepmc2_fileno, args) # print('[i] new file', hepmc2_output_name) hepmc2_writer = pythiaext.Pythia8HepMC2Wrapper( hepmc2_output_name) hepmc2_fileno = hepmc2_fileno + 1 if args.ml: if ml_root_file: if ml_root_ntuple_parts: ml_root_ntuple_parts.Write() ml_root_ntuple_ev.Write() ml_root_file.Write() ml_root_file.Purge() ml_root_file.Close() ml_root_file = None ml_root_ntuple = None ml_root_output_name = format_output_file( ml_root_output_base, ml_root_fileno, args) ml_root_file = ROOT.TFile(ml_root_output_name, 'recreate') ml_root_ntuple_parts = ROOT.TNtuple( 'tree_Particle_gen', 'particles from PYTHIA8 - {} jets'.format( parton_type_from_args(args)), 'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticlePID' ) ml_root_ntuple_ev = ROOT.TNtuple( 'tree_Event_gen', 'event info from PYTHIA8 - {} jets'.format( parton_type_from_args(args)), 'run_number:ev_id:xsec:code') ml_root_fileno = ml_root_fileno + 1 if args.hepmc: hepmc2_writer.fillEvent(pythia) if args.ml: parts_pythia_h = pythiafjext.vectorize_select( pythia, [pythiafjext.kFinal], 0, True) ml_root_ntuple_ev.Fill(run_number, event_number, pythia.info.sigmaGen(), pythia.info.code()) for p in parts_pythia_h: ml_root_ntuple_parts.Fill( run_number, event_number, p.pt(), p.eta(), p.phi(), pythiafjext.getPythia8Particle(p).id()) pythia.stat() pythia.settings.writeFile( format_output_file('pythia_gen_qorg.cmnd', -1, args)) if args.ml: if ml_root_file: if ml_root_ntuple_parts: ml_root_ntuple_parts.Write() ml_root_ntuple_ev.Write() ml_root_file.Write() ml_root_file.Purge() ml_root_file.Close()