コード例 #1
0
ファイル: jets_epps16.py プロジェクト: roylemmon/pyjetty
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()
コード例 #2
0
ファイル: reco_jets_epps16.py プロジェクト: roylemmon/pyjetty
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"))
コード例 #3
0
def main():
    nevents = 1000
    sconfig_pythia = [
        "Beams:eCM = 8000.", "HardQCD:all = on", "PhaseSpace:pTHatMin = 100.",
        "PDF:pset = LHAPDF6:EPPS16nlo_CT14nlo_Pb208/0"
    ]
    #  "LHAPDF6:member = 901300"]
    # "PDF:pSet = EPPS16nlo_CT14nlo_Pb208"]
    if 'gen' in sys.argv:
        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 'write' in sys.argv:
        pythia = create_and_init_pythia(sconfig_pythia)
        if not pythia:
            return
        pyhepmcwriter = mp.Pythia8HepMCWrapper("test_pythia8_hepmc.dat")
        all_jets = []
        for iEvent in tqdm(range(nevents), 'event'):
            if not pythia.next(): continue
            pyhepmcwriter.fillEvent(pythia)
        print("[i] done writing to test_pythia8_hepmc.dat")

    if 'read' in sys.argv:
        import pyhepmc_ng
        input = pyhepmc_ng.ReaderAsciiHepMC2("test_pythia8_hepmc.dat")
        if input.failed():
            print("[error] unable to read from test_pythia8_hepmc.dat")
            return
        event = pyhepmc_ng.GenEvent()
        pbar = tqdm(range(nevents))
        while not input.failed():
            p = input.read_event(event)
            if input.failed():
                # print ("[i] reading done.")
                break
            pbar.update()
コード例 #4
0
ファイル: hepmc2antuple_tn.py プロジェクト: roylemmon/pyjetty
    def main(self):

        if self.hepmc == 3:
            input_hepmc = pyhepmc_ng.ReaderAscii(self.input)
        if self.hepmc == 2:
            input_hepmc = pyhepmc_ng.ReaderAsciiHepMC2(self.input)

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

        event_hepmc = pyhepmc_ng.GenEvent()

        while not input_hepmc.failed():
            ev = input_hepmc.read_event(event_hepmc)
            if input_hepmc.failed():
                break

            self.fill_event(event_hepmc)
            self.increment_event()
            if self.nev > 0 and self.ev_id > self.nev:
                break

        self.finish()
コード例 #5
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()
コード例 #6
0
ファイル: hepmc2antuple_tn.py プロジェクト: roylemmon/pyjetty
def main():
    parser = argparse.ArgumentParser(
        description='hepmc to ALICE Ntuple format',
        prog=os.path.basename(__file__))
    parser.add_argument('-i',
                        '--input',
                        help='input file',
                        default='',
                        type=str,
                        required=True)
    parser.add_argument('-o',
                        '--output',
                        help='output root file',
                        default='',
                        type=str,
                        required=True)
    parser.add_argument('--as-data',
                        help='write as data - tree naming convention',
                        action='store_true',
                        default=False)
    parser.add_argument('--hepmc',
                        help='what format 2 or 3',
                        default=2,
                        type=int)
    parser.add_argument('--nev', help='number of events', default=-1, type=int)
    args = parser.parse_args()

    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)

    outf = ROOT.TFile(args.output, 'recreate')
    outf.cd()
    tdf = ROOT.TDirectoryFile('PWGHF_TreeCreator', 'PWGHF_TreeCreator')
    tdf.cd()
    if args.as_data:
        t_p = ROOT.TNtuple(
            'tree_Particle', 'tree_Particle',
            'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticlePID')
    else:
        t_p = ROOT.TNtuple(
            'tree_Particle_gen', 'tree_Particle_gen',
            'run_number:ev_id:ParticlePt:ParticleEta:ParticlePhi:ParticlePID')
    t_e = ROOT.TNtuple('tree_event_char', 'tree_event_char',
                       'run_number:ev_id:z_vtx_reco:is_ev_rej')

    # run number will be a double - file size in MB
    run_number = os.path.getsize(args.input) / 1.e6
    ev_id = 0

    # unfortunately pyhepmc_ng does not provide the table
    # pdt = pyhepmc_ng.ParticleDataTable()
    # use ROOT instead
    pdg = ROOT.TDatabasePDG()

    event_hepmc = pyhepmc_ng.GenEvent()

    if args.nev > 0:
        pbar = tqdm.tqdm(range(args.nev))
    else:
        pbar = tqdm.tqdm()

    while not input_hepmc.failed():
        ev = input_hepmc.read_event(event_hepmc)
        if input_hepmc.failed():
            break

        fill_event(run_number, ev_id, event_hepmc, t_e, t_p, pdg)

        ev_id = ev_id + 1
        pbar.update()
        if args.nev > 0 and ev_id > args.nev:
            break

    outf.Write()
    outf.Close()