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 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 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()
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_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())