def __init__(self, outputname, enable_aa_trees=False): self.outf = ROOT.TFile(outputname, 'recreate') self.outf.cd() self.t = ROOT.TTree('t', 't') self.tw = RTreeWriter(tree=self.t) self.taaw = None self.tembw = None self.tembwp = None if enable_aa_trees: self.taa = ROOT.TTree('taa', 'taa') self.taaw = RTreeWriter(tree=self.taa) self.temb = ROOT.TTree('temb', 'temb') self.tembw = RTreeWriter(tree=self.temb) self.tembp = ROOT.TTree('tembp', 'tembp') self.tembwp = RTreeWriter(tree=self.tembp) self.jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0) self.dy_groomer_alphas = [0.1, 1.0, 2.0] self.sd_zcuts = [0.1, 0.2, 0.3, 0.4] self.gshop = fjcontrib.GroomerShop(self.jet_def_lund) self.gshop_emb = fjcontrib.GroomerShop(self.jet_def_lund) print(self.gshop.description()) self.sds = [] sd01 = fjcontrib.SoftDrop(0, 0.1, 1.0) sd02 = fjcontrib.SoftDrop(0, 0.2, 1.0) self.sds.append(sd01) self.sds.append(sd02)
def __init__(self, **kwargs): self.configure_from_args( fout=None, output='hjet_tree.root', trigger_range = [6, 7], jet_ana = None) super(HJetTree, self).__init__(**kwargs) if self.fout is None: self.fout = r.TFile(self.output, 'RECREATE') self.fout.cd() self.tree_name='tjet_0{}_{}_{}'.format(int(10*self.jet_ana.jet_R), int(self.trigger_range[0]), int(self.trigger_range[1])) self.tree_writer = RTreeWriter(tree_name=self.tree_name, fout=self.fout)
def initialize_output(self, output_name=None): if output_name: if self.output_filename != output_name: self.output_filename = output_name if self.outf: if self.outf.GetName() != self.output_filename: pinfo('closing output file', self.outf.GetName()) self.outf.Write() self.outf.Close() self.outf = None else: return True if self.outf is None: self.outf = ROOT.TFile(self.output_filename, 'recreate') self.outf.cd() self.tdet = ROOT.TTree('tdet', 'tdet') self.twdet = RTreeWriter( tree=self.tdet, name='Output Tree detector level pp simulation') self.tpp = ROOT.TTree('tpp', 'tpp') self.twpp = RTreeWriter(tree=self.tpp, name='Output Tree pp simulation') self.th = ROOT.TTree('th', 'th') self.twh = RTreeWriter( tree=self.th, name='Output Tree pp simulation embedded into PbPb') pinfo('new output file', self.outf.GetName())
def __init__(self, **kwargs): self.configure_from_args(jet_particle_eta_max = 0.9, output='hjet.root', trigger_ranges = [ [0, 1e3], [6, 7] , [12, 22], [20, 30] ], jet_Rs = [0.2, 0.4, 0.6]) super(HJetAnalysis, self).__init__(**kwargs) self.fout = r.TFile(self.output, 'RECREATE') self.fout.cd() self.event_output = RTreeWriter(tree_name='tev', fout=self.fout) self.hmV0M = r.TH1F('hmV0M', 'hmV0M', 1000, 0, 1000) self.hmV0A = r.TH1F('hmV0A', 'hmV0A', 1000, 0, 1000) self.hmV0C = r.TH1F('hmV0C', 'hmV0C', 1000, 0, 1000) self.jet_particle_selector = fj.SelectorAbsEtaMax(self.jet_particle_eta_max) self.jet_ans = [] self.hjet_ts = [] for jR in self.jet_Rs: j_ana = JetAnalysis(jet_R=jR, particle_eta_max=self.jet_particle_eta_max) self.jet_ans.append(j_ana) for tr in self.trigger_ranges: hjet_tree = HJetTree(fout=self.fout, trigger_range=tr, jet_ana=j_ana) self.hjet_ts.append(hjet_tree)
class EmbeddingOutput(MPBase): def __init__(self, **kwargs): self.configure_from_args(args=None) super(EmbeddingOutput, self).__init__(**kwargs) self.copy_attributes(self.args) self.outf = None self.sd = fjcontrib.SoftDrop(0, self.sd_zcut, self.jetR) def initialize_output(self, output_name=None): if output_name: if self.output_filename != output_name: self.output_filename = output_name if self.outf: if self.outf.GetName() != self.output_filename: pinfo('closing output file', self.outf.GetName()) self.outf.Write() self.outf.Close() self.outf = None else: return True if self.outf is None: self.outf = ROOT.TFile(self.output_filename, 'recreate') self.outf.cd() self.tdet = ROOT.TTree('tdet', 'tdet') self.twdet = RTreeWriter( tree=self.tdet, name='Output Tree detector level pp simulation') self.tpp = ROOT.TTree('tpp', 'tpp') self.twpp = RTreeWriter(tree=self.tpp, name='Output Tree pp simulation') self.th = ROOT.TTree('th', 'th') self.twh = RTreeWriter( tree=self.th, name='Output Tree pp simulation embedded into PbPb') pinfo('new output file', self.outf.GetName()) def close(self): if self.outf: pinfo('closing output file', self.outf.GetName()) self.outf.Write() self.outf.Close() self.outf = None def _fill_det_level(self, jet): self.twdet.fill_branch('pt_det', jet.pt()) self.twdet.fill_branch('pt_phi', jet.phi()) self.twdet.fill_branch('pt_eta', jet.eta()) def fill_det_level(self, iev=-1, jets=[]): if len(jets) > 0: self.twdet.fill_branch('iev', iev) _tmp = [self._fill_det_level(j) for j in jets] self.twdet.fill_tree() def fill_pp_pairs(self, iev, jm): self.twpp.fill_branch('iev', iev) self.twpp.fill_branch('det', jm[0]) self.twpp.fill_branch('part', jm[1]) self.twpp.fill_branch('dpt', jm[0].pt() - jm[1].pt()) self.twpp.fill_branch('dR', jm[0].delta_R(jm[1])) self.twpp.fill_tree() def fill_emb_3(self, iev, jm): # pdebug('@fill: jm[0]', jm[0], 'jm[1]', jm[1], 'jm[2]', jm[2]) self.twh.fill_branch('iev', iev) self.twh.fill_branch('det', jm[0]) self.twh.fill_branch('part', jm[1]) self.twh.fill_branch('hybr', jm[2]) self.twh.fill_branch('dpt_pp', jm[0].pt() - jm[1].pt()) self.twh.fill_branch('dpt_emb', jm[2].pt() - jm[0].pt()) self.twh.fill_branch('dR_pp', jm[0].delta_R(jm[1])) self.twh.fill_branch('dR_emb', jm[0].delta_R(jm[2])) sd0 = self.sd.result(jm[0]) self.twh.fill_branch('sd_det', sd0) sd0_pe1 = fj.PseudoJet() sd0_pe2 = fj.PseudoJet() sd0_has_parents = sd0.has_parents(sd0_pe1, sd0_pe2) self.twh.fill_branch('sd_det_p1', sd0_pe1) self.twh.fill_branch('sd_det_p2', sd0_pe2) sdi0 = fjcontrib.get_SD_jet_info(sd0) self.twh.fill_branch('sd_det_zg', sdi0.z) self.twh.fill_branch('sd_det_Rg', sdi0.dR) sd1 = self.sd.result(jm[1]) self.twh.fill_branch('sd_part', sd1) sd1_pe1 = fj.PseudoJet() sd1_pe2 = fj.PseudoJet() sd1_has_parents = sd1.has_parents(sd1_pe1, sd1_pe2) self.twh.fill_branch('sd_part_p1', sd1_pe1) self.twh.fill_branch('sd_part_p2', sd1_pe2) sdi1 = fjcontrib.get_SD_jet_info(sd1) self.twh.fill_branch('sd_part_zg', sdi1.z) self.twh.fill_branch('sd_part_Rg', sdi1.dR) sd2 = self.sd.result(jm[2]) self.twh.fill_branch('sd_emb', sd2) sd2_pe1 = fj.PseudoJet() sd2_pe2 = fj.PseudoJet() sd2_has_parents = sd2.has_parents(sd2_pe1, sd2_pe2) self.twh.fill_branch('sd_emb_p1', sd2_pe1) self.twh.fill_branch('sd_emb_p2', sd2_pe2) sdi2 = fjcontrib.get_SD_jet_info(sd2) self.twh.fill_branch('sd_emb_zg', sdi2.z) self.twh.fill_branch('sd_emb_Rg', sdi2.dR) m02_1 = -1 m02_2 = -1 if sd0_has_parents and sd2_has_parents: m02_1 = fjtools.matched_pt(sd2_pe1, sd0_pe1) m02_2 = fjtools.matched_pt(sd2_pe2, sd0_pe2) self.twh.fill_branch('sd_det_emb_mpt1', m02_1) self.twh.fill_branch('sd_det_emb_mpt2', m02_2) self.twh.fill_branch('sd_det_split', sd0_has_parents) self.twh.fill_branch('sd_part_split', sd1_has_parents) self.twh.fill_branch('sd_emb_split', sd2_has_parents) self.twh.fill_tree()
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 generate(): parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__)) parser.add_argument( '-d', '--output-dir', help='output directory name - use with default file name', default='.', type=str) parser.add_argument('-o', '--output', help='output file name', default='', type=str) parser.add_argument('--overwrite', default=False, action='store_true') parser.add_argument('--jet-R', help='jet finder R', default=0.4, type=float) parser.add_argument('--charged', default=False, action='store_true') parser.add_argument('--runid', default=0, type=int) parser.add_argument('--tranges', default='6-7', help='hadron trigger ranges min-max,min1-max1,...', type=str) group = parser.add_mutually_exclusive_group() group.add_argument('--inel', default=False, action='store_true') group.add_argument('--hard', default=False, action='store_true') pyconf.add_standard_pythia_args(parser) args = parser.parse_args() # mycfg = ['PhaseSpace:pThatMin = 6'] mycfg = [] if '.root' not in args.output: rs = '0{}'.format(int(args.jet_R * 10)) if args.jet_R >= 1: rs = '{}'.format(int(args.jet_R)) args.output = '{}/h_jet_R{}'.format(args.output_dir, rs) if args.charged: args.output = '{}/h_jet_ch_R{}'.format(args.output_dir, rs) if args.tranges: args.output += '_tranges_{}'.format(args.tranges.replace(',', "_")) if args.runid >= 0: args.output += '_runid_{}'.format(args.runid) args.py_seed = 1000 + args.runid if args.py_noMPI: args.output += '_noMPI' if args.py_noISR: args.output += '_noISR' if args.py_noHadron: args.output += '_noHadr' if args.py_minbias > 0: args.output += '_minbias' if args.py_pthatmin > 0: args.output += '_pthatmin_{}'.format(args.py_pthatmin) if args.inel: args.output += '_inel' mycfg.append("SoftQCD:inelastic = on") mycfg.append("HardQCD:all = off") else: if args.hard: args.output += '_biasref_{}'.format(args.py_biasref) args.output += '_hard' mycfg.append("HardQCD:all = on") else: mycfg.append("HardQCD:all = on") args.output += '_hard' args.output += '.root' print(args) if os.path.exists(args.output): if not args.overwrite: print('[w] output file', args.output, 'exists - skipping.') return print('[i] output file', args.output) # print the banner first fj.ClusterSequence.print_banner() print() #and in the code i do not allow decay of the following particles: mycfg.append("310:mayDecay = off") # //K0s mycfg.append("3122:mayDecay = off") # //labda0 mycfg.append("3112:mayDecay = off") # //sigma- mycfg.append("3212:mayDecay = off") # //sigma0 mycfg.append("3222:mayDecay = off") # //sigma+ mycfg.append("3312:mayDecay = off") # //xi- mycfg.append("3322:mayDecay = off") # //xi+ mycfg.append("3334:mayDecay = off") # //omega- print('[i] additional settings', mycfg) pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if not pythia: return if args.nev < 1: args.nev = 1 tpairs = [] st = args.tranges.split(',') try: for _s in st: t0 = float(_s.split('-')[0]) t1 = float(_s.split('-')[1]) tpairs.append([t0, t1]) except: print( '[e] something is not quite right with trigger ranges... stop here.' ) jet_particle_eta_max = 0.9 jet_particle_selector = fj.SelectorAbsEtaMax(jet_particle_eta_max) v0det = V0Detector() j_ana = JetAnalysis(jet_R=args.jet_R, particle_eta_max=jet_particle_eta_max) fout = r.TFile(args.output, 'RECREATE') fout.cd() hmV0M = r.TH1F('hmV0M', 'hmV0M', 1000, 0, 1000) hmV0A = r.TH1F('hmV0A', 'hmV0A', 1000, 0, 1000) hmV0C = r.TH1F('hmV0C', 'hmV0C', 1000, 0, 1000) event_output = RTreeWriter(tree_name='evT', fout=fout) jet_output = RTreeWriter(tree_name='jetT', fout=fout) hjet_sets = [] for t in tpairs: hjet_output = RTreeWriter(tree_name='hjetT_{}_{}'.format( int(t[0]), int(t[1])), fout=fout) hjet = HJet(trigger_range=[t[0], t[1]]) print(hjet) hjset = HJetSet(hjet, hjet_output) hjet_sets.append(hjset) for ev_id in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue if args.charged: parts = pythiafjext.vectorize_select( pythia, [pythiafjext.kFinal, pythiafjext.kCharged]) else: parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal]) event_output.clear() jet_output.clear() _tmp = [s.hjet_output.clear() for s in hjet_sets] jet_parts = jet_particle_selector(parts) ev_s = pythia.info.sigmaGen() ev_w = pythia.info.weight() ev_code = pythia.info.code() pthard = pythia.info.pTHat() mTot = len(parts) mCB = len(jet_parts) v0det.analyze_event(parts) event_output.fill_branch('ev_id', ev_id) event_output.fill_branch('weight', ev_w) event_output.fill_branch('sigma', ev_s) event_output.fill_branch('code', ev_code) event_output.fill_branch('pthard', pthard) event_output.fill_branch('mTot', mTot) event_output.fill_branch('mCB', mCB) v0det.fill_branches(event_output) event_output.fill_tree() hmV0M.Fill(v0det.V0_mult) hmV0A.Fill(v0det.V0A_mult) hmV0C.Fill(v0det.V0C_mult) if len(parts) < 1: continue if True in [s.analyze_event(jet_parts) for s in hjet_sets]: j_ana.analyze_event(jet_parts) jet_output.fill_branch('ev_id', ev_id) jet_output.fill_branch('weight', ev_w) jet_output.fill_branch('sigma', ev_s) jet_output.fill_branch('code', ev_code) jet_output.fill_branch('pthard', pthard) jet_output.fill_branch('mTot', mTot) jet_output.fill_branch('mCB', mCB) j_ana.fill_branches(jet_output) jet_output.fill_tree() for s in hjet_sets: if s.hjet.trigger_particle: s.hjet_output.fill_branch('ev_id', ev_id) s.hjet_output.fill_branch('weight', ev_w) s.hjet_output.fill_branch('sigma', ev_s) s.hjet_output.fill_branch('code', ev_code) s.hjet_output.fill_branch('pthard', pthard) s.hjet_output.fill_branch('mTot', mTot) s.hjet_output.fill_branch('mCB', mCB) s.hjet.fill_branches(s.hjet_output, j_ana.jets) v0det.fill_branches(s.hjet_output) j_ana.fill_branches(s.hjet_output) s.hjet_output.fill_tree() pythia.stat() fout.Write() fout.Close() print('[i] written', fout.GetName())
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('--ignore-mycfg', help="ignore some settings hardcoded here", default=False, action='store_true') parser.add_argument('--output', default="output.root", type=str) parser.add_argument('--alpha', default=0, type=float) parser.add_argument('--dRmax', default=0.0, type=float) parser.add_argument('--zcut', default=0.1, type=float) parser.add_argument('--overwrite', help="overwrite output", default=False, action='store_true') parser.add_argument('--embed', help='run embedding from a file list', default='', type=str) parser.add_argument('--SDsignal', help='embed only SD signal prongs', default=False, action='store_true') parser.add_argument('--SDsignal-single', help='embed only SD signal - only leading prong!', default=False, action='store_true') parser.add_argument('--efficiency', help='apply charged particle efficiency', default=False, action='store_true') parser.add_argument('--benchmark', help='benchmark pthat setting - 80 GeV', default=False, action='store_true') parser.add_argument('--csjet', help='constituent subtration jet-by-jet', default=False, action='store_true') args = parser.parse_args() if args.output == 'output.root': args.output = 'output_alpha_{}_dRmax_{}_SDzcut_{}.root'.format( args.alpha, args.dRmax, args.zcut) if args.py_seed >= 0: args.output = 'output_alpha_{}_dRmax_{}_SDzcut_{}_seed_{}.root'.format( args.alpha, args.dRmax, args.zcut, args.py_seed) if args.embed: args.output = args.output.replace('.root', '_emb.root') if args.efficiency: args.output = args.output.replace('.root', '_effi.root') if args.SDsignal: args.output = args.output.replace('.root', '_SDsignal.root') if args.SDsignal_single: args.output = args.output.replace('.root', '_SDsignal_single.root') if args.csjet: args.output = args.output.replace('.root', '_csjet.root') 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 # print the banner first fj.ClusterSequence.print_banner() print() # set up our jet definition and a jet selector jet_R0 = 0.4 jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) print(jet_def) 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 = args.zcut sd = fjcontrib.SoftDrop(0, sd_zcut, jet_R0) jarho = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta) ja = JetAnalysis(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta) be = None embd = None if len(args.embed) > 0: embd = DataBackgroundIO(file_list=args.embed) print(embd) else: be = BoltzmannEvent(mean_pt=0.6, multiplicity=2000 * max_eta * 2, max_eta=max_eta, max_pt=100) print(be) cs = None if args.dRmax > 0: cs = CEventSubtractor(alpha=args.alpha, max_distance=args.dRmax, max_eta=max_eta, bge_rho_grid_size=0.25, max_pt_correct=100) print(cs) csjet = None if args.csjet: csjet = CSubtractorJetByJet(max_eta=max_eta, bge_rho_grid_size=0.25) print(csjet) parts_selector = fj.SelectorAbsEtaMax(max_eta) if args.nev < 1: args.nev = 1 outf = ROOT.TFile(args.output, 'recreate') outf.cd() t = ROOT.TTree('t', 't') tw = RTreeWriter(tree=t) te = ROOT.TTree('te', 'te') twe = RTreeWriter(tree=te) # effi_pp = AliceChargedParticleEfficiency(csystem='pp') effi_PbPb = None if args.efficiency: effi_PbPb = AliceChargedParticleEfficiency(csystem='PbPb') print(effi_PbPb) ### EVENT LOOP STARTS HERE for iev in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue parts_pythia = pythiafjext.vectorize_select( pythia, [pythiafjext.kFinal, pythiafjext.kCharged]) parts_gen = parts_selector(parts_pythia) if effi_PbPb: parts = effi_PbPb.apply_efficiency(parts_gen) else: parts = parts_gen signal_jets = fj.sorted_by_pt(jet_selector(jet_def(parts))) if len(signal_jets) < 1: continue for sjet in signal_jets: if args.SDsignal or args.SDsignal_single: sd_sjet = sd.result(sjet) pe1 = fj.PseudoJet() pe2 = fj.PseudoJet() has_parents = sd_sjet.has_parents(pe1, pe2) if has_parents: jparts = fj.vectorPJ() pe1.set_user_index(0) pe2.set_user_index(1) if args.SDsignal_single: if pe1.pt() > pe2.pt(): jparts.push_back(pe1) else: jparts.push_back(pe2) else: jparts.push_back(pe1) jparts.push_back(pe2) sjets = fj.sorted_by_pt(jet_selector(jet_def(jparts))) if len(sjets) == 1: sjet = sjets[0] else: continue else: continue if embd: bg_parts = embd.load_event(offset=10000) # for p in bg_parts: # print(p.user_index()) else: bg_parts = be.generate(offset=10000) # for p in bg_parts: # print(p.user_index()) full_event = bg_parts tmp = [full_event.push_back(psj) for psj in sjet.constituents()] if cs: cs_parts = cs.process_event(full_event) rho = cs.bge_rho.rho() jarho.analyze_event(cs_parts) tmp = [ fill_tree_data(ej, twe, sd, rho, iev, pythia.info.weight(), pythia.info.sigmaGen()) for ej in jarho.jets ] tmp = [ fill_tree_matched(sjet, ej, tw, sd, rho, iev, pythia.info.weight(), pythia.info.sigmaGen()) for ej in jarho.jets ] else: jarho.analyze_event(full_event) rho = jarho.rho if csjet: #_csjet = fjcontrib.ConstituentSubtractor(jarho.bg_estimator) # subtr_jets = [_csjet.result(ej) for ej in jarho.jets] csjet.set_event_particles(full_event) #subtr_jets = [csjet.process_jet(ej) for ej in jarho.jets] #print ('jbyj cs', len(subtr_jets), 'from', len(jarho.jets)) #subtr_jets_wconstits = [_j for _j in subtr_jets if _j.has_constituents()] #for _j in subtr_jets_wconstits: # print(len(_j.constituents())) subtr_jets_wconstits = csjet.process_jets(jarho.jets) japerjet = JetAnalysisPerJet( jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta, input_jets=subtr_jets_wconstits) # for _j in japerjet.jets: # for _c in _j.constituents(): # if _c.user_index() >= 0: # print('user index kept?', _c.user_index()) # # else: # # print('user index kept?', _c.user_index(), _c.pt()) # _sd_j = sd.result(_j) # https://phab.hepforge.org/source/fastjetsvn/browse/contrib/contribs/RecursiveTools/trunk/Recluster.cc L 270 # tmp = [fill_tree_matched(sjet, ej, tw, sd, rho, iev, pythia.info.sigmaGen()) for ej in subtr_jets_wcs] tmp = [ fill_tree_data(ej, twe, sd, rho, iev, pythia.info.weight(), pythia.info.sigmaGen()) for ej in japerjet.jets ] tmp = [ fill_tree_matched(sjet, ej, tw, sd, rho, iev, pythia.info.weight(), pythia.info.sigmaGen()) for ej in japerjet.jets ] else: tmp = [ fill_tree_data(ej, twe, sd, rho, iev, pythia.info.weight(), pythia.info.sigmaGen()) for ej in jarho.jets ] tmp = [ fill_tree_matched(sjet, ej, tw, sd, rho, iev, pythia.info.weight(), pythia.info.sigmaGen()) for ej in jarho.jets ] pythia.stat() outf.Write() outf.Close() print('[i] written', outf.GetName())
def generate(): parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__)) parser.add_argument( '-d', '--output-dir', help='output directory name - use with default file name', default='.', type=str) parser.add_argument('-o', '--output', help='output file name', default='', type=str) parser.add_argument('--overwrite', default=False, action='store_true') parser.add_argument('--t-min', help='trigger pt min', default=6., type=float) parser.add_argument('--t-max', help='trigger pt max', default=7., type=float) parser.add_argument('--jet-R', help='jet finder R', default=0.4, type=float) parser.add_argument('--charged', default=False, action='store_true') parser.add_argument('--runid', default=0, type=int) pyconf.add_standard_pythia_args(parser) args = parser.parse_args() # mycfg = ['PhaseSpace:pThatMin = 6'] mycfg = [] if '.root' not in args.output: rs = '0{}'.format(int(args.jet_R * 10)) if args.jet_R >= 1: rs = '{}'.format(int(args.jet_R)) args.output = '{}/h_jet_R{}_t{}_{}'.format(args.output_dir, rs, args.t_min, args.t_max) if args.charged: args.output = '{}/h_jet_ch_R{}_t{}_{}'.format( args.output_dir, rs, int(args.t_min), int(args.t_max)) if args.py_pthatmin > 0: args.output += '_pthatmin{}'.format(int(args.py_pthatmin)) if args.py_noMPI: args.output += '_noMPI' if args.py_noISR: args.output += '_noISR' if args.py_noHadron: args.output += '_noHadr' if args.py_minbias > 0: args.output += '_minbias' if args.runid > 0: args.output += '_runid_{}'.format(args.runid) args.py_seed = 1000 + args.runid args.output += '.root' if os.path.exists(args.output): if not args.overwrite: print('[w] output file', args.output, 'exists - skipping.') return # print the banner first fj.ClusterSequence.print_banner() print() mycfg.append("SoftQCD:inelastic = on") # Andreas' recommendation #and in the code i do not allow decay of the following particles: mycfg.append("310:mayDecay = off") # //K0s mycfg.append("3122:mayDecay = off") # //labda0 mycfg.append("3112:mayDecay = off") # //sigma- mycfg.append("3212:mayDecay = off") # //sigma0 mycfg.append("3222:mayDecay = off") # //sigma+ mycfg.append("3312:mayDecay = off") # //xi- mycfg.append("3322:mayDecay = off") # //xi+ mycfg.append("3334:mayDecay = off") # //omega- print('[i] additional settings', mycfg) pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if not pythia: return if args.nev < 1: args.nev = 1 jet_particle_eta_max = 0.9 jet_particle_selector = fj.SelectorAbsEtaMax(jet_particle_eta_max) fout = r.TFile(args.output, 'RECREATE') fout.cd() hmV0M = r.TH1F('hmV0M', 'hmV0M', 1000, 0, 1000) hmV0A = r.TH1F('hmV0A', 'hmV0A', 1000, 0, 1000) hmV0C = r.TH1F('hmV0C', 'hmV0C', 1000, 0, 1000) event_output = RTreeWriter(tree_name='evT', fout=fout) jet_output = RTreeWriter(tree_name='jetT', fout=fout) hjet_output = RTreeWriter(tree_name='hjetT', fout=fout) v0det = V0Detector() j_ana = JetAnalysis(jet_R=args.jet_R, particle_eta_max=jet_particle_eta_max) hjet = HJet(trigger_range=[args.t_min, args.t_max]) for ev_id in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue if args.charged: parts = pythiafjext.vectorize_select( pythia, [pythiafjext.kFinal, pythiafjext.kCharged]) else: parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal]) event_output.clear() jet_output.clear() hjet_output.clear() jet_parts = jet_particle_selector(parts) ev_w = pythia.info.sigmaGen() ev_code = pythia.info.code() pthard = pythia.info.pTHat() mTot = len(parts) mCB = len(jet_parts) v0det.analyze_event(parts) event_output.fill_branch('ev_id', ev_id) event_output.fill_branch('weight', ev_w) event_output.fill_branch('code', ev_code) event_output.fill_branch('pthard', pthard) event_output.fill_branch('mTot', mTot) event_output.fill_branch('mCB', mCB) v0det.fill_branches(event_output) event_output.fill_tree() hmV0M.Fill(v0det.V0_mult) hmV0A.Fill(v0det.V0A_mult) hmV0C.Fill(v0det.V0C_mult) if len(parts) < 1: continue hjet.analyze_event(jet_parts) if hjet.trigger_particle: j_ana.analyze_event(jet_parts) jet_output.fill_branch('ev_id', ev_id) jet_output.fill_branch('weight', ev_w) jet_output.fill_branch('code', ev_code) jet_output.fill_branch('pthard', pthard) jet_output.fill_branch('mTot', mTot) jet_output.fill_branch('mCB', mCB) j_ana.fill_branches(jet_output) jet_output.fill_tree() hjet_output.fill_branch('ev_id', ev_id) hjet_output.fill_branch('weight', ev_w) hjet_output.fill_branch('code', ev_code) hjet_output.fill_branch('pthard', pthard) hjet_output.fill_branch('mTot', mTot) hjet_output.fill_branch('mCB', mCB) hjet.fill_branches(hjet_output, j_ana.jets) v0det.fill_branches(hjet_output) j_ana.fill_branches(hjet_output) hjet_output.fill_tree() pythia.stat() fout.Write() fout.Close() print('[i] written', fout.GetName())
def main(agrs): fj.ClusterSequence.print_banner() print() ja = JetAnalysis(jet_R=args.jet_R, jet_algorithm=fj.antikt_algorithm, particle_eta_max=args.max_eta) grfj = GridFastJet(grid_size=0.01, jet_R=args.jet_R, jet_algorithm=fj.antikt_algorithm, particle_eta_max=args.max_eta) print(grfj) grfjR = GridFastJetROOT(grid_size=0.01, jet_R=args.jet_R, jet_algorithm=fj.antikt_algorithm, particle_eta_max=args.max_eta) print(grfjR) pythia = None if args.pythia: mycfg = [] pythia = pyconf.create_and_init_pythia_from_args(args, mycfg) if not pythia: print("[e] pythia initialization failed.") return fout = ROOT.TFile(args.output, 'recreate') fout.cd() ts = RTreeWriter(tree_name='ts', fout=fout) tg = RTreeWriter(tree_name='tg', fout=fout) tgR = RTreeWriter(tree_name='tgR', fout=fout) for iev in tqdm.tqdm(range(args.nev)): if pythia: if not pythia.next(): continue parts_pythia = pythiafjext.vectorize_select( pythia, [pythiafjext.kFinal, pythiafjext.kCharged]) ja.analyze_event(parts=parts_pythia) if len(ja.jets) < 1: continue if ja.jets[0].pt() < 10: continue ts.fill_branch('j', ja.jets) ts.fill_branch('p', ja.particles) ts.fill_tree() if args.grid: grfj.analyze_event(parts=parts_pythia) tg.fill_branch('j', grfj.jets) tg.fill_branch('p', grfj.particles) tg.fill_tree() if args.gridR: grfjR.analyze_event(parts=parts_pythia) tgR.fill_branch('j', grfjR.jets) tgR.fill_branch('p', grfjR.particles) tgR.fill_tree() ts.write_and_close()
def main(): parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__)) parser.add_argument('--output', default="output.root", type=str) parser.add_argument('datalist', help='run through a file list', default='', type=str) parser.add_argument('--alpha', default=0, type=float) parser.add_argument('--dRmax', default=0.0, type=float) parser.add_argument('--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=-100, type=float) parser.add_argument('--nev', help='number of events to run', default=0, type=int) args = parser.parse_args() if args.output == 'output.root': args.output = 'output_data_alpha_{}_dRmax_{}_SDzcut_{}.root'.format( args.alpha, args.dRmax, args.zcut) if args.jetptcut > -100: args.output = 'output_data_alpha_{}_dRmax_{}_SDzcut_{}_jpt_{}.root'.format( args.alpha, args.dRmax, args.zcut, args.jetptcut) 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 # print the banner first fj.ClusterSequence.print_banner() print() # set up our jet definition and a jet selector jet_R0 = 0.4 jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) print(jet_def) if args.benchmark: 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: jet_selector = fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0) sd_zcut = args.zcut sd = fjcontrib.SoftDrop(0, sd_zcut, jet_R0) ja = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta) data = DataIO(file_list=args.datalist) print(data) cs = None if args.dRmax > 0: cs = CEventSubtractor(alpha=args.alpha, max_distance=args.dRmax, max_eta=max_eta, bge_rho_grid_size=0.25, max_pt_correct=100) print(cs) parts_selector = fj.SelectorAbsEtaMax(max_eta) outf = ROOT.TFile(args.output, 'recreate') outf.cd() t = ROOT.TTree('t', 't') tw = RTreeWriter(tree=t) # need to change this for data to drive... delta_t = 1e-6 start_t = time.time() iev = 0 while data.load_event(): iev = iev + 1 _data_parts = data.particles if cs: cs_parts = cs.process_event(_data_parts) rho = cs.bge_rho.rho() ja.analyze_event(cs_parts) else: ja.analyze_event(_data_parts) rho = ja.rho tmp = [ fill_tree_data(j, tw, sd, rho, iev, 1.) for j in ja.jets if j.pt() > args.jetptcut ] if iev % 1000 == 0: delta_t = time.time() - start_t print('[i] processing event', iev, ' - ev/sec = ', iev / delta_t, 'elapsed = ', delta_t) if args.nev > 0: if iev > args.nev: break print('[i] processed events', iev, ' - ev/sec = ', iev / delta_t, 'elapsed = ', delta_t) outf.Write() outf.Close() print('[i] written', outf.GetName())
class HJetTree(MPBase): def __init__(self, **kwargs): self.configure_from_args( fout=None, output='hjet_tree.root', trigger_range = [6, 7], jet_ana = None) super(HJetTree, self).__init__(**kwargs) if self.fout is None: self.fout = r.TFile(self.output, 'RECREATE') self.fout.cd() self.tree_name='tjet_0{}_{}_{}'.format(int(10*self.jet_ana.jet_R), int(self.trigger_range[0]), int(self.trigger_range[1])) self.tree_writer = RTreeWriter(tree_name=self.tree_name, fout=self.fout) def analyze_event(self, jet_parts): self.trigger_particle = None t_selector = fj.SelectorPtRange(self.trigger_range[0], self.trigger_range[1]) t_candidates = t_selector(jet_parts) if len(t_candidates) < 1: return False self.trigger_particle = random.choice(t_candidates) return True def fill_branches(self, ev_id, weight): if self.jet_ana is None: print('[e] no jet ana in hjettree...') return False if self.trigger_particle: if self.trigger_particle: self.dphis = [self.trigger_particle.delta_phi_to(j) for j in self.jet_ana.jets] self.detas = [self.trigger_particle.eta() - j.eta() for j in self.jet_ana.jets] else: self.dphis = [] self.detas = [] self.tree_writer.fill_branch('ev_id', ev_id) self.tree_writer.fill_branch('t_w', weight) self.tree_writer.fill_branch('t_pt', self.trigger_particle.pt()) self.tree_writer.fill_branch('t_phi', self.trigger_particle.phi()) self.tree_writer.fill_branch('t_eta', self.trigger_particle.eta()) self.tree_writer.fill_branch('dphi', self.dphis) self.tree_writer.fill_branch('deta', self.detas) self.jet_ana.fill_branches(self.tree_writer, ev_id = None, weight = None) def fill_tree(self): self.tree_writer.fill_tree() def write_and_close_file(self): self.fout.Write() self.fout.Close()
class HJetAnalysis(MPBase): def __init__(self, **kwargs): self.configure_from_args(jet_particle_eta_max = 0.9, output='hjet.root', trigger_ranges = [ [0, 1e3], [6, 7] , [12, 22], [20, 30] ], jet_Rs = [0.2, 0.4, 0.6]) super(HJetAnalysis, self).__init__(**kwargs) self.fout = r.TFile(self.output, 'RECREATE') self.fout.cd() self.event_output = RTreeWriter(tree_name='tev', fout=self.fout) self.hmV0M = r.TH1F('hmV0M', 'hmV0M', 1000, 0, 1000) self.hmV0A = r.TH1F('hmV0A', 'hmV0A', 1000, 0, 1000) self.hmV0C = r.TH1F('hmV0C', 'hmV0C', 1000, 0, 1000) self.jet_particle_selector = fj.SelectorAbsEtaMax(self.jet_particle_eta_max) self.jet_ans = [] self.hjet_ts = [] for jR in self.jet_Rs: j_ana = JetAnalysis(jet_R=jR, particle_eta_max=self.jet_particle_eta_max) self.jet_ans.append(j_ana) for tr in self.trigger_ranges: hjet_tree = HJetTree(fout=self.fout, trigger_range=tr, jet_ana=j_ana) self.hjet_ts.append(hjet_tree) def analyze_event(self, ev_id, parts, pythia=None): # V0 multiplicity and event props v0det = V0Detector(particles=parts) self.hmV0M.Fill(v0det.V0_mult) self.hmV0A.Fill(v0det.V0A_mult) self.hmV0C.Fill(v0det.V0C_mult) self.event_output.fill_branch('ev_id', ev_id) self.event_output.fill_branch('mV0A', v0det.V0A_mult) self.event_output.fill_branch('mV0C', v0det.V0C_mult) self.event_output.fill_branch('mV0', v0det.V0_mult) if pythia: ev_w = pythia.info.sigmaGen() pthard = pythia.info.pTHat() else: ev_w = -1 pthard = -1 self.event_output.fill_branch('weight', ev_w) self.event_output.fill_branch('pthard', pthard) jet_parts = self.jet_particle_selector(parts) mTot = len(parts) mCB = len(jet_parts) self.event_output.fill_branch('mTot', mTot) self.event_output.fill_branch('mCB', mCB) self.event_output.fill_tree() trigger_present = False for hj in self.hjet_ts: if hj.analyze_event(jet_parts): trigger_present = True if trigger_present: for ja in self.jet_ans: ja.analyze_event(jet_parts) for hj in self.hjet_ts: if hj.trigger_particle: hj.fill_branches(ev_id, ev_w) hj.tree_writer.fill_branch('ev_id', ev_id) hj.tree_writer.fill_branch('mV0A', v0det.V0A_mult) hj.tree_writer.fill_branch('mV0C', v0det.V0C_mult) hj.tree_writer.fill_branch('mV0', v0det.V0_mult) hj.tree_writer.fill_branch('mTot', mTot) hj.tree_writer.fill_branch('mCB', mCB) hj.tree_writer.fill_branch('weight', ev_w) hj.tree_writer.fill_branch('pthard', pthard) hj.fill_tree() def finalize(self): self.fout.Write() self.fout.Close() print('[i] written', self.fout.GetName())
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())
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('--ignore-mycfg', help="ignore some settings hardcoded here", default=False, action='store_true') parser.add_argument('--output', default="output.root", type=str) parser.add_argument('--alpha', default=0, type=float) parser.add_argument('--dRmax', default=0.25, type=float) parts.add_argument('--zcut', default=0.1, type=float) parser.add_argument('--overwrite', help="overwrite output", default=False, action='store_true') args = parser.parse_args() if args.output == 'output.root': args.output = 'output_alpha_{}_dRmax_{}_SDzcut_{}.root'.format( args.alpha, args.dRmax, args.zcut) if os.path.isfile(args.output): if not args.overwrite: print('[i] output', args.output, 'exists - use --overwrite to do just that...') return # print the banner first fj.ClusterSequence.print_banner() print() # set up our jet definition and a jet selector jet_R0 = 0.4 jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0) jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax( 100.0) & fj.SelectorAbsEtaMax(1 - 1.05 * jet_R0) jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(1 - 1.05 * jet_R0) # jet_selector = fj.SelectorPtMin(10.0) & fj.SelectorPtMax(20.0) & fj.SelectorAbsEtaMax(1 - 1.05 * jet_R0) # jet_selector_cs = fj.SelectorPtMin(0.0) & fj.SelectorAbsEtaMax(1 - 1.05 * jet_R0) print(jet_def) mycfg = ['PhaseSpace:pThatMin = 80', 'PhaseSpace:pThatMax = -1'] # mycfg = [] 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 = args.zcut sd = fjcontrib.SoftDrop(0, sd_zcut, jet_R0) max_eta = 1 cs = CEventSubtractor(alpha=args.alpha, max_distance=args.dRmax, max_eta=max_eta, bge_rho_grid_size=0.25, max_pt_correct=100) be = BoltzmannEvent(mean_pt=0.6, multiplicity=1000 * max_eta * 2, max_eta=max_eta, max_pt=100) parts_selector = fj.SelectorAbsEtaMax(max_eta) if args.nev < 100: args.nev = 100 outf = ROOT.TFile(args.output, 'recreate') outf.cd() tn = ROOT.TNtuple( 'tn', 'tn', 'n:pt:phi:eta:cspt:csphi:cseta:dR:dpt:rg:csrg:z:csz:dRg:dzg') hpt = ROOT.TH1F('hpt', 'hpt', 40, 00, 160) hptcs = ROOT.TH1F('hptcs', 'hptcs', 40, 00, 160) hdpt = ROOT.TH1F('hdpt', 'hdpt', 40, -50, 50) hrg = ROOT.TH1F('hrg', 'hrg', 40, -1.1, 0.9) hrgcs = ROOT.TH1F('hrgcs', 'hrgcs', 40, -1.1, 0.9) hdrg = ROOT.TH1F('hdrg', 'hdrg', 40, -1.1, 0.9) hdz = ROOT.TH1F('hdz', 'hdz', 40, -1.1, 0.9) hdzz = ROOT.TH2F('hdzz', 'hdzz', 40, -1.1, 0.9, 40, -1.1, 0.9) hdphi = ROOT.TH1F('hdphi', 'hdphi', 90, -ROOT.TMath.Pi(), ROOT.TMath.Pi()) t = ROOT.TTree('t', 't') tw = RTreeWriter(tree=t) for i in tqdm.tqdm(range(args.nev)): if not pythia.next(): continue parts_pythia = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal]) parts = parts_selector(parts_pythia) signal_jets = fj.sorted_by_pt(jet_selector(jet_def(parts))) if len(signal_jets) < 1: continue bg_parts = be.generate(offset=10000) full_event = bg_parts sjet = signal_jets[0] lc = [full_event.push_back(psj) for psj in sjet.constituents()] # idxs = [psj.user_index() for psj in sjet.constituents()] # print('pythia jet:', idxs) cs_parts = cs.process_event(full_event) cs_signal_jets = fj.sorted_by_pt(jet_selector_cs(jet_def(cs_parts))) emb_jets = fj.sorted_by_pt(jet_selector_cs(jet_def(full_event))) # max_eta_part = max([j.eta() for j in full_event]) # print ('number of particles', len(full_event), max_eta_part) # mean_pt = sum([j.pt() for j in bg_parts]) / len(bg_parts) # print ('mean pt in bg', mean_pt) # print ('number of CS particles', len(cs_parts)) sd_signal_jet = sd.result(sjet) sd_info_signal_jet = fjcontrib.get_SD_jet_info(sd_signal_jet) # for j in cs_signal_jets: for j in emb_jets: if matched_pt(sjet, j) <= 0.5: continue sd_j = sd.result(j) sd_info_j = fjcontrib.get_SD_jet_info(sd_j) rho = cs.bge_rho.rho() tn.Fill(i, sjet.pt(), sjet.phi(), sjet.eta(), j.pt(), j.phi(), j.eta(), j.delta_R(sjet), j.pt() - sjet.pt(), sd_info_signal_jet.dR, sd_info_j.dR, sd_info_j.dR - sd_info_signal_jet.dR, sd_info_signal_jet.z, sd_info_j.z, sd_info_j.z - sd_info_signal_jet.z) hpt.Fill(sjet.pt()) hptcs.Fill(j.pt()) hdpt.Fill(j.pt() - sjet.pt()) hrg.Fill(sd_info_signal_jet.dR) hrgcs.Fill(sd_info_j.dR) hdphi.Fill(sjet.delta_phi_to(j)) if sd_info_j.dR > 0 and sd_info_signal_jet.dR > 0: hdrg.Fill(sd_info_j.dR - sd_info_signal_jet.dR) if sd_info_j.z > 0 and sd_info_signal_jet.z > 0: hdz.Fill(sd_info_j.z - sd_info_signal_jet.z) hdzz.Fill(sd_info_j.z, sd_info_j.z - sd_info_signal_jet.z) # for i,j in enumerate(signal_jets): # j_sd = sd.result(j) # sd_info = fjcontrib.get_SD_jet_info(j_sd) # # print(" |-> SD jet params z={0:10.3f} dR={1:10.3f} mu={2:10.3f}".format(sd_info.z, sd_info.dR, sd_info.mu)) pythia.stat() outf.Write() outf.Close()