def create_root_file(output_name, tree_name, truth_vertex_info, track_info, truth_track_info): tfile = ROOT.TFile(output_name, "RECREATE") ttree = ROOT.TTree(tree_name, "") writer = Writer(ttree) #truth vertex info writer.add("pvr_x") writer.add("pvr_y") writer.add("pvr_z") writer.add("ntrks_prompt") #leave SV stuff empty for now writer.add("svr_x") writer.add("svr_y") writer.add("svr_z") writer.add("svr_pvr") #truth info for tracks writer.add("prt_pid") #leave empty for now writer.add("prt_hits") #leave empty for now #fill these below writer.add("prt_px") writer.add("prt_py") writer.add("prt_pz") writer.add("prt_e") #can be zeros writer.add("prt_pid") #leave empty for now writer.add("prt_hits") #can be zeros writer.add("prt_x") writer.add("prt_y") writer.add("prt_z") writer.add("prt_pvr") #actual track info writer.add("recon_x") writer.add("recon_y") writer.add("recon_z") writer.add("recon_tx") writer.add("recon_ty") writer.add("recon_chi2") # fill with ones for now #POCA info (fill below with zeros for now) writer.add("recon_pocax") writer.add("recon_pocay") writer.add("recon_pocaz") writer.add("recon_sigmapocaxy") n_events = len(truth_vertex_info.groupby("evt")) counter = 1 for vrtx, trk, truthtrk in zip(truth_vertex_info.groupby("evt"), track_info.groupby("evt"), truth_track_info.groupby("evt")): prt_pvr = [] trk_nlist = truthtrk[1].index.get_level_values("num").tolist() trk_nlists = vrtx[1]["trk_nums"].tolist() vrtx_nlist = vrtx[1].index.get_level_values("num").tolist() nfound = 0 truth_dict = {} for i in range(len(trk_nlists)): for j in trk_nlists[i]: truth_dict[trk_nlist[j]] = vrtx_nlist[i] nfound += 1 prt_pvr_list = list(truth_dict.values()) print('Event ', vrtx[0], ' Processed (', counter, '/', n_events, ')') xlist = vrtx[1]["x"].tolist() ylist = vrtx[1]["y"].tolist() zlist = vrtx[1]["z"].tolist() nlist = vrtx[1]["n_trks"].tolist() for i in range(len(xlist)): writer["pvr_x"].append(xlist[i]) writer["pvr_y"].append(ylist[i]) writer["pvr_z"].append(zlist[i]) writer["ntrks_prompt"].append(nlist[i]) # currently not used d0 = np.asarray(trk[1]["d0"]).tolist() phi = np.asarray(trk[1]["phi"]).tolist() z0 = np.asarray(trk[1]["z0"]).tolist() cottheta = np.asarray(trk[1]["cottheta"]).tolist() arr1 = np.multiply(d0, np.cos(np.add(phi, -np.pi / 2))) arr2 = np.multiply(d0, np.sin(np.add(phi, -np.pi / 2))) loc = np.stack((np.multiply(d0, np.cos(np.add(phi, -np.pi / 2))), np.multiply(d0, np.sin(np.add(phi, -np.pi / 2))), z0)) direc = np.stack((np.cos(phi), np.sin(phi))) direc = np.divide(direc, np.linalg.norm(direc, axis=1).reshape(2, 1)) for i in range(len(loc[0])): writer["recon_x"].append(loc[0][i]) writer["recon_y"].append(loc[1][i]) writer["recon_z"].append(loc[2][i]) writer["recon_tx"].append(direc[0][i]) writer["recon_ty"].append(direc[1][i]) writer["recon_chi2"].append(0) writer["recon_pocax"].append(1) writer["recon_pocay"].append(1) writer["recon_pocaz"].append(1) writer["recon_sigmapocaxy"].append(1) pxlist = truthtrk[1]["px"].tolist() pylist = truthtrk[1]["py"].tolist() pzlist = truthtrk[1]["pz"].tolist() xlist = truthtrk[1]["vertexX"].tolist() ylist = truthtrk[1]["vertexY"].tolist() zlist = truthtrk[1]["vertexZ"].tolist() for i in range(len(pxlist)): writer["prt_px"].append(pxlist[i]) writer["prt_py"].append(pylist[i]) writer["prt_pz"].append(pzlist[i]) writer["prt_e"].append(0) writer["prt_x"].append(xlist[i]) writer["prt_y"].append(ylist[i]) writer["prt_z"].append(zlist[i]) writer["prt_hits"].append(0) writer["prt_pvr"].append(prt_pvr_list[i]) ttree.Fill() writer.clear() counter += 1 ttree.Print() ttree.Write(ttree.GetName(), ROOT.TObject.kOverwrite) tfile.Close()
def run(lname, tEvt): name = "{}.root".format(lname) myhash = abs(hash(lname)) % 900000000 print("Producing", tEvt, "tracks to", name, "with hash:", myhash) # Initialize Pythia. random = ROOT.TRandom3(myhash) pythia = pythia8.Pythia("", False) pythia.readString("Random:seed = {}".format(myhash)) pythia.readString("Print:quiet = on") pythia.readString("SoftQCD:all = on") pythia.init() module = velo.ModuleMaterial("dat/run3.root") rffoil = velo.FoilMaterial("dat/run3.root") scatter = Scatter() # Create the output TFile and TTree. tfile = ROOT.TFile(name, "RECREATE") ttree = ROOT.TTree("data", "") # New interface # Create the writer handler and add branches writer = Writer(ttree) writer.add("pvr_x") writer.add("pvr_y") writer.add("pvr_z") writer.add("svr_x") writer.add("svr_y") writer.add("svr_z") writer.add("svr_pvr") writer.add("hit_x") writer.add("hit_y") writer.add("hit_z") writer.add("hit_prt") writer.add("prt_pid") writer.add("prt_px") writer.add("prt_py") writer.add("prt_pz") writer.add("prt_e") writer.add("prt_x") writer.add("prt_y") writer.add("prt_z") writer.add("prt_hits") writer.add("prt_pvr") writer.add("ntrks_prompt") number_rejected_events = 0 # Fill the events. ttime = 0 for iEvt in range(tEvt): start = time.time() ipv = 0 npv = 0 target_npv = np.random.poisson(7.6) iEvt += 1 while npv < target_npv: if not pythia.next(): continue # All distance measurements are in units of mm xPv, yPv, zPv = ( random.Gaus(0, 0.055), random.Gaus(0, 0.055), random.Gaus(100, 63), ) # normal LHCb operation # pvr x and y spead can be found https://arxiv.org/pdf/1410.0149.pdf page 42. z dependent ## [-1000,-750, -500, -250] # mm writer["pvr_x"].append(xPv) writer["pvr_y"].append(yPv) writer["pvr_z"].append(zPv) number_of_detected_particles = 0 # find heavy flavor SVs for prt in pythia.event: if not heavyFlavor(prt.id()): continue # TODO: require particles with hits from the SVs writer["svr_x"].append(prt.xDec() + xPv) writer["svr_y"].append(prt.yDec() + yPv) writer["svr_z"].append(prt.zDec() + zPv) writer["svr_pvr"].append(ipv) for prt in pythia.event: if not prt.isFinal or prt.charge() == 0: continue if not prtStable(prt.id()): continue if abs(prt.zProd()) > 1000: continue if (prt.xProd() ** 2 + prt.yProd() ** 2) ** 0.5 > 40: continue if prt.pAbs() < 0.1: continue prt.xProd( prt.xProd() + xPv ) # Need to change the origin of the event before getting the hits prt.yProd(prt.yProd() + yPv) prt.zProd(prt.zProd() + zPv) hits = Hits(module, rffoil, scatter, prt) if len(hits) == 0: continue if len(hits) > 2 and abs(zPv - prt.zProd()) < 0.001: number_of_detected_particles += 1 # if prt.pAbs() < 0.2: print 'slow!', prt.pAbs(), prt.id() writer["prt_pid"].append(prt.id()) writer["prt_px"].append(prt.px()) writer["prt_py"].append(prt.py()) writer["prt_pz"].append(prt.pz()) writer["prt_e"].append(prt.e()) writer["prt_x"].append(prt.xProd()) writer["prt_y"].append(prt.yProd()) writer["prt_z"].append(prt.zProd()) writer["prt_pvr"].append(ipv) writer["prt_hits"].append(len(hits)) for xHit, yHit, zHit in hits: # xHit_recorded, yHit_recorded, zHit_recorded = np.random.uniform(-0.0275,0.0275)+xHit, np.random.uniform(-0.0275,0.0275)+yHit, zHit # normal xHit_recorded, yHit_recorded, zHit_recorded = ( np.random.normal(0, 0.012) + xHit, np.random.normal(0, 0.012) + yHit, zHit, ) # normal writer["hit_x"].append(xHit_recorded) writer["hit_y"].append(yHit_recorded) writer["hit_z"].append(zHit_recorded) writer["hit_prt"].append(len(writer["prt_e"]) - 1) # if number_of_detected_particles < 5: iEvt -= 1; number_rejected_events+=1; continue writer["ntrks_prompt"].append(number_of_detected_particles) ipv += 1 if number_of_detected_particles > 0: npv += 1 itime = time.time() - start ttime += itime if iEvt % 100 == 0: print( "{} Evt {}/{}, {:3} PVs, {:3} tracks in {:.3} s".format( name, iEvt, tEvt, npv, len(writer["pvr_z"]), itime ) ) writer.write() # Write and close the TTree and TFile. ttree.Print() ttree.Write(ttree.GetName(), ROOT.TObject.kOverwrite) tfile.Close() return ttime