Exemple #1
0
    def __init__(self, **kwargs):
        self.configure_from_args(jet_R=0.4,
                                 jet_algorithm=fj.antikt_algorithm,
                                 particle_eta_max=0.9,
                                 jet_pt_min=0.0,
                                 particles=None,
                                 explicit_ghosts=False)
        super(JetAnalysis, self).__init__(**kwargs)

        self.particle_selector = fj.SelectorAbsEtaMax(self.particle_eta_max)

        self.jet_eta_max = self.particle_eta_max - self.jet_R * 1.05
        # self.jet_eta_max = self.particle_eta_max - 0.05
        self.jet_def = fj.JetDefinition(self.jet_algorithm, self.jet_R)
        self.jet_selector = fj.SelectorPtMin(
            self.jet_pt_min) & fj.SelectorAbsEtaMax(self.jet_eta_max)
        if self.explicit_ghosts:
            self.jet_area_def = fj.AreaDefinition(
                fj.active_area_explicit_ghosts,
                fj.GhostedAreaSpec(self.particle_eta_max))
        else:
            self.jet_area_def = fj.AreaDefinition(
                fj.active_area, fj.GhostedAreaSpec(self.particle_eta_max))

        if self.particles:
            self.analyze_event(self.particles)
Exemple #2
0
    def __init__(self, **kwargs):
        self.configure_from_args(jet_R=0.4,
                                 jet_algorithm=fj.antikt_algorithm,
                                 particle_eta_max=0.9,
                                 particles=None)
        super(JetAnalysis, self).__init__(**kwargs)

        self.bg_rho_range = fj.SelectorAbsEtaMax(self.particle_eta_max)
        self.bg_jet_def = fj.JetDefinition(fj.kt_algorithm, self.jet_R)
        self.bg_area_def = fj.AreaDefinition(
            fj.active_area_explicit_ghosts,
            fj.GhostedAreaSpec(self.particle_eta_max))
        self.bg_estimator = fj.JetMedianBackgroundEstimator(
            self.bg_rho_range, self.bg_jet_def, self.bg_area_def)
        self.rho = 0

        self.jet_eta_max = self.particle_eta_max - self.jet_R * 1.05
        self.jet_def = fj.JetDefinition(self.jet_algorithm, self.jet_R)
        self.jet_selector = fj.SelectorPtMin(0.0) & fj.SelectorAbsEtaMax(
            self.jet_eta_max)
        self.jet_area_def = fj.AreaDefinition(
            fj.active_area_explicit_ghosts,
            fj.GhostedAreaSpec(self.particle_eta_max))

        if self.particles:
            self.analyze_event(self.particles)
Exemple #3
0
    def __init__(self, **kwargs):
        self.configure_from_args(tree_name='tree_Particle',
                                 tree_name_gen='tree_Particle_gen',
                                 args=None)
        super(Embedding, self).__init__(**kwargs)
        self.copy_attributes(self.args)
        self.jet_def = fj.JetDefinition(fj.antikt_algorithm, self.jetR)
        if self.benchmark:
            self.jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax(
                100.0) & fj.SelectorAbsEtaMax(self.max_eta - 1.05 * self.jetR)
            # jet_selector_cs = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * self.jetR)
        else:
            self.jet_selector = fj.SelectorAbsEtaMax(self.max_eta -
                                                     1.05 * self.jetR)
        self.parts_selector = fj.SelectorAbsEtaMax(self.max_eta)

        self.output = EmbeddingOutput(args=self.args)
        # self.output.copy_attributes(self)

        self.sd = fjcontrib.SoftDrop(0, self.sd_zcut, self.jetR)

        self.ja_part = JetAnalysis(jet_R=self.jetR,
                                   jet_algorithm=fj.antikt_algorithm,
                                   jet_pt_min=5.,
                                   particle_eta_max=self.max_eta)
        self.ja_det = JetAnalysis(jet_R=self.jetR,
                                  jet_algorithm=fj.antikt_algorithm,
                                  jet_pt_min=self.jetptcut,
                                  particle_eta_max=self.max_eta)
        self.ja_hybrid = JetAnalysis(jet_R=self.jetR,
                                     jet_algorithm=fj.antikt_algorithm,
                                     jet_pt_min=5.,
                                     particle_eta_max=self.max_eta)

        self.dataPbPb = DataBackgroundIO(name='Data PbPb',
                                         file_list=self.datalistAA)
        self.det_sim = DataIO(name='Sim Pythia Detector level',
                              file_list=self.simulationpp,
                              random_file_order=False)
        self.part_sim = DataIO(name='Sim Pythia Particle level',
                               file_list=self.simulationpp,
                               random_file_order=False,
                               tree_name='tree_Particle_gen')

        self.cs = None
        if self.dRmax > 0:
            self.cs = CEventSubtractor(alpha=self.alpha,
                                       max_distance=self.dRmax,
                                       max_eta=self.max_eta,
                                       bge_rho_grid_size=0.25,
                                       max_pt_correct=100)
Exemple #4
0
def main():
	parser = argparse.ArgumentParser(description='pythia8 fastjet on the fly', prog=os.path.basename(__file__))
	pyconf.add_standard_pythia_args(parser)
	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.4
	jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
	jet_selector = fj.SelectorPtMin(5.0) & fj.SelectorAbsEtaMax(2)
	print(jet_def)

	all_jets = []

	# 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
		parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
		jets = jet_selector(jet_def(parts))
		all_jets.extend(jets)

	pythia.stat()

	jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
	lund_gen = fjcontrib.LundGenerator(jet_def_lund)

	print('making lund diagram for all jets...')
	lunds = [lund_gen.result(j) for j in all_jets]

	print('listing lund plane points... Delta, kt - for {} selected jets'.format(len(all_jets)))
	for l in lunds:
		print ('- jet pT={0:5.2f} eta={1:5.2f}'.format(l[0].pair().perp(), l[0].pair().eta()))
		print ('  Deltas={}'.format([s.Delta() for s in l]))
		print ('  kts={}'.format([s.Delta() for s in l]))
		print ( )

	print('[i] reclustering and using soft drop...')
	jet_def_rc = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
	print('Reclustering:', jet_def_rc)

	rc = fjcontrib.Recluster(jet_def_rc, True)
	sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
	for i,j in enumerate(all_jets):
		j_rc = rc.result(j)
		print()
		print('- [{0:3d}] orig pT={1:10.3f} reclustered pT={2:10.3f}'.format(i, j.perp(), j_rc.perp()))
		j_sd = sd.result(j)
		print('  |-> after soft drop pT={0:10.3f} delta={1:10.3f}'.format(j_sd.perp(), j_sd.perp() - j.perp()))
		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))
Exemple #5
0
	def __init__(self, **kwargs):
		self.fout = None
		super(HFAnalysisInvMass, self).__init__(**kwargs)
		self.fout = ROOT.TFile(self.name+'.root', 'recreate')
		self.fout.cd()
		# self.hinvmass = ROOT.TH1F('hinvmass', 'hinvmass', 400, 1.5, 2.5)
		# self.hinvmass.Sumw2()
		# self.hinvmasspt = ROOT.TH2F('hinvmasspt', 'hinvmasspt', 400, 1.5, 2.5, 50, 2, 12)
		# self.hinvmasspt.Sumw2()
		self.tw = treewriter.RTreeWriter(tree_name='d0', fout=self.fout)
		# jet stuff

		max_eta = 0.9
		self.parts_selector = fj.SelectorPtMin(0.15) & fj.SelectorPtMax(100.0) & fj.SelectorAbsEtaMax(max_eta)

		jet_R0 = 0.4
		self.jet_selector = fj.SelectorPtMin(5.0) & fj.SelectorPtMax(100.0) & fj.SelectorAbsEtaMax(max_eta - 1.05 * jet_R0)
		self.jarho = JetAnalysisWithRho(jet_R=jet_R0, jet_algorithm=fj.antikt_algorithm, particle_eta_max=max_eta)
Exemple #6
0
def main(args):
    nevents = args.nevents
    sconfig_pythia = get_pythia_config(args)
    if args.generate:
        pythia = create_and_init_pythia(sconfig_pythia)
        if not pythia:
            return
        all_jets = []
        for iEvent in tqdm(range(nevents), 'event'):
            if not pythia.next(): continue
        print("[i] done generating")

    if args.write:
        pythia = create_and_init_pythia(sconfig_pythia)
        if not pythia:
            return
        pyhepmcwriter = mp.Pythia8HepMCWrapper(args.write)
        all_jets = []
        for iEvent in tqdm(range(nevents), 'event'):
            if not pythia.next(): continue
            pyhepmcwriter.fillEvent(pythia)
        print("[i] done writing to {}".format(args.write))

    if args.read:
        import pyhepmc_ng
        input = pyhepmc_ng.ReaderAsciiHepMC2(args.read)
        if input.failed():
            print("[error] unable to read from {}".format(args.read))
            return

        # print the banner first
        fj.ClusterSequence.print_banner()
        print()
        jet_R0 = 0.4
        jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
        jet_selector = fj.SelectorPtMin(100.0) & fj.SelectorPtMax(
            200.0) & fj.SelectorAbsEtaMax(1)

        event = pyhepmc_ng.GenEvent()
        pbar = tqdm(range(nevents))
        while not input.failed():
            e = input.read_event(event)
            if input.failed():
                break
            fjparts = []
            for i, p in enumerate(event.particles):
                if p.status == 1:
                    psj = fj.PseudoJet(p.momentum.px, p.momentum.py,
                                       p.momentum.pz, p.momentum.e)
                    psj.set_user_index(i)
                    fjparts.append(psj)
            jets = jet_selector(jet_def(fjparts))
            pbar.update()
Exemple #7
0
 def get_event(self):
     if not self.pythia:
         return None
     parts_selector = fj.SelectorAbsEtaMax(self.eta_max)
     jet_selector = fj.SelectorPtMin(
         self.jet_pt_min) & fj.SelectorAbsEtaMax(self.eta_max -
                                                 1.05 * self.jet_R0)
     jet_def = fj.JetDefinition(fj.antikt_algorithm, self.jet_R0)
     while True:
         if not self.pythia.next():
             continue
         self.parts_pythia = pythiafjext.vectorize_select(
             self.pythia, [pythiafjext.kFinal])
         parts_gen = parts_selector(self.parts_pythia)
         signal_jets = fj.sorted_by_pt(
             jet_selector(jet_def(self.parts_pythia)))
         if len(signal_jets) < 1:
             continue
         else:
             break
     return self.parts_pythia
Exemple #8
0
def main():
    parser = argparse.ArgumentParser(description='pythia8 in python',
                                     prog=os.path.basename(__file__))
    pyconf.add_standard_pythia_args(parser)
    args = parser.parse_args()

    mycfg = []
    pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)

    max_eta_hadron = 3
    jet_R0 = 0.4
    jet_selector = fj.SelectorPtMin(100.0) & fj.SelectorPtMax(
        125.0) & fj.SelectorAbsEtaMax(max_eta_hadron - 1.05 * jet_R0)
    parts_selector_h = fj.SelectorAbsEtaMax(max_eta_hadron)

    fj.ClusterSequence.print_banner()
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)

    pbar = tqdm.tqdm(range(args.nev))
    for i in pbar:
        if not pythia.next():
            pbar.update(-1)
            continue

        parts_pythia_h = pythiafjext.vectorize_select(pythia,
                                                      [pythiafjext.kFinal], 0,
                                                      False)
        parts_pythia_h_selected = parts_selector_h(parts_pythia_h)

        jets_h = fj.sorted_by_pt(jet_selector(
            jet_def(parts_pythia_h_selected)))

        if len(jets_h) < 1:
            continue

        # do your things with jets here...

    pythia.stat()
    pythia.settings.writeFile(args.py_cmnd_out)
    def init_jet_tools(self):

        for jetR in self.jetR_list:
            jetR_str = str(jetR).replace('.', '')

            if not self.no_tree:
                # Initialize tree writer
                name = 'particle_unscaled_R%s' % jetR_str
                t = ROOT.TTree(name, name)
                setattr(self, "t_R%s" % jetR_str, t)
                tw = RTreeWriter(tree=t)
                setattr(self, "tw_R%s" % jetR_str, tw)

            # set up our jet definition and a jet selector
            jet_def = fj.JetDefinition(fj.antikt_algorithm, jetR)
            setattr(self, "jet_def_R%s" % jetR_str, jet_def)
            print(jet_def)

        pwarning('max eta for particles after hadronization set to',
                 self.max_eta_hadron)
        parts_selector_h = fj.SelectorAbsEtaMax(self.max_eta_hadron)

        for jetR in self.jetR_list:
            jetR_str = str(jetR).replace('.', '')

            jet_selector = fj.SelectorPtMin(5.0) & \
                           fj.SelectorAbsEtaMax(self.max_eta_hadron - jetR)
            setattr(self, "jet_selector_R%s" % jetR_str, jet_selector)

            #max_eta_parton = self.max_eta_hadron + 2. * jetR
            #setattr(self, "max_eta_parton_R%s" % jetR_str, max_eta_parton)
            #pwarning("Max eta for partons with jet R =", jetR, "set to", max_eta_parton)
            #parts_selector_p = fj.SelectorAbsEtaMax(max_eta_parton)
            #setattr(self, "parts_selector_p_R%s" % jetR_str, parts_selector_p)

            count1 = 0  # Number of jets rejected from ch-h matching
            setattr(self, "count1_R%s" % jetR_str, count1)
            count2 = 0  # Number of jets rejected from h-p matching
            setattr(self, "count2_R%s" % jetR_str, count2)
Exemple #10
0
    def init_jet_tools(self):
        
        for jetR in self.jetR_list:
            jetR_str = str(jetR).replace('.', '')      
            
            # set up our jet definition and a jet selector
            jet_def = fj.JetDefinition(fj.antikt_algorithm, jetR)
            setattr(self, "jet_def_R%s" % jetR_str, jet_def)
            print(jet_def)

        pwarning('max eta for particles after hadronization set to', self.max_eta_hadron)
        parts_selector_h = fj.SelectorAbsEtaMax(self.max_eta_hadron)

        for jetR in self.jetR_list:
            jetR_str = str(jetR).replace('.', '')
            
            jet_selector = fj.SelectorPtMin(5.0) & \
                           fj.SelectorAbsEtaMax(self.max_eta_hadron - jetR)
            setattr(self, "jet_selector_R%s" % jetR_str, jet_selector)

            count1 = 0  # Number of jets rejected from ch-h matching
            setattr(self, "count1_R%s" % jetR_str, count1)
            count2 = 0  # Number of jets rejected from h-p matching
            setattr(self, "count2_R%s" % jetR_str, count2)
Exemple #11
0
    def init(self):
        """
        must be called before next(). Initializes the generator, and creates
        all reused fastjet objects (selectors, jet definition). Returns the 
        status of the pythia generator
        """
        # first initialize Pythia
        self.gen.readString("Beams:eCM = 200.0")
        self.gen.readString("HardQCD:all = on")
        self.gen.readString("Random:setSeed = on")
        self.gen.readString('Random:seed = ' + str(self.seed))
        self.gen.readString('PhaseSpace:pTHatMin = ' + str(self.pt_hat_min))
        self.gen.readString('PhaseSpace:pTHatMax = ' + str(self.pt_hat_max))
        for setting in self.gen_args:
            try:
                self.gen.readString(setting)
            except:
                print('pythia read string failed')

        # create jet definition
        self.jet_def = fastjet.JetDefinition(self.jet_alg, self.jet_radius)

        # create fastjet selectors for tracks and jets
        self.track_selector = (
            fastjet.SelectorAbsEtaMax(self.track_eta_max) *
            fastjet.SelectorPtRange(self.track_pt_min, self.track_pt_max))
        self.jet_selector = (fastjet.SelectorAbsEtaMax(self.jet_eta_max) *
                             fastjet.SelectorPtMin(self.jet_pt_min))

        self.initialized = self.gen.init()

        if not self.star_sim and not self.gauss_smear and not self.tpc_sim:
            print('Warning: no detector simulation method has been chosen',
                  'disabling detector simulation')

        return self.initialized
Exemple #12
0
def main_make(args):
	if os.path.isfile(args.output):
		if not args.overwrite:
			print('[i] output', args.output, 'exists - use --overwrite to do just that...')
			return

	print(args)

	# alice specific
	max_eta = 0.9

	outf = ROOT.TFile(args.output, 'recreate')
	outf.cd()

	bgestim = BackgroundEstimator(name='bg_estim', grid_size=0.1, particle_eta_max=max_eta)
	print(bgestim)

	be = None
	data = None
	if len(args.data) > 0:
		data = DataIO(file_list=args.data)
		print(data)
	else:
		be = BoltzmannEvent(mean_pt=0.6, multiplicity=2000 * max_eta * 2, max_eta=max_eta, max_pt=100)
		print(be)

	parts_selector = fj.SelectorAbsEtaMax(max_eta)

	if args.nev < 1:
		args.nev = 1

	### EVENT LOOP STARTS HERE
	for iev in tqdm.tqdm(range(args.nev)):
		if data:
			bg_parts = data.load_event(offset=10000)
		else:
			if be:
				bg_parts = be.generate(offset=10000)
		bgestim.fill_grid(bg_parts)

	outf.cd()
	if be:
		be.write()
	bgestim.write()
	outf.Write()
	outf.Close()
	print('[i] written', outf.GetName())
Exemple #13
0
def main():
    input_file = "$HOME/data/jetscape/test_out.hepmc"
    if len(sys.argv) > 1:
        input_file = sys.argv[1]

    input_file = os.path.expandvars(input_file)

    print('[i] reading from:', input_file)
    # input = pyhepmc_ng.ReaderAsciiHepMC2(input_file)
    input = pyhepmc_ng.ReaderAscii(input_file)
    if input.failed():
        print("[error] unable to read from {}".format(input_file))
        return
    nevents = 1000

    # print the banner first
    fj.ClusterSequence.print_banner()
    print()
    jet_R0 = 0.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    jet_selector = fj.SelectorPtMin(0.0) & fj.SelectorPtMax(
        200.0) & fj.SelectorAbsEtaMax(3)

    all_jets = []
    event = pyhepmc_ng.GenEvent()
    pbar = tqdm.tqdm(range(nevents))
    while not input.failed():
        e = input.read_event(event)
        if input.failed():
            break
        fjparts = []
        for i, p in enumerate(event.particles):
            if p.status == 1 and not p.end_vertex:
                psj = fj.PseudoJet(p.momentum.px, p.momentum.py, p.momentum.pz,
                                   p.momentum.e)
                psj.set_user_index(i)
                fjparts.append(psj)
        jets = jet_selector(jet_def(fjparts))
        all_jets.append([[j.pt(), j.eta()] for j in jets])
        pbar.update()
        for j in jets:
            hjetpt.Fill(j.perp())
        if pbar.n >= nevents:
            break
    pbar.close()
Exemple #14
0
def main_jets(args):
	outfname = '{}_output.root'.format(args.read).replace('.dat', '')
	print("output file: {}".format(outfname))
	foutput = r.TFile(outfname, "recreate")
	foutput.cd()
	lbins = logbins(1., 100, 10)
	hjetpt = r.TH1F('hjetpt', 'hjetpt', 10, lbins)

	input = pyhepmc_ng.ReaderAsciiHepMC2(args.read)
	if input.failed():
		print ("[error] unable to read from {}".format(args.read))
		return

	# print the banner first
	fj.ClusterSequence.print_banner()
	print()
	jet_R0 = 0.4
	jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
	jet_selector = fj.SelectorPtMin(0.0) & fj.SelectorPtMax(200.0) & fj.SelectorAbsEtaMax(3)

	all_jets = []
	event = pyhepmc_ng.GenEvent()
	pbar = tqdm(range(args.nevents))
	while not input.failed():
		e = input.read_event(event)
		if input.failed():
			break
		fjparts = []
		for i,p in enumerate(event.particles):
			if p.status == 1 and not p.end_vertex:
				psj = fj.PseudoJet(p.momentum.px, p.momentum.py, p.momentum.pz, p.momentum.e)
				psj.set_user_index(i)
				fjparts.append(psj)
		jets = jet_selector(jet_def(fjparts))
		all_jets.append([ [j.pt(), j.eta()] for j in jets])
		pbar.update()
		for j in jets:
			hjetpt.Fill(j.perp())
		if pbar.n >= args.nevents:
			break
	foutput.Write()
	foutput.Close()
	joblib.dump(all_jets, outfname.replace(".root", ".joblib"))
Exemple #15
0
def main():
    fj.ClusterSequence.print_banner()
    print()

    sconfig_pythia = [
        "Beams:eCM = 8000.", "HardQCD:all = on", "PhaseSpace:pTHatMin = 20."
    ]
    pythia = create_and_init_pythia(sconfig_pythia)
    jet_R0 = 0.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    jet_selector = fj.SelectorPtMin(20.0) & fj.SelectorPtMax(
        40.0) & fj.SelectorAbsEtaMax(1)
    # sd = rt.SoftDrop(0, 0.1, 1.0)

    all_jets = []
    for iEvent in tqdm(range(1000), 'event'):
        if not pythia.next():
            continue
        parts = pyfj.vectorize(pythia, True, -1, 1, True)
        jets = jet_selector(jet_def(parts))
        all_jets.extend(jets)

    print_jets(all_jets)
Exemple #16
0
	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)
Exemple #17
0
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 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('--user-seed', help='PYTHIA starting seed', default=1111, type=int)
    parser.add_argument('--output', default="output.root", type=str)
    parser.add_argument('--sd_beta', help='SoftDrop beta', default=None, type=float)
    parser.add_argument('--jetR', help='Jet radius/resolution parameter', default=0.4, type=float)
    parser.add_argument('--no-match-level', help="Save simulation for only one level with " + \
                        "no matching. Options: 'p', 'h', 'ch'", default=None, type=str)
    parser.add_argument('--p-ch-MPI', help="Match between parton level (no MPI) and charged " + \
                        "hadron level (with MPI) for ALICE response characterization",
                        action='store_true', default=False)
    args = parser.parse_args()

    level = args.no_match_level
    if args.p_ch_MPI:
        if level:
            print("ERROR: --no-match-level and --p-ch-MPI cannot be set simultaneously.")
            exit(1)
    if level not in [None, 'p', 'h', 'ch']:
        print("ERROR: Unrecognized type %s. Please use 'p', 'h', or 'ch'" % args.type_only)
        exit(1)

    # Angularity beta values
    betas = [1, 1.5, 2, 3]

    if args.user_seed < 0:
        args.user_seed = 1111
    pinfo('user seed for pythia', args.user_seed)
    # mycfg = ['PhaseSpace:pThatMin = 100']
    mycfg = ['Random:setSeed=on', 'Random:seed={}'.format(args.user_seed)]
    mycfg.append('HadronLevel:all=off')

    # Have at least 1 event
    if args.nev < 1:
        args.nev = 1

    if args.p_ch_MPI:
        d = vars(args)
        d['py_noMPI'] = True
        pythia_noMPI = pyconf.create_and_init_pythia_from_args(args, mycfg)
        args_MPI = copy.deepcopy(args)
        d = vars(args_MPI)
        d['py_noMPI'] = False
        pythia = pyconf.create_and_init_pythia_from_args(args_MPI, mycfg)
    else:
        pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)

    # print the banner first
    fj.ClusterSequence.print_banner()
    print()
    # set up our jet definition and a jet selector
    jet_R0 = args.jetR
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    print(jet_def)

    # hadron level - ALICE
    max_eta_hadron = 0.9
    pwarning('max eta for particles after hadronization set to', max_eta_hadron)
    parts_selector_h = fj.SelectorAbsEtaMax(max_eta_hadron)
    jet_selector = fj.SelectorPtMin(5.0) #& fj.SelectorAbsEtaMax(max_eta_hadron - jet_R0)

    max_eta_parton = max_eta_hadron + 2. * jet_R0
    pwarning('max eta for partons set to', max_eta_parton)
    parts_selector_p = fj.SelectorAbsEtaMax(max_eta_parton)

    if not args.sd_beta is None:
        args.output = args.output.replace('.root', '_sdbeta{}.root'.format(args.sd_beta))
    outf = ROOT.TFile(args.output, 'recreate')
    outf.cd()
    t = ROOT.TTree('t', 't')
    tw = RTreeWriter(tree=t)

    count1 = 0
    count2 = 0

    # event loop
    for iev in range(args.nev):  #tqdm.tqdm(range(args.nev)):
        if not pythia.next():
            if args.p_ch_MPI:
                pythia_noMPI.next()
            continue
        if args.p_ch_MPI and not pythia_noMPI.next():
            continue

        parts_pythia_p = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, True)
        parts_pythia_p_selected = parts_selector_p(parts_pythia_p)

        if args.p_ch_MPI:
            parts_pythia_p = pythiafjext.vectorize_select(pythia_noMPI, [pythiafjext.kFinal], 0, True)
            parts_pythia_p_selected = parts_selector_p(parts_pythia_p)

        hstatus = pythia.forceHadronLevel()
        if not hstatus:
            #pwarning('forceHadronLevel false event', iev)
            continue
        # parts_pythia_h = pythiafjext.vectorize_select(
        #     pythia, [pythiafjext.kHadron, pythiafjext.kCharged])
        parts_pythia_h = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, True)
        parts_pythia_h_selected = parts_selector_h(parts_pythia_h)

        parts_pythia_hch = pythiafjext.vectorize_select(
            pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, True)
        parts_pythia_hch_selected = parts_selector_h(parts_pythia_hch)

        # pinfo('debug partons...')
        # for p in parts_pythia_p_selected:
        #   pyp = pythiafjext.getPythia8Particle(p)
        #   print(pyp.name())
        # pinfo('debug hadrons...')
        # for p in parts_pythia_h_selected:
        #   pyp = pythiafjext.getPythia8Particle(p)
        #   print(pyp.name())
        # pinfo('debug ch. hadrons...')
        # for p in parts_pythia_hch_selected:
        #   pyp = pythiafjext.getPythia8Particle(p)
        #   print(pyp.name())

        # parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
        jets_p = fj.sorted_by_pt(jet_selector(jet_def(parts_pythia_p)))
        jets_h = fj.sorted_by_pt(jet_selector(jet_def(parts_pythia_h)))
        jets_ch = fj.sorted_by_pt(jet_selector(jet_def(parts_pythia_hch)))

        if level:  # Only save info at one level w/o matching
            jets = locals()["jets_%s" % level]
            for jet in jets:
                tw.fill_branch('iev', iev)
                tw.fill_branch(level, jet)
                kappa = 1
                for beta in betas:
                    label = str(beta).replace('.', '')
                    tw.fill_branch('l_%s_%s' % (level, label),
                                   lambda_beta_kappa(jet, jet_R0, beta, kappa))
            continue

        if not args.sd_beta is None:
            sd = fjcontrib.SoftDrop(args.sd_beta, 0.1, jet_R0)

        for i,jchh in enumerate(jets_ch):

            # match hadron (full) jet
            drhh_list = []
            for j, jh in enumerate(jets_h):
                drhh = jchh.delta_R(jh)
                if drhh < jet_R0 / 2.:
                    drhh_list.append((j,jh))
            if len(drhh_list) != 1:
                count1 += 1
            else:  # Require unique match
                j, jh = drhh_list[0]

                # match parton level jet
                dr_list = []
                for k, jp in enumerate(jets_p):
                    dr = jh.delta_R(jp)
                    if dr < jet_R0 / 2.:
                        dr_list.append((k, jp))
                if len(dr_list) != 1:
                    count2 += 1
                else:
                    k, jp = dr_list[0]

                    # pwarning('event', iev)
                    # pinfo('matched jets: ch.h:', jchh.pt(), 'h:', jh.pt(),
                    #       'p:', jp.pt(), 'dr:', dr)

                    tw.fill_branch('iev', iev)
                    tw.fill_branch('ch', jchh)
                    tw.fill_branch('h', jh)
                    tw.fill_branch('p', jp)

                    kappa = 1
                    for beta in betas:
                        label = str(beta).replace('.', '')
                        tw.fill_branch("l_ch_%s" % label,
                                       lambda_beta_kappa(jchh, jet_R0, beta, kappa))
                        tw.fill_branch("l_h_%s" % label,
                                       lambda_beta_kappa(jh, jet_R0, beta, kappa))
                        tw.fill_branch("l_p_%s" % label,
                                       lambda_beta_kappa(jp, jet_R0, beta, kappa))

                    if not args.sd_beta is None:
                        jchh_sd = sd.result(jchh)
                        jchh_sd_info = fjcontrib.get_SD_jet_info(jchh_sd)
                        jh_sd = sd.result(jh)
                        jh_sd_info = fjcontrib.get_SD_jet_info(jh_sd)
                        jp_sd = sd.result(jp)
                        jp_sd_info = fjcontrib.get_SD_jet_info(jp_sd)

                        tw.fill_branch('p_zg', jp_sd_info.z)
                        tw.fill_branch('p_Rg', jp_sd_info.dR)
                        tw.fill_branch('p_thg', jp_sd_info.dR/jet_R0)
                        tw.fill_branch('p_mug', jp_sd_info.mu)

                        tw.fill_branch('h_zg', jh_sd_info.z)
                        tw.fill_branch('h_Rg', jh_sd_info.dR)
                        tw.fill_branch('h_thg', jh_sd_info.dR/jet_R0)
                        tw.fill_branch('h_mug', jh_sd_info.mu)

                        tw.fill_branch('ch_zg', jchh_sd_info.z)
                        tw.fill_branch('ch_Rg', jchh_sd_info.dR)
                        tw.fill_branch('ch_thg', jchh_sd_info.dR/jet_R0)
                        tw.fill_branch('ch_mug', jchh_sd_info.mu)

            #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))

    tw.fill_tree()
    if args.p_ch_MPI:
        pythia_noMPI.stat()
    pythia.stat()

    print("%i jets cut at first match criteria; %i jets cut at second match criteria." % 
          (count1, count2))

    outf.Write()
    outf.Close()
Exemple #19
0
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('--nw',
                        help="no warn",
                        default=True,
                        action='store_true')
    parser.add_argument('--ignore-mycfg',
                        help="ignore some settings hardcoded here",
                        default=False,
                        action='store_true')
    parser.add_argument('--enable-background',
                        help="enable background calc",
                        default=False,
                        action='store_true')
    parser.add_argument('--output',
                        help="output file name",
                        default='leadsj_vs_zloss.root',
                        type=str)
    parser.add_argument('--jetptmin',
                        help="jet pt minimum",
                        default=-1,
                        type=float)
    parser.add_argument('--jetptmax',
                        help="jet pt maximum",
                        default=1e6,
                        type=float)
    parser.add_argument('--eta', help="jet eta max", default=2.4, type=float)
    parser.add_argument(
        '--kt',
        help="use kT algorithm instead of anti-kT for the subjets",
        default=False,
        action='store_true')

    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.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    jet_selector = fj.SelectorPtMin(args.py_pthatmin) & fj.SelectorPtMax(
        1000.0) & fj.SelectorAbsEtaMax(args.eta - jet_R0)
    mycfg = []
    if args.jetptmin > 0:
        mycfg = ['PhaseSpace:pThatMin = {}'.format(args.jetptmin)]
        jet_selector = fj.SelectorPtMin(args.jetptmin) & fj.SelectorPtMax(
            args.jetptmax) & fj.SelectorAbsEtaMax(args.eta - jet_R0)
    print(jet_def)

    if args.ignore_mycfg:
        mycfg = []
    pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
    if not pythia:
        perror("pythia initialization failed.")
        return

    nbins = 20
    # sjrs = [0.001 + x * 0.04 for x in range(0, nbins)]
    sjrs = logbins(0.001, jet_R0, nbins)
    print(sjrs)
    print('log(1/r) :', [ROOT.TMath.Log(1 / r) for r in sjrs])
    sjdefs = dict()
    for sjr in sjrs:
        if args.kt:
            _jet_def = fj.JetDefinition(fj.kt_algorithm, sjr)
        else:
            _jet_def = fj.JetDefinition(fj.antikt_algorithm, sjr)
        sjdefs[sjr] = _jet_def

    # tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = 'leadsj_vs_x.root')
    tw = treewriter.RTreeWriter(name='lsjvszloss', file_name=args.output)
    tw.fout.cd()
    h_zloss_r_q = dict()
    h_zloss_r_g = dict()
    for sjr in sjrs:
        sname = 'h_zloss_glue_{}'.format(sjr)
        _h_zloss_r_g = ROOT.TH1F(sname, sname, len(sjrs), 0., 1.)
        h_zloss_r_g[sjr] = _h_zloss_r_g
        sname = 'h_zloss_quark_{}'.format(sjr)
        _h_zloss_r_q = ROOT.TH1F(sname, sname, len(sjrs), 0., 1.)
        h_zloss_r_q[sjr] = _h_zloss_r_q

    lbins = logbins(ROOT.TMath.Log(1. / jet_R0), ROOT.TMath.Log(1. / sjrs[0]),
                    nbins)
    print('lbins:', lbins)

    sname = 'prof_zloss_vs_r_any'
    prof_a = ROOT.TProfile(sname, sname, nbins, 0, jet_R0)
    prof_a_log = ROOT.TProfile(sname + '_log', sname + '_log', nbins, lbins)

    sname = 'prof_zloss_vs_r_glue'
    prof_g = ROOT.TProfile(sname, sname, nbins, 0, jet_R0)
    prof_g_log = ROOT.TProfile(sname + '_log', sname + '_log', nbins, lbins)

    sname = 'prof_zloss_vs_r_quark'
    prof_q = ROOT.TProfile(sname, sname, nbins, 0, jet_R0)
    prof_q_log = ROOT.TProfile(sname + '_log', sname + '_log', nbins, lbins)
    # prof_q_log = ROOT.TProfile(sname+'_log', sname+'_log', nbins, ROOT.TMath.Log(1./jet_R0), ROOT.TMath.Log(1./sjrs[0]))

    sname = 'h2_zloss_vs_r_glue'
    h2_zloss_r_g = ROOT.TH2F(sname, sname, nbins, 0., jet_R0, len(sjrs), 0.,
                             1.)
    sname = 'h2_zloss_vs_r_quark'
    h2_zloss_r_q = ROOT.TH2F(sname, sname, nbins, 0., jet_R0, len(sjrs), 0.,
                             1.)

    # loop

    if args.nev < 100:
        args.nev = 100
    for i in tqdm.tqdm(range(args.nev)):
        if not pythia.next():
            continue
        # parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
        partons = pythiafjext.vectorize_select(pythia, [pythiafjext.kParton],
                                               0, True)
        parts = pythiafjext.vectorize_select(
            pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
        # parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)
        jets = jet_selector(jet_def(parts))

        # for j in tqdm.tqdm(jets):
        for j in jets:
            j_type = match_dR(j, partons, jet_R0 / 2.)
            if j_type[0] is None:
                if args.nw:
                    continue
                pwarning('Jet with no parton label')
                continue

            tw.fill_branch("j", j)
            for sjr in sjrs:
                rc_jets = fj.sorted_by_pt(sjdefs[sjr](j.constituents()))
                tw.fill_branch("sjr{}".format(sjr), rc_jets[0])
                zloss = 1. - rc_jets[0].perp() / j.perp()
                tw.fill_branch("sjr{}_zloss".format(sjr), zloss)
                tw.fill_branch("ppid", j_type[0])
                tw.fill_branch("pquark", j_type[1])
                tw.fill_branch("pglue", j_type[2])

                prof_a.Fill(sjr, zloss)
                prof_a_log.Fill(ROOT.TMath.Log(1. / sjr), zloss)

                if j_type[1]:
                    h_zloss_r_q[sjr].Fill(zloss)
                    h2_zloss_r_q.Fill(sjr, zloss)
                    prof_q.Fill(sjr, zloss)
                    prof_q_log.Fill(ROOT.TMath.Log(1. / sjr), zloss)
                if j_type[2]:
                    h_zloss_r_g[sjr].Fill(zloss)
                    h2_zloss_r_g.Fill(sjr, zloss)
                    prof_g.Fill(sjr, zloss)
                    prof_g_log.Fill(ROOT.TMath.Log(1. / sjr), zloss)

            tw.fill_tree()

    pythia.stat()
    tw.write_and_close()
Exemple #20
0
    return None


sconfig_pythia = [
    "Beams:eCM = 8000.", "HardQCD:all = on", "PhaseSpace:pTHatMin = 100."
]
pythia = create_and_init_pythia(sconfig_pythia)

# 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(100.0) & fj.SelectorPtMax(
    200.0) & fj.SelectorAbsEtaMax(1)
sd = fjcontrib.SoftDrop(0, 0.1, 1.0)

all_jets = []
for iEvent in tqdm(range(10), 'event'):
    if not pythia.next(): continue
    parts = pyfj.vectorize(pythia, True, -1, 1, False)
    jets = jet_selector(jet_def(parts))
    all_jets.extend(jets)

all_sd_jets = [sd.result(j) for j in all_jets]
pts = [j.pt() for j in all_jets]
sd_pts = [j.pt() for j in all_sd_jets]
sd_delta_pt = [delta for delta in deltas(all_jets, all_sd_jets)]
nangs0 = [pyfj.angularity(j, 0.) for j in all_jets]
Exemple #21
0
def main():
    mycfg = []
    ssettings = "--py-ecm 5000 --py-minbias --user-seed=100000"
    args = get_args_from_settings(ssettings)
    pythia_mb = pyconf.create_and_init_pythia_from_args(args, mycfg)

    mycfg = []
    ssettings = "--py-ecm 5000 --user-seed=100000 --nev 100000"
    args = get_args_from_settings(ssettings)
    pythia_hard = pyconf.create_and_init_pythia_from_args(args, mycfg)

    max_eta_hadron = 1
    parts_selector_h = fj.SelectorAbsEtaMax(max_eta_hadron)
    jet_R0 = 0.4
    jet_selector = fj.SelectorPtMin(20.0) & fj.SelectorPtMax(
        200.0) & fj.SelectorAbsEtaMax(max_eta_hadron - 1.05 * jet_R0)
    n_pileup = 1  #5

    # 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)

    fout = ROOT.TFile(args.output, 'recreate')
    fout.cd()
    tj_delta = ROOT.TNtuple(
        "tj_delta", "tj_delta",
        "pt:eta:phi:L11:L21:L31:ptm:etam:phim:L11m:L21m:L31m:mpt")
    tj_no_pup = ROOT.TNtuple("tj_no_pup", "tj_no_pup",
                             "pt:eta:phi:L11:L21:L31")
    tj_pup = ROOT.TNtuple("tj_pup", "tj_pup", "pt:eta:phi:L11:L21:L31")
    hmult_hard = ROOT.TH1F("hmult_hard", "hmult_hard", 300, 0, 300)
    hmult_pup = ROOT.TH1F("hmult_pup", "hmult_pup", 300, 0, 300)
    hpt_acc_hard = ROOT.TProfile2D("hpt_acc_hard", "hpt_acc_hard;#eta;#varphi",
                                   50, -1, 1, 50, 0,
                                   ROOT.TMath.Pi() * 2.)
    hpt_acc_pup = ROOT.TProfile2D("hpt_acc_pup", "hpt_acc_pup;#eta;#varphi",
                                  50, -1, 1, 50, 0,
                                  ROOT.TMath.Pi() * 2.)

    for n in tqdm(range(args.nev)):
        if not pythia_hard.next():
            continue
        parts_pythia_h = pythiafjext.vectorize_select(
            pythia_hard, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
        parts_pythia_h_selected = parts_selector_h(parts_pythia_h)

        parts_pileup = None
        for ipile in range(n_pileup):
            while not pythia_mb.next():
                continue
            parts_pythia_h_ue = pythiafjext.vectorize_select(
                pythia_mb, [pythiafjext.kFinal, pythiafjext.kCharged], 10000,
                False)
            parts_pythia_h_selected_ue = parts_selector_h(parts_pythia_h_ue)
            if parts_pileup is None:
                parts_pileup = parts_pythia_h_selected_ue
            else:
                parts_pileup += parts_pythia_h_selected_ue

        mult_hard = len(parts_pythia_h_selected)
        mult_ue = len(parts_pileup)

        jets_h = fj.sorted_by_pt(jet_selector(
            jet_def(parts_pythia_h_selected)))
        jets_h_w_ue = fj.sorted_by_pt(
            jet_selector(jet_def(parts_pileup + parts_pythia_h_selected)))

        if len(jets_h) < 1:
            continue

        fill_matched(jets_h, jets_h_w_ue, tj_no_pup, tj_pup, tj_delta, jet_R0)

        hmult_hard.Fill(mult_hard)
        hmult_pup.Fill(mult_ue)

        _tmp = [
            hpt_acc_hard.Fill(p.eta(), p.phi(), p.perp())
            for p in parts_pythia_h_selected
        ]
        _tmp = [
            hpt_acc_pup.Fill(p.eta(), p.phi(), p.perp()) for p in parts_pileup
        ]

    pythia_hard.stat()
    pythia_mb.stat()

    fout.Write()
    fout.Close()
    print('[i] written ', fout.GetName())
Exemple #22
0
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('--nw', help="no warn", default=False, action='store_true')
	parser.add_argument('--ignore-mycfg', help="ignore some settings hardcoded here", default=False, action='store_true')
	parser.add_argument('--enable-background', help="enable background calc", default=False, action='store_true')
	parser.add_argument('--output', help="output file name", default='leadsj_vs_x_output.root', type=str)

	# for background
	parser.add_argument('--cent-bin', help="centraility bin 0 is the  0-5 percent most central bin", type=int, default=0)
	parser.add_argument('--seed', help="pr gen seed", type=int, default=1111)
	parser.add_argument('--harmonics', help="set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only)", 
						type=int, default=5)
	parser.add_argument('--eta', help="set eta range must be uniform (e.g. abs(eta) < 0.9, which is ALICE TPC fiducial acceptance)",
						type=float, default=0.9)
	parser.add_argument('--qa', help="PrintOutQAHistos", default=False, action='store_true')

	parser.add_argument('--dRmax', default=0.25, type=float)
	parser.add_argument('--alpha', default=0, type=float)


	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.4
	jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
	jet_selector = fj.SelectorPtMin(args.py_pthatmin) & fj.SelectorPtMax(1000.0) & fj.SelectorAbsEtaMax(args.eta - jet_R0)
	# jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(200.0) &fj.SelectorAbsEtaMax(1)
	print(jet_def)

	all_jets = []

	# mycfg = ['PhaseSpace:pThatMin = 80']
	# mycfg = ['PhaseSpace:pThatMin = 40']	
	mycfg = ['']	
	if args.ignore_mycfg:
		mycfg = []
	pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
	if not pythia:
		perror("pythia initialization failed.")
		return

	jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
	lund_gen = fjcontrib.LundGenerator(jet_def_lund)
	print (lund_gen.description())
	dy_groomer = fjcontrib.DynamicalGroomer(jet_def_lund)
	print (dy_groomer.description())

	# sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
	sd01 = fjcontrib.SoftDrop(0, 0.1, 1.0)
	print (sd01)
	sd02 = fjcontrib.SoftDrop(0, 0.2, 1.0)
	print (sd02)

	# jet_def_rc01 = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
	# jet_def_rc02 = fj.JetDefinition(fj.cambridge_algorithm, 0.2)
	# print (jet_def_rc01)
	# print (jet_def_rc02)
	# rc = fjcontrib.Recluster(jet_def_rc, True)

	jet_def_rc01 = fj.JetDefinition(fj.antikt_algorithm, 0.1)
	jet_def_rc02 = fj.JetDefinition(fj.antikt_algorithm, 0.2)
	print (jet_def_rc01)
	print (jet_def_rc02)
	#rc = fjcontrib.Recluster(jet_def_rc, True)

	# tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = 'leadsj_vs_x.root')
	tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = args.output)
	tgbkg = None
	be = None
	if args.enable_background:
		# ROOT.gSystem.Load("libpyjetty_TennGen.dylib")
		# tgbkg = ROOT.TennGen() # //constructor
		# tgbkg.SetCentralityBin(args.cent_bin) # //centraility bin 0 is the  0-5 % most central bin
		# tgbkg.SetRandomSeed(args.seed) # //setting the seed
		# tgbkg.SetHarmonics(args.harmonics) # // set harmonics flag (0 : v1 - v5) , (1 : v2 - v5) , (2: v3 - v5) , (3: v1 - v4) , (4: v1 - v3) , (5: uniform dN/dphi no harmonics) , (6 : v1 - v2 , v4 - v5) , (7 : v1 - v3 , v5) , (8 : v1 , v3 - v5) , (9 : v1 only) , (10 : v2 only) , (11 : v3 only) , (12 : v4 only) , (13 : v5 only)
		# tgbkg.SetEtaRange(args.eta) # //set eta range must be uniform (e.g. |eta| < 0.9, which is ALICE TPC fiducial acceptance)
		# tgbkg.PrintOutQAHistos(args.qa) #
		# tgbkg.InitializeBackground() #

		from pyjetty.mputils import BoltzmannEvent
		be = BoltzmannEvent(mean_pt=0.7, multiplicity=2000 * args.eta * 2, max_eta=max_eta, max_pt=100)
		print(be)

		from pyjetty.mputils import CEventSubtractor, CSubtractorJetByJet
		cs = CEventSubtractor(alpha=args.alpha, max_distance=args.dRmax, max_eta=args.eta, bge_rho_grid_size=0.25, max_pt_correct=100)
		print(cs)


	if args.nev < 100:
		args.nev = 100
	for i in tqdm.tqdm(range(args.nev)):
		if not pythia.next():
			continue
		# parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
		partons = pythiafjext.vectorize_select(pythia, [pythiafjext.kParton], 0, True)
		parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
		# parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)
		jets = jet_selector(jet_def(parts))

		# for j in tqdm.tqdm(jets):
		for j in jets:
			j_type = match_dR(j, partons, jet_R0 / 2.)
			if j_type[0] is None:
				if args.nw:
					continue
				pwarning('Jet with no parton label')
				continue

			j_sd02 = sd02.result(j)
			sd02_info = fjcontrib.get_SD_jet_info(j_sd02)
			j_sd01 = sd01.result(j)
			sd01_info = fjcontrib.get_SD_jet_info(j_sd01)

			rc_sjets01 = fj.sorted_by_pt(jet_def_rc01(j.constituents()))
			rc_sjets02 = fj.sorted_by_pt(jet_def_rc02(j.constituents()))
			tw.fill_branches(	j 			= j, 
								lund 		= [ls for ls in lund_gen.result(j)], 
								dyg1 		= dy_groomer.result(j, 1), 

								sd01 		= j_sd01, 
								sd01_z 		= sd01_info.z, 
								sd01_mu 	= sd01_info.mu, 
								sd01_Delta 	= sd01_info.dR, 

								sd02 		= j_sd02, 
								sd02_z 		= sd02_info.z, 
								sd02_mu 	= sd02_info.mu, 
								sd02_Delta 	= sd02_info.dR, 

								# breaking compatibility
								# sd 			= j_sd, 
								# sd_z 		= sd_info.z, 
								# sd_mu 		= sd_info.mu, 
								# sd_Delta 	= sd_info.dR, 

								lsjet01 	= rc_sjets01[0],
								nsjet01    	= len(rc_sjets01),
								sjet01     	= rc_sjets01,
								lsjet02 	= rc_sjets02[0],
								nsjet02    	= len(rc_sjets02),
								sjet02     	= rc_sjets02,

								ppid       	= j_type[0],
								pquark     	= j_type[1],
								pglue      	= j_type[2], # this is redundancy

								pycode 		= pythia.info.code(),
								pysigmagen  = pythia.info.sigmaGen(),
								pysigmaerr  = pythia.info.sigmaErr(),
								pyid1       = pythia.info.id1pdf(),
								pyid2       = pythia.info.id1pdf(),
								pyx1 	    = pythia.info.x1pdf(),
								pyx2       	= pythia.info.x2pdf(),
								pypdf1      = pythia.info.pdf1(),
								pyQfac 		= pythia.info.QFac(),
								pyalphaS 	= pythia.info.alphaS(),

								pypthat 	= pythia.info.pTHat(),
								pymhat 		= pythia.info.mHat()

							 )
			if be:
				bg_parts = be.generate(offset=10000)
				full_event = bg_parts
				tmp = [full_event.push_back(psj) for psj in j.constituents()]
				if cs:
					cs_parts = cs.process_event(full_event)
					rho = cs.bge_rho.rho()
					bg_jets = fj.sorted_by_pt(jet_def(cs_parts))
					for bj in bg_jets:
						if fjtools.matched_pt(bj, j) > 0.5:
							pass

			tw.fill_tree()

	pythia.stat()

	tw.write_and_close()
def main():
	parser = argparse.ArgumentParser(description='pythia8 in python', prog=os.path.basename(__file__))
	parser.add_argument('-i', '--input', help='input file', default='low', type=str, required=True)
	parser.add_argument('--nev', help='number of events', default=-1, type=int)
	parser.add_argument('--jetptcut', help='jet pt cut', default=10., type=float)	
	args = parser.parse_args()	

	files = find_files(args.input, "*.txt")
	print(files)
	print('[i] number of files found {}'.format(len(files)))

	# jet finder
	# print the banner first
	fj.ClusterSequence.print_banner()
	print()
	jet_R0 = 0.4
	jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
	jet_selector = fj.SelectorPtMin(10.0) & fj.SelectorAbsEtaMax(3)

	if args.nev < 0:
		args.nev = len(files)
	if args.nev > len(files):
		args.nev = len(files)	

	print("[i] running on {}".format(files[:args.nev]))
	all_jets = []
	for fin in tqdm.tqdm(files[:args.nev]):
		input_hepmc = hepmc3ext.YuukaRead(fin)
		if input_hepmc.failed():
			print ("[error] unable to read from {}".format(fin))
			sys.exit(1)
		if input_hepmc.nextEvent():
			jets = find_jets_hepmc(jet_def, jet_selector, input_hepmc, accept_status=[62])
			all_jets.extend(jets)

	jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
	lund_gen = fjcontrib.LundGenerator(jet_def_lund)

	print('[i] making lund diagram for all jets...')
	lunds = [lund_gen.result(j) for j in all_jets]

	print('[i] reclustering and using soft drop...')
	jet_def_rc = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
	print('[i] Reclustering:', jet_def_rc)

	all_jets_sd = []
	rc = fjcontrib.Recluster(jet_def_rc, True)
	sd = fjcontrib.SoftDrop(0, 0.1, 1.0)

	fout = ROOT.TFile('hepmc_jetreco.root', 'recreate')
	lbins = logbins(1., 500, 50)
	hJetPt04 = ROOT.TH1D("hJetPt04", "hJetPt04", 50, lbins)
	hJetPt04sd = ROOT.TH1D("hJetPt04sd", "hJetPt04sd", 50, lbins)
	[hJetPt04.Fill(j.perp()) for j in all_jets]
	[hJetPt04sd.Fill(j.perp()) for j in all_jets_sd]
	hLund = ROOT.TH2D("hLund", "hLund", 60, 0, 6, 100, -4, 5)
	lunds = [lund_gen.result(j) for j in all_jets if j.perp() > args.jetptcut]
	jsel = [j for j in all_jets if j.perp() > args.jetptcut]
	print('[i] {} jets above {} GeV/c'.format(len(jsel), args.jetptcut))
	for l in lunds:
		for s in l:
			hLund.Fill(math.log(1./s.Delta()), math.log(s.kt()))
	if len(jsel) > 0:
		hLund.Scale(1./len(jsel))
	fout.Write()
Exemple #24
0
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')

    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.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    jet_selector = fj.SelectorPtMin(100.0) & fj.SelectorAbsEtaMax(1)
    print(jet_def)

    all_jets = []

    mycfg = ['PhaseSpace:pThatMin = 100']
    if args.ignore_mycfg:
        mycfg = []
    pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
    if not pythia:
        print("[e] pythia initialization failed.")
        return
    if args.nev < 10:
        args.nev = 10
    for i in tqdm.tqdm(range(args.nev)):
        if not pythia.next():
            continue
        attach_pythia_particle_info = True
        parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal],
                                             attach_pythia_particle_info)
        jets = jet_selector(jet_def(parts))
        for j in jets:
            gshop = fjcontrib.GroomerShop(j)
            # note these LundDeclustering objects can be streamed to a tree using RTreeStreamer
            dg01 = gshop.dynamical(0.1)
            idg01 = gshop.index(dg01)
            dg20 = gshop.dynamical(2.0)
            idg20 = gshop.index(dg20)
            max_kt = gshop.max_kt()
            imax_kt = gshop.index(max_kt)
            max_pt_softer = gshop.max_pt_softer()
            # check if the same split selected:
            # if dg01.pair() == dg20.pair():
            # or use indices
            if idg01 == idg20 and idg01 > 0:
                print('- interesting jet?:')
                print('  dg01         :', dg01.as_string(), 'index dg01:',
                      idg01)
                print('  dg20         :', dg20.as_string(), 'index dg20:',
                      idg20)
                print('  max_kt       :', max_kt.as_string(), 'index max_kt:',
                      imax_kt)
                print('  max_pt_softer:', max_pt_softer.as_string())
                print('  max_kappa    :', gshop.max_kappa().as_string())
                print('  max_tf       :', gshop.max_tf().as_string())
                print('  min_tf       :', gshop.min_tf().as_string())
                print('  max_z        :', gshop.max_z().as_string())

            print('softdrop check for jet:', j)
            sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
            j_sd = sd.result(j)
            #print('  |-> after soft drop pT={0:10.3f} delta={1:10.3f}'.format(j_sd.perp(), j_sd.perp() - j.perp()))
            sd_info = fjcontrib.get_SD_jet_info(j_sd)
            print("  |-> SD jet params          z={} dR={} mu={}".format(
                sd_info.z, sd_info.dR, sd_info.mu))
            print('  |-> GroomerShop::soft_drop',
                  gshop.soft_drop(0, 0.1, 1.0).as_string())
            # or call with no radius param - will use max allowed
            # print('  |-> GroomerShop::soft_drop', gshop.soft_drop(0, 0.1).as_string())

    pythia.stat()
Exemple #26
0
def main():
    parser = argparse.ArgumentParser(
        description='read jewel to lund jets object and pickle with joblib',
        prog=os.path.basename(__file__))
    parser.add_argument('-d',
                        '--dir',
                        help='input directory',
                        default='',
                        type=str,
                        required=True)
    parser.add_argument('-n',
                        '--nfiles',
                        help='number of files',
                        default=-1,
                        type=int,
                        required=False)
    parser.add_argument('-p',
                        '--pattern',
                        help='input directory',
                        default='',
                        type=str,
                        required=True)
    parser.add_argument('--nev', help='number of events', default=-1, type=int)
    parser.add_argument('--clean',
                        help='remove output before writing',
                        action='store_true',
                        default=False)
    args = parser.parse_args()

    files = find_files(args.dir, args.pattern)

    # jet finder
    # print the banner first
    fj.ClusterSequence.print_banner()
    print()
    jet_R0 = 0.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(
        500.0) & fj.SelectorAbsEtaMax(2)

    jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
    lund_gen = fjcontrib.LundGenerator(jet_def_lund)

    sd = fjcontrib.SoftDrop(0, 0.1, 1.0)

    if args.nfiles > 0:
        rfiles = files[:args.nfiles]
    else:
        rfiles = files
    print("[i] found {} files - running for {} files".format(
        len(files), len(rfiles)))
    for ifile, fn in enumerate(tqdm.tqdm(rfiles)):
        # now lets read the HEPMC file and do some jet finding
        input_hepmc = hepmc2wrap.ReadHepMCFile(fn)
        if input_hepmc.failed():
            print("[error] unable to read from {}".format(fn))
            sys.exit(1)

        all_jets = []
        for iev in tqdm.tqdm(range(args.nev)):
            if input_hepmc.failed():
                break
            if input_hepmc.NextEvent():
                jets_hepmc = find_jets_hepmc(jet_def,
                                             jet_selector,
                                             input_hepmc,
                                             final=True)
                all_jets.extend(jets_hepmc)

        lunds = [lund_gen.result(j) for j in all_jets]
        lunds_pickle = [
            lundjet.LundJet(j, lund_gen.result(j)) for j in all_jets
        ]

        all_jets_sd = [sd.result(j) for j in all_jets]
        lunds_sd = [lund_gen.result(j) for j in all_jets_sd]
        lunds_sd_pickle = [
            lundjet.LundJet(j, lund_gen.result(j)) for j in all_jets_sd
        ]

        joblib_file = 'lund_{}.joblib'.format(os.path.basename(fn))
        it = 0
        while os.path.exists(joblib_file):
            if args.clean:
                os.unlink(joblib_file)
                break
            joblib_file = 'lund_{}_{}.joblib'.format(os.path.basename(fn), it)
            it += 1
        joblib.dump(lunds_pickle, joblib_file)

        joblib_file_sd = 'lund_sd_{}.joblib'.format(os.path.basename(fn))
        it = 0
        while os.path.exists(joblib_file_sd):
            if args.clean:
                os.unlink(joblib_file_sd)
                break
            joblib_file_sd = 'lund_sd_{}_{}.joblib'.format(
                os.path.basename(fn), it)
            it += 1
        joblib.dump(lunds_sd_pickle, joblib_file_sd)
Exemple #27
0
def main():
    parser = argparse.ArgumentParser(description='pythia8 in python',
                                     prog=os.path.basename(__file__))
    parser.add_argument('-i',
                        '--input',
                        help='input file',
                        default='low',
                        type=str,
                        required=True)
    parser.add_argument('--hepmc',
                        help='what format 2 or 3',
                        default=2,
                        type=int)
    parser.add_argument('--nev', help='number of events', default=10, type=int)
    args = parser.parse_args()

    ###
    # now lets read the HEPMC file and do some jet finding
    if args.hepmc == 3:
        input_hepmc = pyhepmc_ng.ReaderAscii(args.input)
    if args.hepmc == 2:
        input_hepmc = pyhepmc_ng.ReaderAsciiHepMC2(args.input)

    if input_hepmc.failed():
        print("[error] unable to read from {}".format(args.input))
        sys.exit(1)

    # jet finder
    # print the banner first
    fj.ClusterSequence.print_banner()
    print()
    jet_R0 = 0.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    jet_selector = fj.SelectorPtMin(10.0) & fj.SelectorPtMax(
        500.0) & fj.SelectorAbsEtaMax(3)

    all_jets = []
    event_hepmc = pyhepmc_ng.GenEvent()
    pbar = tqdm.tqdm(range(args.nev))
    while not input_hepmc.failed():
        ev = input_hepmc.read_event(event_hepmc)
        if input_hepmc.failed():
            break
        jets_hepmc = find_jets_hepmc(jet_def, jet_selector, event_hepmc)
        all_jets.extend(jets_hepmc)
        pbar.update()
        if pbar.n >= args.nev:
            break

    jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
    lund_gen = fjcontrib.LundGenerator(jet_def_lund)

    print('[i] making lund diagram for all jets...')
    lunds = [lund_gen.result(j) for j in all_jets]

    print('[i] listing lund plane points... Delta, kt - for {} selected jets'.
          format(len(all_jets)))
    for l in lunds:
        if len(l) < 1:
            continue
        print('- jet pT={0:5.2f} eta={1:5.2f}'.format(l[0].pair().perp(),
                                                      l[0].pair().eta()))
        print('  Deltas={}'.format([s.Delta() for s in l]))
        print('  kts={}'.format([s.kt() for s in l]))
        print()

    print('[i] reclustering and using soft drop...')
    jet_def_rc = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
    print('[i] Reclustering:', jet_def_rc)

    all_jets_sd = []
    rc = fjcontrib.Recluster(jet_def_rc, True)
    sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
    for i, j in enumerate(all_jets):
        j_rc = rc.result(j)
        print()
        print('- [{0:3d}] orig pT={1:10.3f} reclustered pT={2:10.3f}'.format(
            i, j.perp(), j_rc.perp()))
        j_sd = sd.result(j)
        print('  |-> after soft drop pT={0:10.3f} delta={1:10.3f}'.format(
            j_sd.perp(),
            j_sd.perp() - j.perp()))
        all_jets_sd.append(j_sd)
        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))

    fout = ROOT.TFile('hepmc_jetreco.root', 'recreate')
    lbins = logbins(1., 500, 50)
    hJetPt04 = ROOT.TH1D("hJetPt04", "hJetPt04", 50, lbins)
    hJetPt04sd = ROOT.TH1D("hJetPt04sd", "hJetPt04sd", 50, lbins)
    [hJetPt04.Fill(j.perp()) for j in all_jets]
    [hJetPt04sd.Fill(j.perp()) for j in all_jets_sd]
    hLund = ROOT.TH2D("hLund", "hLund", 60, 0, 6, 100, -4, 5)
    lunds = [lund_gen.result(j) for j in all_jets if j.perp() > 100]
    j100 = [j for j in all_jets if j.perp() > 100]
    print('{} jets above 100 GeV/c'.format(len(j100)))
    for l in lunds:
        for s in l:
            hLund.Fill(math.log(1. / s.Delta()), math.log(s.kt()))
    fout.Write()
Exemple #28
0
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')

    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.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    jet_selector = fj.SelectorPtMin(80.0) & fj.SelectorPtMax(
        100.0) & fj.SelectorAbsEtaMax(1)
    # jet_selector = fj.SelectorPtMin(40.0) & fj.SelectorPtMax(200.0) &fj.SelectorAbsEtaMax(1)
    print(jet_def)

    all_jets = []

    mycfg = ['PhaseSpace:pThatMin = 80']
    # mycfg = ['PhaseSpace:pThatMin = 40']
    if args.ignore_mycfg:
        mycfg = []
    pythia = pyconf.create_and_init_pythia_from_args(args, mycfg)
    if not pythia:
        perror("pythia initialization failed.")
        return

    jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, 1.0)
    lund_gen = fjcontrib.LundGenerator(jet_def_lund)
    print(lund_gen.description())
    dy_groomer = fjcontrib.DynamicalGroomer(jet_def_lund)
    print(dy_groomer.description())
    # sd = fjcontrib.SoftDrop(0, 0.1, 1.0)
    sd = fjcontrib.SoftDrop(0, 0.2, 1.0)
    print(sd)

    # jet_def_rc01 = fj.JetDefinition(fj.cambridge_algorithm, 0.1)
    # jet_def_rc02 = fj.JetDefinition(fj.cambridge_algorithm, 0.2)
    # print (jet_def_rc01)
    # print (jet_def_rc02)
    # rc = fjcontrib.Recluster(jet_def_rc, True)

    jet_def_rc01 = fj.JetDefinition(fj.antikt_algorithm, 0.1)
    jet_def_rc02 = fj.JetDefinition(fj.antikt_algorithm, 0.2)
    print(jet_def_rc01)
    print(jet_def_rc02)
    #rc = fjcontrib.Recluster(jet_def_rc, True)

    # tw = treewriter.RTreeWriter(name = 'lsjvsx', file_name = 'leadsj_vs_x.root')
    tw = treewriter.RTreeWriter(name='lsjvsx',
                                file_name='leadsj_vs_x_bias80.root')

    if args.nev < 100:
        args.nev = 100
    for i in tqdm.tqdm(range(args.nev)):
        if not pythia.next():
            continue
        # parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
        partons = pythiafjext.vectorize_select(pythia, [pythiafjext.kParton],
                                               0, True)
        parts = pythiafjext.vectorize_select(
            pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, False)
        # parts = pythiafjext.vectorize_select(pythia, [pythiafjext.kFinal], 0, False)
        jets = jet_selector(jet_def(parts))

        # for j in tqdm.tqdm(jets):
        for j in jets:
            j_type = match_dR(j, partons, jet_R0 / 2.)
            j_sd = sd.result(j)
            sd_info = fjcontrib.get_SD_jet_info(j_sd)
            rc_sjets01 = fj.sorted_by_pt(jet_def_rc01(j.constituents()))
            rc_sjets02 = fj.sorted_by_pt(jet_def_rc02(j.constituents()))
            tw.fill_branches(
                j=j,
                lund=[ls for ls in lund_gen.result(j)],
                dyg1=dy_groomer.result(j, 1),
                sd=j_sd,
                sd_z=sd_info.z,
                sd_mu=sd_info.mu,
                sd_Delta=sd_info.dR,
                lsjet01=rc_sjets01[0],
                nsjet01=len(rc_sjets01),
                sjet01=rc_sjets01,
                lsjet02=rc_sjets02[0],
                nsjet02=len(rc_sjets02),
                sjet02=rc_sjets02,
                ppid=j_type[0],
                pquark=j_type[1],
                pglue=j_type[2]  # this is redundancy
            )
            tw.fill_tree()

    pythia.stat()

    tw.write_and_close()
Exemple #29
0
def main():
    parser = argparse.ArgumentParser(description='jetscape in python', \
                                     prog=os.path.basename(__file__))
    parser.add_argument('-i', '--input', help='input file', \
                        default='low', type=str, required=True)
    parser.add_argument('--nev', help='number of events', \
                        default=1000, type=int)
    args = parser.parse_args()

    plot_hadrons = False
    plot_final_state_partons = True
    plot_nth_partons = False
    include_thermal_background = False

    # Initialize histogram dictionary
    hDict = initializeHistograms()

    # Use pyhepmc_ng to parse the HepMC file
    input_hepmc = pyhepmc_ng.ReaderAscii(args.input)
    if input_hepmc.failed():
        print("[error] unable to read from {}".format(args.input))
        sys.exit(1)

    # jet finder
    fj.ClusterSequence.print_banner()
    print()
    jetR = 0.4
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jetR)
    jet_selector = fj.SelectorPtMin(50.0) & fj.SelectorAbsEtaMax(2)

    # Jet re-clustering definition, for primary Lund plane
    jet_def_lund = fj.JetDefinition(fj.cambridge_algorithm, jetR)
    lund_gen = fjcontrib.LundGenerator(jet_def_lund)

    # Load thermal background: array of dataframes, one per event
    if include_thermal_background:
        thermal_background_array = get_thermal_background()

    # Loop through events
    all_jets_hadron = []
    all_jets_parton = []
    kt_shower_list = []
    event_hepmc = pyhepmc_ng.GenEvent()
    pbar = tqdm.tqdm(range(args.nev))
    while not input_hepmc.failed():
        ev = input_hepmc.read_event(event_hepmc)
        if input_hepmc.failed():
            nstop = pbar.n
            pbar.close()
            print('End of HepMC file at event {} '.format(nstop))
            break

        if plot_hadrons:
            hadrons = get_hadrons(event_hepmc)
            jets_hadron = find_jets(jet_def, jet_selector, hadrons)
            all_jets_hadron.extend(jets_hadron)
        if plot_final_state_partons:
            parton_list = get_final_partons(event_hepmc, hDict)
        elif plot_nth_partons:
            # Get the first parton in the shower (it is apparently the child of the initiating parton)
            first_parton = []
            first_parton.append(
                get_first_parton(event_hepmc, hDict).children[0])
            #print('-----------------------------')
            #print('First parton E: {}'.format(first_parton[0].momentum.e))

            n = 3
            parton_list = get_nth_partons(first_parton, n)
            #print('Total number of partons: {}'.format(len(parton_list)))

        fj_particles = []
        for parton in parton_list:
            fj_particles.append(
                fj.PseudoJet(parton.momentum.px, parton.momentum.py,
                             parton.momentum.pz, parton.momentum.e))

        # Append thermal background
        if include_thermal_background:
            thermal_background_df = thermal_background_array[0]
            thermal_background_fjparticles = get_fjparticles(
                thermal_background_df)
            fj_particles.extend(thermal_background_fjparticles)

        jets_parton = find_jets(jet_def, jet_selector, fj_particles)
        all_jets_parton.extend(jets_parton)

        # Analyze shower history
        max_kt = get_max_kt_shower(event_hepmc, hDict)
        kt_shower_list.append(max_kt)

        pbar.update()

        if pbar.n >= args.nev:
            pbar.close()
            print('{} event limit reached'.format(args.nev))
            break

    print('Constructing histograms...')

    if plot_hadrons:
        n_jets_hadron = len(all_jets_hadron)
        print('n_jets_hadron: {}'.format(n_jets_hadron))
    n_jets_parton = len(all_jets_parton)
    print('n_jets_parton: {}'.format(n_jets_parton))

    # Fill histogram
    if plot_hadrons:
        [fill_jet_histogram(hDict, jet) for jet in all_jets_hadron]

        # Fill Lund diagram
        # Note: l in lunds, l is the list of splittings in a given jet (following hardest splitting)
        lunds_hadron = [lund_gen.result(jet) for jet in all_jets_hadron]
        [
            fill_lund_histogram(hDict, "hLundHadron", splitting_list)
            for splitting_list in lunds_hadron
        ]
        hDict['hLundHadron'].Scale(1. / n_jets_hadron, "width")

    lunds_parton = [lund_gen.result(jet) for jet in all_jets_parton]
    [
        fill_lund_histogram(hDict, "hLundParton", splitting_list)
        for splitting_list in lunds_parton
    ]
    hDict['hLundParton'].Scale(1. / n_jets_parton, "width")

    for splitting_list in lunds_parton:
        max_kt_recluster = 0
        for split in splitting_list:
            kt = split.kt()
            if kt > max_kt_recluster:
                max_kt_recluster = kt
        hDict['hMaxKT_2D'].Fill(np.log(max_kt_recluster), np.log(kt))

    plot_histograms(hDict)
Exemple #30
0
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('--user-seed',
                        help='pythia seed',
                        default=1111,
                        type=int)
    parser.add_argument('--output', default="output.root", type=str)
    parser.add_argument('--beta', help='sd beta', default=0, type=float)
    parser.add_argument('--jetR', help='jet radius', default=0.4, type=float)
    args = parser.parse_args()

    if args.user_seed < 0:
        args.user_seed = 1111
    pinfo('user seed for pythia', args.user_seed)
    # mycfg = ['PhaseSpace:pThatMin = 100']
    mycfg = ['Random:setSeed=on', 'Random:seed={}'.format(args.user_seed)]
    mycfg.append('HadronLevel:all=off')
    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.jetR
    jet_def = fj.JetDefinition(fj.antikt_algorithm, jet_R0)
    print(jet_def)

    # hadron level - ALICE
    max_eta_hadron = 0.9
    pwarning('max eta for particles after hadronization set to',
             max_eta_hadron)
    parts_selector_h = fj.SelectorAbsEtaMax(max_eta_hadron)
    jet_selector = fj.SelectorPtMin(20.0) & fj.SelectorPtMax(
        100.0) & fj.SelectorAbsEtaMax(max_eta_hadron - 1.05 * jet_R0)

    max_eta_parton = max_eta_hadron + 2. * jet_R0
    pwarning('max eta for partons set to', max_eta_parton)
    parts_selector_p = fj.SelectorAbsEtaMax(max_eta_parton)

    outf = ROOT.TFile(
        args.output.replace('.root', '_beta{}.root'.format(args.beta)),
        'recreate')
    outf.cd()
    t = ROOT.TTree('t', 't')
    tw = RTreeWriter(tree=t)

    # event loop
    for iev in tqdm.tqdm(range(args.nev)):
        if not pythia.next():
            continue

        parts_pythia_p = pythiafjext.vectorize_select(pythia,
                                                      [pythiafjext.kFinal], 0,
                                                      True)
        parts_pythia_p_selected = parts_selector_p(parts_pythia_p)

        hstatus = pythia.forceHadronLevel()
        if not hstatus:
            pwarning('forceHadronLevel false event', iev)
            continue
        # parts_pythia_h = pythiafjext.vectorize_select(pythia, [pythiafjext.kHadron, pythiafjext.kCharged])
        parts_pythia_h = pythiafjext.vectorize_select(pythia,
                                                      [pythiafjext.kFinal], 0,
                                                      True)
        parts_pythia_h_selected = parts_selector_h(parts_pythia_h)

        parts_pythia_hch = pythiafjext.vectorize_select(
            pythia, [pythiafjext.kFinal, pythiafjext.kCharged], 0, True)
        parts_pythia_hch_selected = parts_selector_h(parts_pythia_hch)

        # pinfo('debug partons...')
        # for p in parts_pythia_p_selected:
        # 	pyp = pythiafjext.getPythia8Particle(p)
        # 	print(pyp.name())
        # pinfo('debug hadrons...')
        # for p in parts_pythia_h_selected:
        # 	pyp = pythiafjext.getPythia8Particle(p)
        # 	print(pyp.name())
        # pinfo('debug ch. hadrons...')
        # for p in parts_pythia_hch_selected:
        # 	pyp = pythiafjext.getPythia8Particle(p)
        # 	print(pyp.name())

        # parts = pythiafjext.vectorize(pythia, True, -1, 1, False)
        jets_p = fj.sorted_by_pt(jet_def(parts_pythia_p))
        jets_h = fj.sorted_by_pt(jet_def(parts_pythia_h))
        jets_ch_h = fj.sorted_by_pt(jet_selector(jet_def(parts_pythia_hch)))

        sd = fjcontrib.SoftDrop(args.beta, 0.1, jet_R0)

        for i, jchh in enumerate(jets_ch_h):
            # match hadron (full) jet
            for j, jh in enumerate(jets_h):
                drhh = jchh.delta_R(jh)
                if drhh < jet_R0 / 2.:
                    # match parton level jet
                    for k, jp in enumerate(jets_p):
                        dr = jh.delta_R(jp)
                        if dr < jet_R0 / 2.:
                            jchh_sd = sd.result(jchh)
                            jchh_sd_info = fjcontrib.get_SD_jet_info(jchh_sd)
                            jh_sd = sd.result(jh)
                            jh_sd_info = fjcontrib.get_SD_jet_info(jh_sd)
                            jp_sd = sd.result(jp)
                            jp_sd_info = fjcontrib.get_SD_jet_info(jp_sd)

                            # pwarning('event', iev)
                            # pinfo('matched jets: ch.h:', jchh.pt(), 'h:', jh.pt(), 'p:', jp.pt(), 'dr:', dr)

                            tw.fill_branch('iev', iev)
                            tw.fill_branch('ch', jchh)
                            tw.fill_branch('h', jh)
                            tw.fill_branch('p', jp)

                            tw.fill_branch('p_zg', jp_sd_info.z)
                            tw.fill_branch('p_Rg', jp_sd_info.dR)
                            tw.fill_branch('p_thg', jp_sd_info.dR / jet_R0)
                            tw.fill_branch('p_mug', jp_sd_info.mu)

                            tw.fill_branch('h_zg', jh_sd_info.z)
                            tw.fill_branch('h_Rg', jh_sd_info.dR)
                            tw.fill_branch('h_thg', jh_sd_info.dR / jet_R0)
                            tw.fill_branch('h_mug', jh_sd_info.mu)

                            tw.fill_branch('ch_zg', jchh_sd_info.z)
                            tw.fill_branch('ch_Rg', jchh_sd_info.dR)
                            tw.fill_branch('ch_thg', jchh_sd_info.dR / jet_R0)
                            tw.fill_branch('ch_mug', jchh_sd_info.mu)

                            tw.fill_tree()
            #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()