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
