CreateDataFast.py

 
from itertools import compress
import pandas as pd
import numpy as np
import ROOT
from pathlib import Path
path = Path('~/Dokumente/jun22_Test/').expanduser()

EventsPerFile = 50
Nfiles = 5
events_per_file = 50
P_e = np.load("P_e.npy")
m_e = 0.00051099891 #GeV

def CreatePapa():
    FileNUMBERS = range(1,Nfiles +1)


   
    for i in FileNUMBERS:

        try:
            infile= ROOT.TFile.Open("~/Dokumente/jun22_Test/TrackTree"+ str(i).zfill(5)+ ".root","read")
            rec =  infile.Get("t4")
        except Exception:
                  print("Cant open File"+ "~/Dokumente/jun22_Test/TrackTree"+ str(i).zfill(5)+ ".root") 
                  continue
              
                
        if not hasattr(rec , 'GetEntry'):
            continue
        
        EventList =   np.ndarray(shape=(EventsPerFile,54, 700), dtype='float')
        for event in range(EventsPerFile):
            rec.GetEntry(event)
            print('\r'+str((i-FileNUMBERS[0] + event/EventsPerFile)*100/len(FileNUMBERS))[:4]+"%  of "+ str(len(FileNUMBERS)*events_per_file) +" Events", end='\r')
            Dict = {}
            Dict["TrackNumber"] = int(rec.FindLeaf("TrackNumber").GetValue(0))*np.ones((700))
    
            for leaf in  ["fPx", "fPy", "fPz","NoTrdHit", "TRDe1", "TRDe2", "TRDe3", "TRDe4","fLength","ToFTime","fAaxis", "fBaxis", "fPhi","fCenterX","fCenterY","NoRichHit","fChi2","RPfX","RPfY","StsHts","Ann","MvdHts","fCharge","PDG","MotherPDG","MotherId",  "XclosestHitMvd", "YclosestHitMvd", "ZclosestHitMvd", "XclosestHit", "YclosestHit", "ZclosestHit", "chi2NDFtrd", "chi2NDFsts","fTrackChi2NDF","DalitzCode", "StsHitsMC", "MvdHitsMC","TrdHitsMC"]:
                Dict[leaf] = np.array([ rec.FindLeaf(leaf).GetValue(track) for track in range(700)])
            
            Dict["distance"] = np.array([np.sqrt((RPx-Cx)**2+(RPy-Cy)**2) for RPx ,RPy, Cx, Cy in zip(Dict["RPfX"], Dict["RPfY"], Dict["fCenterX"],Dict["fCenterY"])])
            Dict["E"] = np.array([ np.sqrt(m_e**2+ Px**2+Py**2+Pz**2) for Px, Py, Pz in zip(Dict["fPx"],Dict["fPy"],Dict["fPz"])])
            Dict["P"] = np.array([np.sqrt(Px**2 + Py**2 + Pz**2)   for Px, Py, Pz in zip(Dict["fPx"],Dict["fPy"],Dict["fPz"])])
            Dict["theta"] = np.array([np.arctan2(Pz,np.sqrt(Px**2 + Py**2))    for Px, Py, Pz in zip(Dict["fPx"],Dict["fPy"],Dict["fPz"])])
            Dict["phi"] = np.array([ np.arctan2(Py,Px)  for Px, Py, Pz in zip(Dict["fPx"],Dict["fPy"],Dict["fPz"])])
            Dict["Pt"] = np.array([ np.sqrt(Py**2+Px**2)  for Px, Py, Pz in zip(Dict["fPx"],Dict["fPy"],Dict["fPz"])])
            Dict["Likelihood"] = []
            for e1, e2, e3, e4, Nhit in zip(Dict["TRDe1"] , Dict["TRDe2"] ,Dict["TRDe3"] ,Dict["TRDe4"], Dict["NoTrdHit"] ):
                index = 10**6*4*np.array([e1, e2, e3, e4])[:int(Nhit)]
                index[index> 319] = 319
                P_ele = np.prod([P_e[n] for n in index.astype(int)])
                P_pio =  np.prod([1-P_e[n] for n in index.astype(int)])
                Dict["Likelihood"].append(P_ele/(P_ele+P_pio))
            Dict["Likelihood"] = np.array(Dict["Likelihood"])
            Dict["beta"] = np.array([l/100/((t)*10**(-9))/(2.99792458*10**8) if t!= 0 else 0 for l, t in zip(Dict["fLength"],Dict["ToFTime"])])
            Dict['DeltaBetaEL'] = np.array([ beta - P/ E  for beta, P, E in zip(Dict["beta"] ,Dict["P"],Dict["E"])  ])
            Dict["fPhi"] = np.array([phi+np.pi/2 for phi in Dict["fPhi"]])
            Dict["fCenterY"] = np.array([y-100 if y > 0 else y+100  for y in  Dict["fCenterY"] ])
            Dict["RadialPosition"] =  np.array([np.sqrt(y**2 + x**2) for x,y in zip(Dict["fCenterX"] ,Dict["fCenterY"] )])
            Dict["RadialAngle"] = np.array([np.arctan2(y, x ) for x,y in zip(Dict["fCenterX"] ,Dict["fCenterY"])])
            Dict["fChi2NDF"] = np.array([chi/(Nhit-5) if Nhit > 5 else chi/999 for chi, Nhit in zip(Dict["fChi2"],Dict["NoRichHit"])])
            Dict["distClosestMvd"] = np.array([np.sqrt(x**2+y**2+z**2) for x,y,z in zip(Dict["XclosestHitMvd"], Dict["YclosestHitMvd"], Dict["ZclosestHitMvd"])])
            Dict["distClosestSts"] = np.array([np.sqrt(x**2+y**2+z**2) for x,y,z in zip(Dict[ "XclosestHit"], Dict["YclosestHit"], Dict["ZclosestHit"])])
            EventList[event,:,:] =  np.array(list(Dict.values()))
            
       # np.save(path/"PapaBIB", np.array(list(Dict.keys())))
       # np.save(path/str(i).zfill(5), EventList)
        
    


    
    
#################################################################################################################

def CratePDI():
    bib = np.load(path/"PapaBIB.npy")
    fileNameList =[ str(i).zfill(5) + ".npy" for i in range(1,100) if i != 89]
    DataArray = [pd.DataFrame({label: np.load(path/fileName).transpose(1, 0, 2).reshape(54,-1)[ i, :] for i,label in enumerate(bib) }) for fileName  in fileNameList]
    Data = pd.concat(DataArray)
    
    PDG = {'Electron': 11, 'Pion': 211, 'Myon':13, 'Kaon': 321, 'Proton': 2212, 'Hadron': 0}
    for label, value in PDG.items():
        if label == 'Hadron':
            TF = [True if abs(pdg) != 11 else False for pdg in Data["PDG"]]
            Dict1 = {}
            for keyD, valueD in Data.items():
                Dict1[keyD] = np.array(list(compress(valueD, TF)))
            np.save(path/label ,np.array(list(Dict1.values())))
        else:
            TF = [True if abs(pdg) == value else False for pdg in Data["PDG"]]
            Dict1 = {}
            for keyD, valueD in Data.items():
                Dict1[keyD] = np.array(list(compress(valueD, TF)))
            np.save(path/label ,np.array(list(Dict1.values())))
        
    np.save(path/"PID_BIB",np.array(list(Data.keys())))


CreatePapa()
#CratePDI()