many plots

Tracker/TrackerRecAlgs/PairVertex/scripts/NtupleReader.py

@@ -2,11 +2,12 @@
 from PlotHelper import *
 from Utilities import *
 from glob import glob
+from copy import deepcopy
 
 #CUTS
 #attr: [>,<,==,!=] if none --> skip
 REALCUT={
-    'longTracks':[1,None,None,None],}
+    'longTracks':[None,None,2,None],}
 #     'Track_nLayers':[6,None,None,None],
 #     'Track_nDoF':[0,None,None,None],
 #     'Track_Chi2':[0,None,None,None],
@@ -14,50 +15,225 @@ REALCUT={
 #     'Preshower1_charge':[0,None,None,None]   
 # }
 
+{'longTracks':[None,None,2,None]}
+
+BASELINE={
+        'VetoNu0_charge':[None,40,None,None],
+        'VetoNu1_charge':[None,40,None,None],
+        'VetoSt10_charge':[None,40,None,None],
+        'VetoSt20_charge':[None,40,None,None],
+        'VetoSt21_charge':[None,40,None,None],
+        'Preshower0_charge':[2.5,None,None,None],
+        'Preshower1_charge':[2.5,None,None,None],
+        'Calo_total_E_EM':[0,None,None,None]      
+}
+
+BASELINETEST={
+        'VetoNu0_charge':[None,40,None,None],
+        'VetoNu1_charge':[None,40,None,None],
+        'VetoSt10_charge':[None,40,None,None],
+        'VetoSt20_charge':[None,40,None,None],
+        'VetoSt21_charge':[None,40,None,None],
+        'Preshower0_charge':[2.5,None,None,None],
+        'Preshower1_charge':[2.5,None,None,None],
+        'Calo_total_E_EM':[0,None,None,None]  
+}
+
+BASELINEL=[]
+
+VETOCUTS={
+        'VetoNu0_charge':[None,None,0,None],
+        'VetoNu1_charge':[None,None,0,None],
+        'VetoSt10_charge':[None,None,0,None],
+        'VetoSt11_charge':[None,None,0,None],
+        'VetoSt20_charge':[None,None,0,None]}
+
+WAVEFORMD={
+    'Preshower0_charge':[0,None,None,None],
+    'Preshower1_charge':[0,None,None,None] }
+
 #LAMBDAS AND CONDITIONALS
 
 REALLAMBDAS=[
     lambda event: event.Timing0_charge+event.Timing1_charge+event.Timing2_charge+event.Timing3_charge>0,
-    lambda event: event.Calo0_charge+event.Calo1_charge+event.Calo2_charge+event.Calo3_charge>20000]
+    lambda event: event.Calo0_charge+event.Calo1_charge+event.Calo2_charge+event.Calo3_charge>0]
 #     lambda event: AnyCaloThreshold,
 #     lambda event: all([event.Track_nLayers[i]>6 for i in event.longTracks])
 # ]
 
-def GeneralPlotter(arrDict,dcuts=[],lcuts=[],showplot=True,savefig=False):
+WAVEFORMLAMBDAS=[
+    lambda event: event.Timing0_charge+event.Timing1_charge+event.Timing2_charge+event.Timing3_charge>0,
+    lambda event: event.Calo0_charge+event.Calo1_charge+event.Calo2_charge+event.Calo3_charge>0]
+
+def BoolTimingBaseline(event):
+    th=70
+    return (event.Timing0_charge+event.Timing1_charge+event.Timing2_charge+event.Timing3_charge>th)
+
+def BoolGoodTracks(event):
+    gt_i=[]
+    for i in range(event.longTracks):
+        tests=[]
+        if event.Track_nDoF[i]==0: continue
+        else: tests+=[event.Track_Chi2[i]/event.Track_nDoF[i]<25]
+        tests+=[event.Track_nLayers[i]>=7]
+        R=95
+        tests+=[sqrt(event.Track_x0[i]**2+event.Track_y0[i]**2)<R and sqrt(event.Track_x1[i]**2+event.Track_y1[i]**2)<R and  sqrt(event.Track_X_atVetoNu[i]**2+event.Track_Y_atVetoNu[i]**2)<R]
+        if all(tests): gt_i+=[i]
+    
+    return len(gt_i)>=1
+
+    #iterate over all tracks
+    #count tracks that are good
+    #nlayers, chi2/dof, r<100mm
+
+def BoolLowCollAreDaughters(event):
+    if event.longTracks<2: return False
+
+    angleIndexSorted=[i for i in range(event.longTracks)]
+    theta=[] #in order of track number
+    for tr in range(event.longTracks):
+        theta+=[arctan(event.Track_x0[tr]/event.Track_z0[tr])]
+
+    angleIndexSorted.sort(key=lambda elem: theta[elem]) #lowest to highest px/pz angle
+    #print(f"IndeciesSorted:{angleIndexSorted} \nTheta:{theta} \nEventBarcode: {event.t_barcode}\n")
+
+    diff_i=0
+    diff_v=theta[angleIndexSorted[0]]-theta[angleIndexSorted[1]] #if there are only 2 tracks, they make the smallest diff
+    
+    for i in range(event.longTracks-1):
+        dv=theta[angleIndexSorted[i]]-theta[angleIndexSorted[i+1]]
+        if dv<diff_v:
+            diff_v=dv
+            diff_i=i
+
+    i_lowcol1=angleIndexSorted[diff_i] #lowest diff
+    i_lowcol2=angleIndexSorted[diff_i+1]
+    test= event.t_barcode[i_lowcol1] in [2,3] and event.t_barcode[i_lowcol2] in [2,3] and event.t_barcode[i_lowcol1]!=event.t_barcode[i_lowcol2]
+    #if not test: print(f"IndeciesSorted:{angleIndexSorted} \nTheta:{theta} \nEventBarcode: {event.t_barcode}\n")
+    return test
+
+    #collinearity of all tracks vs collin specifically of matched daughters
+    #for events with tracks>2: are the matched daughters the most collin?
+
+def GetThetaPxPz(event):
+    theta=[] #in order of track number
+    for tr in range(event.longTracks):
+        theta+=[1000*arctan(event.Track_px0[tr]/event.Track_pz0[tr])]
+    return theta
+
+def GetThetaPyPz(event):
+    theta=[] #in order of track number
+    for tr in range(event.longTracks):
+        theta+=[1000*arctan(event.Track_py0[tr]/event.Track_pz0[tr])]
+    return theta
+
+def BoolLayersHit(event):
+    return all([event.tracksNLayersHit[i]>=7 for i in range(len(event.Track_nLayers))])
+
+def BoolBothDaughtersFiducial(event):
+    dp=[0,0]
+    for tp in range(len(event.truthParticleBarcode)):
+        #print(event.truthParticleBarcode)
+        if event.truthParticleBarcode[tp]==2: dp[0]=event.truthParticleIsFiducial[tp]
+        if event.truthParticleBarcode[tp]==3: dp[1]=event.truthParticleIsFiducial[tp]
+        if event.truthParticleBarcode[tp]>3: break
+    return all(dp)
+
+def BoolDaughtersMatched(event):
+    #both daughter particles are matched to tracks
+    if event.longTracks>=2:
+        return 2 in event.t_barcode and 3 in event.t_barcode
+    else: return False
+
+def IntDaughtersMatched(event):
+    #both daughter particles are matched to tracks
+    dm=0
+    if 2 in event.t_barcode: dm+=1
+    if 3 in event.t_barcode: dm+=1
+        
+    #if len(event.t_barcode)>2: print(event.t_barcode)
+    
+    return dm
+
+def BoolDM0(event):
+    return IntDaughtersMatched(event)>0
+
+def BoolDM1(event):
+    return IntDaughtersMatched(event)==1
+
+def BoolDM2(event):
+    return IntDaughtersMatched(event)==2
+
+def IntBothDaughtersFiducial(event):
+    dp=0
+    for tp in range(len(event.truthParticleBarcode)):
+        #print(event.truthParticleBarcode)
+        if event.truthParticleBarcode[tp] in [2,3]: dp+=1
+        if event.truthParticleBarcode[tp]>3: break
+    return dp
+
+
+def GeneralPlotter(t,arrDict,dcuts=[],lcuts=[],showplot=True,savefig=False,titleextra='',label=None):
     nans=0
+    for name in arrDict.keys():
+        ph=''
+        if 'lkwargs' in list(arrDict[name].keys()): ph=arrDict[name]['lkwargs']
+        arrDict[name]=deepcopy(MainArrDict[name])
+        if 'lkwargs' in list(arrDict[name].keys()): arrDict[name]['lkwargs']=ph
+        print(arrDict[name]['array'])
+
+    x=0
     ecount=0
     for event in t:
-        
-        if ApplyAllCuts(event,dcuts,lcuts):
+        ecount+=1
+        if ApplyAllCuts(event,dcuts,lcuts): #Check that event passes given cuts before doing anything with it
             for name in arrDict.keys():
+                #if x%10**5==0: print(x)
+                #Get value that will go in arrray for plotting
                 value=nan
                 if 'lkwargs' in list(arrDict[name].keys()): 
+                    #print(arrDict[name]['lkwargs'])
                     value = arrDict[name]['lambda'](event,**arrDict[name]['lkwargs'])
                 else: value = arrDict[name]['lambda'](event)
+
+                #Check that all values to be added to the plot-array are valid
                 checkReal=True
                 if type(value)==list: checkReal = all([v != inf and v !=nan and v !=-1*inf for v in value])
                 else: checkReal = value != inf and value !=nan and value !=-1*inf
-
-                if checkReal and type(value)==list:
-                    if arrDict[name]['plotterfunc'] in [configBasicHist]: arrDict[name]['array']+=value
-                    else:
+                
+                #If the value is valid, add it to the plot array.
+                if checkReal and type(value)==list: 
+                    if arrDict[name]['plotterfunc'] in [configBasicHist]: arrDict[name]['array']+=value #when plotting on a subset of items like tracks instead of events
+                    else: #MultiHist
                         for iv in range(len(value)):
                             arrDict[name]['array'][iv]+=[value[iv]]
-                elif checkReal: arrDict[name]['array']+=[value]
+                elif checkReal: arrDict[name]['array']+=[value] #Basic histogram
+                elif 'SepHist' in list(arrDict[name].keys()): #are subarrays allowed to have different lengths?
+                    checkList=[v != inf and v !=nan and v !=-1*inf for v in value] #if yes, check each value returned for nans (may be nan if cut; actual nancount doesnt apply here)
+                    for ich in len(checkList):
+                        if checkList[ich]: arrDict[name]['array'][ich]+=[value[ich]]
+
                 else: nans+=1
-        if ecount%100000==0:print(ecount,' events processed')
-        ecount+=1
 
     print(f"# of NaNs:{nans}")
     for name in arrDict.keys():
-        isPrintout= arrDict[name]['plotterfunc'] in [PrintBoolPercent]
+        isPrintout= arrDict[name]['plotterfunc'] in [PrintBoolPercent] #Is it text-only output?
         print(name)
-        print('LEN:',len(arrDict[name]['array']))
-        arrDict[name]['plotterfunc']( arrDict[name]['array'], **arrDict[name]['plotkwargs'])
+        #print(arrDict[name]['descr'])
+        print('LEN:',len(arrDict[name]['array'])) #how many events remain after the cut?
+        if label:
+            arrDict[name]['plotkwargs'].update({'label':label})
+            arrDict[name]['plotkwargs']['pckwargs'].update({'legend':True})
+
+        arrDict[name]['plotterfunc']( arrDict[name]['array'], **arrDict[name]['plotkwargs']) #plot using the given function
         if not isPrintout:
-            plt.title(user_input+name)
+            plt.title(user_input+name+titleextra)
             if savefig: plt.savefig(user_input+name)
             if showplot: plt.show()
+
+        print()
+    print("EVENTS: ",ecount)
+    
 COUNT=0
 def ReadoutTest(event):
     value = event.longTracks
@@ -68,19 +244,89 @@ def ReadoutTest(event):
     if value>1: print(value)
     return value
 
+def OrderChecker(arrDict):
+    x=0
+    for name in arrDict.keys():
+        arrDict[name]=deepcopy(MainArrDict[name])
+        print(arrDict[name]['array'])
+
+    for event in t:
+        x+=1
+        if x%10**5==0: print(x)
+
+#TRACK EFFICIENCY STUFF
+def GetTrackEfficiency(t,dcutN=[],dcutD=[],lcutN=[],lcutD=[],plim=0):
+    narray=0.0
+    darray=0.0
+    for event in t:
+        if ApplyAllCuts(event,dcutN,lcutN) and BoolTrackPCut(event,plim): narray+=1
+            #each track
+            # for tr in range(event.longTracks):
+            #     if event.Track_p0[tr]>plim:
+            #         narray+=1.0
+        if ApplyAllCuts(event,dcutD,lcutD) and BoolBothDaughtersFiducial(event): darray+=1 
+            # for tr in range(event.longTracks):
+            #     darray+=IntBothDaughtersFiducial(event)
+    if darray==0:return 0
+    return (narray/darray)*100
+
+def BoolTrackPCut(event,plim=0):
+    #require all tracks to pass the cut
+    for tr in range(event.longTracks):
+        if event.Track_p0[tr]<plim: return False
+    return True
+
+def GetTrackEffCoords(t):
+    plims = [0,10,20,30,40,50]
+    ycoords=[] #list of lists
+    calocuts=[0,100,200,300,400,500]
+    
+    for cl in calocuts:
+        y=[]
+        for pl in plims:
+            y+=[GetTrackEfficiency(t,dcutN=[BASELINE],lcutN=[lambda event: event.Calo_total_E_EM>cl*(10**3),BoolGoodTracks,BoolTimingBaseline],lcutD=[],plim=pl*(10**3))]
+        ycoords+=[y]
+    labels=[f'CaloEM_Tot > {c}GeV' for c in calocuts]
+    return plims,ycoords,labels
+
+def TrackEfficiencyShow(t,title=''):
+    x,ys,labels=GetTrackEffCoords(t)
+    for i in range(len(ys)):
+        PlotTrackEfficiency(x,ys[i],labels[i])
+    
+    PlotConfigs(xscale='linear',title=title,ylim=(0,100),show=True,xlabel='Upstream Pmag Cut (GeV)',ylabel='Efficiency (%)',legend=True)
+
+#TRACK SEPARATION VS MOMENTUM STUDY
+def ShowTrackSepVsRecoMom(t):
+    x1,x2=([],[]) #track sep
+    y1,y2=([],[]) #reco momentum
+    for event in t:
+        if ApplyAllCuts(event,[BASELINE],[BoolTimingBaseline,BoolGoodTracks]):
+            if IntDaughtersMatched(event)==1:
+                x1+=[sqrt((event.truthd0_x[0]-event.truthd1_x[0])**2+(event.truthd0_y[0]-event.truthd1_y[0])**2)]
+                y1+=[max(event.Track_p0)]
+            if IntDaughtersMatched(event)==2:
+                x2+=[sqrt((event.truthd0_x[0]-event.truthd1_x[0])**2+(event.truthd0_y[0]-event.truthd1_y[0])**2)]
+                y2+=[max(event.Track_p0)]
+    
+    PlotTrackEfficiency(x1,y1,'Matched Tracks==1')
+    PlotTrackEfficiency(x2,y2,'Matched Tracks==2')
+    PlotConfigs(xscale='linear',title='Track Sep vs Max Reco Momentum',show=True,xlabel='Track Separation (mm)',ylabel='Max Upstream Pmag (MeV)',legend=True)
+
+
 #
 MainArrDict={
     'CutTest':{
         'array':[],
         'lambda': ApplyAllCuts,
-        'lkwargs':{'dcuts':[REALCUT],'lcuts':REALLAMBDAS}, 
+        'lkwargs':{'dcuts':[BASELINE],'lcuts':[BoolTimingBaseline, BoolGoodTracks,BoolDM0]}, 
         'plotterfunc': PrintBoolPercent,
         'plotkwargs' : {},
         'descr':'Check how much a cut cuts.'
     },
     'BasicHist':{
         'array':[],
-        'lambda': ReadoutTest,
+        'lambda': IntDaughtersMatched,
         'plotterfunc': configBasicHist,
         'plotkwargs' : {},
         'descr':'Single variable tests.'
@@ -98,30 +344,144 @@ MainArrDict={
         'plotterfunc': configMultiHist,
         'plotkwargs' : {'label':['Lowest','Highest'],'pckwargs':{'xscale':'symlog','xlabel':"Momentum (MeV)",'legend':True}},
         'descr':'Look at 2 highest momentum tracks downstream'
-    },}
+    },
+    'TrackCount':{
+        'array':[],
+        'lambda': lambda event: event.longTracks,
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[0,25]},'pckwargs':{}}
+    },
+    'CaloTotal':{
+        'array':[],
+        'lambda': lambda event: event.Calo_total_E_EM, #event.Calo0_charge+event.Calo1_charge+event.Calo2_charge+event.Calo3_charge,
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'pckwargs':{}}
+    },
+    'Momentum':{
+        'array':[],
+        'lambda': lambda event: [p for p in event.Track_p0],
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[10**3,10**10]},'pckwargs':{'xscale':'symlog','xlabel':"Momentum (MeV)",'ylabel':"#Tracks",'legend':True},'rrange':[10**3,10**12]},#
+        'descr':'Look at momentum tracks upstream'
+    },
+    'TruthMomentum':{
+        'array':[],
+        'lambda': lambda event: [event.truthd0_P[0],event.truthd1_P[0]],
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[10**3,10**10]},'pckwargs':{'xscale':'symlog','xlabel':"Momentum (MeV)",'ylabel':"#Tracks",'legend':True},'rrange':[10**3,10**12]},#
+        'descr':'Look at momentum tracks upstream'
+    },
+    'MaxRecoMomentumUp':{
+        'array':[],
+        'lambda': lambda event: max(event.Track_p0),
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'bins':50,'kwargs':{'range':[10**3,10**11]},'pckwargs':{'xscale':'log','xlabel':"Momentum (MeV)",'legend':True}},
+        'descr':'Look at highest momentum tracks upstream'
+    },
+    'ThPxPz':{
+        'array':[],
+        'lambda':GetThetaPxPz,
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[-5,5]},'pckwargs':{'xlabel':"Theta_PxPz (mrad)",'ylabel':"#Tracks"}},#
+        'descr':'Look at collinearity'
+    },
+    'ThPyPz':{
+        'array':[],
+        'lambda':GetThetaPyPz,
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[-5,5]},'pckwargs':{'xlabel':"Theta_PyPz (mrad)",'ylabel':"#Tracks"}},#
+        'descr':'Look at collinearity'
+    },
+    'CollPxPz':{
+        'array':[],
+        'lambda': BoolLowCollAreDaughters,
+        'plotterfunc': PrintBoolPercent,
+        'plotkwargs' : {},
+        'descr':'Check how much a cut cuts.'
+    },
+    'RecoSep2trk':{
+        'array':[],
+        'lambda': lambda event: sqrt((event.Track_X_atTrig[0]-event.Track_X_atTrig[1])**2+(event.Track_Y_atTrig[0]-event.Track_Y_atTrig[1])**2),
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[0,20]},'pckwargs':{}},
+        'descr':'Single variable tests.'
+    },
+    'TruthSep2trk':{
+        'array':[],
+        'lambda': lambda event: sqrt((event.t_st1_x[0]-event.t_st1_x[1])**2+(event.t_st1_y[0]-event.t_st1_y[1])**2),
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[0,20]},'pckwargs':{}},
+        'descr':'Single variable tests.'
+    },
+    }
 
 TestArrDict=dict()
 testname='HighestMomentumDown' #Any key from MainArrDict
 TestArrDict[testname]=MainArrDict[testname]
 
-def getRunFiles(rnum,dataDir="/eos/experiment/faser/phys/2022/r0011/"):
-    dataDir=f"{dataDir}{rnum}/"
+def getRunFiles(rnum,dataDir="/eos/experiment/faser/phys/2022/r0013/"):
+    dataDir=f"/eos/experiment/faser/sim/mc22/foresee/{rnum}/phy/r0013/" #f"{dataDir}{rnum}/"
     return sorted(glob(f"{dataDir}/Faser*"))
 
-def getRunNums(dataDir="/eos/experiment/faser/phys/2022/r0011/"):
+def getRunNums(dataDir="/eos/experiment/faser/phys/2022/r0013/"):
     return [rn.split(dataDir)[1] for rn in sorted(glob(f"{dataDir}*"))]
 
+def MakePlotDict(entry):
+    TestArrDict=dict()
+    if type(entry)==list:
+        for e in entry:
+            testname=e #Any key from MainArrDict
+            TestArrDict[testname]=deepcopy(MainArrDict[testname])
+    else:
+        testname=entry #Any key from MainArrDict
+        TestArrDict[testname]=deepcopy(MainArrDict[testname])
+    return TestArrDict
 
 #EXECUTE
 if __name__ == "__main__":
     #runnum = '008301'
     t = TChain("nt") 
     nfiles = 0
-    for runnum in getRunNums():
+    fgrab=''
+    for runnum in [f'1100{ri}' for ri in range(22,56)]:#,52,55,35,42,49,36,43,50,37,44,31,45]]:#getRunNums():
+        #if int(runnum[-2:]) in [28,33,38,51]: continue
+        if int(runnum[-2:]) != 38: continue
+        #if int(runnum[-2:]) not in [48,49]: continue#43,44,45,35,36,37,31,32 52,48,49,50,41,42
+        
         files = getRunFiles(runnum)
+        user_input=f'Ntuple_Run{runnum}'
+        
         for f in files:
             print(f)
+            fgrab=f
             nfiles += t.Add(f) #t.Add("/eos/experiment/faser/phys/2022/r0011/"+user_input+".root") 
-            
+        
+    rnlab = fgrab.split("FaserMC-MC22_")[1].split(f'-1100')[0]
+    print(rnlab)
+    #rnlab = "All 5k-event Models"
+
+    #TrackEfficiencyShow(t,title=f'{rnlab} TrackEfficiency (Good Trks>=1)')
+
+    # Momentum plot
+    # for li in [0,5,10,50,100]:
+    #     mpd=MakePlotDict('Momentum')
+    #     mpd['Momentum']['lkwargs']['lim']=li*10**3
+    #     #print(mpd)
+    #     try:
+    #         emcut=100
+    #         GeneralPlotter(t,mpd,dcuts=[],lcuts=[lambda event: event.Calo_total_E_EM>emcut*(10**3)],showplot=False,label=f'Pmag>{li}GeV Cut ')
+    #     except ValueError:
+    #         print("Empty after cut: ",li,"Gev")
+    #         continue
+    # plt.title(f"FaserMC Aee {rnlab} Pmag CaloCut={emcut} GeV")
+    # plt.show()
     
-    GeneralPlotter(TestArrDict,dcuts=[{'longTracks':[None,None,2,None],}],lcuts=[],savefig=True)
+    # GeneralPlotter(t,MakePlotDict('CutTest'),dcuts=[],lcuts=[BoolBothDaughtersFiducial])
\ No newline at end of file
+    #GeneralPlotter(t,MakePlotDict('BasicHist'),dcuts=[BASELINE],lcuts=[BoolTimingBaseline, BoolGoodTracks,BoolBothDaughtersFiducial])
+    GeneralPlotter(t,MakePlotDict('MaxRecoMomentumUp'),dcuts=[BASELINE,{'longTracks':[None,None,1,None],}],lcuts=[BoolTimingBaseline,BoolGoodTracks])
+    #GeneralPlotter(t,MakePlotDict('MaxRecoMomentumUp'),dcuts=[BASELINE],lcuts=[BoolTimingBaseline,BoolGoodTracks,BoolDM2],label='Matched Tracks==2')
+    #GeneralPlotter(t,MakePlotDict('TruthMomentum'),dcuts=[BASELINETEST,REALCUT],lcuts=[BoolDaughtersMatched,BoolBothDaughtersFiducial],label='Truth')
+    #GeneralPlotter(t,MakePlotDict('CollPxPz'),dcuts=[{'longTracks':[None,None,1,None]}],lcuts=[])
+
+    #ShowTrackSepVsRecoMom(t)
+    
\ No newline at end of file

Tracker/TrackerRecAlgs/PairVertex/scripts/NtupleReaderROOT.py

+#!/usr/bin/env python
+import ROOT as R
+from glob import glob
+from copy import deepcopy
+
+#CUTS
+#attr: [>,<,==,!=] if none --> skip
+REALCUT={
+    'longTracks':[None,None,1,None],}
+#     'Track_nLayers':[6,None,None,None],
+#     'Track_nDoF':[0,None,None,None],
+#     'Track_Chi2':[0,None,None,None],
+#     'Preshower0_charge':[0,None,None,None],
+#     'Preshower1_charge':[0,None,None,None]   
+# }
+
+VETOCUTS={
+        'VetoNu0_charge':[None,None,0,None],
+        'VetoNu1_charge':[None,None,0,None],
+        'VetoSt10_charge':[None,None,0,None],
+        'VetoSt11_charge':[None,None,0,None],
+        'VetoSt20_charge':[None,None,0,None]}
+
+WAVEFORMD={
+    'Preshower0_charge':[0,None,None,None],
+    'Preshower1_charge':[0,None,None,None] }
+
+#LAMBDAS AND CONDITIONALS
+
+REALLAMBDAS=[
+    lambda event: event.Timing0_charge+event.Timing1_charge+event.Timing2_charge+event.Timing3_charge>0,
+    lambda event: event.Calo0_charge+event.Calo1_charge+event.Calo2_charge+event.Calo3_charge>0]
+#     lambda event: AnyCaloThreshold,
+#     lambda event: all([event.Track_nLayers[i]>6 for i in event.longTracks])
+# ]
+
+WAVEFORMLAMBDAS=[
+    lambda event: event.Timing0_charge+event.Timing1_charge+event.Timing2_charge+event.Timing3_charge>0,
+    lambda event: event.Calo0_charge+event.Calo1_charge+event.Calo2_charge+event.Calo3_charge>0]
+
+def RootFilterFromDict(dcuts): #accepts dictionary of cuts, returns a string
+    output=''
+    symbols=['>=','<=','==','!=']
+    for k in dcuts.keys():
+        for ib in range(0,4):
+            if dcuts[k][ib]!=None: 
+                o1=''
+                o1+=f'{k}{symbols[ib]}{dcuts[k][ib]}'
+                if output!='': o1=' && '+o1
+                output+=o1
+    
+    return output
+
+COUNT=0
+def ReadoutTest(event):
+    value = event.longTracks
+    global COUNT
+    COUNT+=1
+    if COUNT<1000: print(value)
+    if COUNT==1000: print("Reached 1000 events. Stopping printout except when value>1")
+    if value>1: print(value)
+    return value
+
+def OrderChecker(arrDict):
+    x=0
+    for name in arrDict.keys():
+        arrDict[name]=deepcopy(MainArrDict[name])
+        print(arrDict[name]['array'])
+
+    for event in t:
+        x+=1
+        if x%10**5==0: print(x)
+
+def getRunFiles(rnum,dataDir="/eos/experiment/faser/phys/2022/r0013/"):
+    dataDir=f"{dataDir}{rnum}/"
+    return [sorted(glob(f"{dataDir}/Faser*"))[0]]
+
+def getRunNums(dataDir="/eos/experiment/faser/phys/2022/r0013/"):
+    return [rn.split(dataDir)[1] for rn in sorted(glob(f"{dataDir}*"))]
+
+#EXECUTE
+if __name__ == "__main__":
+    runnum = 110031#'009148'
+    # Utilise multiple cores
+    R.EnableImplicitMT()
+    R.gROOT.SetBatch(True)
+    
+    # Create Dataframe
+    #df = R.RDataFrame("nt", f"/eos/experiment/faser/phys/2022/r0013/{runnum}/Faser-Physics-{runnum}-*-PHYS.root") #DATA
+    df = R.RDataFrame("nt", f"/eos/experiment/faser/sim/mc22/foresee/{runnum}/phy/r0013/FaserMC-MC22_FS_Aee_*-PHYS.root") #SIM
+
+    filterstr=RootFilterFromDict(WAVEFORMD)
+    #df.Filter(filterstr)
+    #df.Draw("Track_p0")
+    histo = df.Histo1D(("Track_p0", "; Pmag_up; NTracks", 100, 0, 10**10), "Track_p0")
+    c = R.TCanvas()
+    histo.Draw()
+    
+    input()
\ No newline at end of file

Tracker/TrackerRecAlgs/PairVertex/scripts/PairVertexReader.py

@@ -327,6 +327,12 @@ MainArrDict={
         'descr':'How many clusters have energy deposits from both daughter particles?',
         'SepHist':'arrays allowed to have different lengths'
     },
+    'TrackCount':{
+        'array':[],
+        'lambda': lambda event: event.TrackCount,
+        'plotterfunc': configBasicHist,
+        'plotkwargs' : {'kwargs':{'range':[0,25]},'pckwargs':{}}
+    },
     
     }
 
@@ -393,18 +399,30 @@ def GeneralPlotter(arrDict,dcuts=[],lcuts=[],showplot=True,savefig=False,titleex
 
         print()
 
+def OrderChecker(arrDict):
+    nans=0
+    for name in arrDict.keys():
+        arrDict[name]=deepcopy(MainArrDict[name])
+        print(arrDict[name]['array'])
+
+    for event in t:
+        x=1
 
 #EXECUTE
 if __name__ == "__main__":
-    DP_fnames=['Ntuple_FaserMC_MC22_FS_Aee_100MeV_1Em5']#  'Ntuple_Aee_10MeV_1Em4_' 'Ntuple_Aee_100MeV_1Em5_','Ntuple_Aee_10MeV_1Em5_', ,'Ntuple_Amm_316MeV_2Em6_' .root
+    DP_fnames=['Ntuple_Aee_100MeV_1Em5_']#  'Ntuple_Aee_10MeV_1Em4_' '','Ntuple_Aee_10MeV_1Em5_', ,'Ntuple_Amm_316MeV_2Em6_' .root
+    #r0013  Ntuple_FaserMC_MC22_FS_Aee_100MeV_1Em5
     for f in DP_fnames:
         user_input=f
         t = TChain("events") 
         nfiles = 0
-        nfiles += t.Add("/eos/home-s/sshively/CALYPZONE/run/"+user_input+".root") 
+        nfiles += t.Add("/eos/home-s/sshively/CALYPZONE/run/BackUp2/"+user_input+".root") 
         print(f)
 
-        GeneralPlotter(MakePlotDict('RecoMaxMomentum'),dcuts=[{}],lcuts=[])
+        OrderChecker(MakePlotDict('TrackCount'))
+        # GeneralPlotter(MakePlotDict('TrackCount'),dcuts=[{}],lcuts=[],label='No Cuts',showplot=False)
+        # GeneralPlotter(MakePlotDict('TrackCount'),dcuts=[WaveformCuts],lcuts=WaveformCuts2,label='WF Cut')
+        # GeneralPlotter(MakePlotDict('TrackCount'),dcuts=[{}],lcuts=[lambda e: e.Calo0_integral+e.Calo1_integral+e.Calo2_integral+e.Calo3_integral>10**5],label='>1 TeV cut',)
 
     
 

Tracker/TrackerRecAlgs/PairVertex/scripts/PlotHelper.py

@@ -6,6 +6,8 @@ from ROOT import TEfficiency, TH1F
 def PlotConfigs(xscale='linear',yscale='linear',xlim=None,ylim=None,title=None,show=False,xticks=None,xlabel=None,ylabel='#Events',legend=False,legloc='best',save=False):
     plt.xscale(xscale)
     plt.yscale(yscale)
+    if xlim: plt.xlim(xlim)
+    if ylim: plt.ylim(ylim)
     if xticks: plt.xticks(xticks)
     if xlabel: plt.xlabel(xlabel)
     if ylabel: plt.ylabel(ylabel)
@@ -28,8 +30,9 @@ def configHitMap(arrays,bins=(50,50),pckwargs={'xlabel':"X (mm)",'ylabel':"Y (mm
 def configBasicHist(arr,bins=25,label=None,pckwargs={'xscale':'linear'},kwargs={},rrange=None):
     #1D Histogram
     if 'xscale' in pckwargs.keys() and 'log' in pckwargs['xscale']: 
-        if rrange:CorrectBinsForLog(arr,bins,pckwargs['xscale'],rrange=rrange)
-        else: bins=CorrectBinsForLog(arr,bins,pckwargs['xscale'])
+        #if rrange:CorrectBinsForLog(arr,bins,pckwargs['xscale'],rrange=rrange)
+        #else: 
+            bins=CorrectBinsForLog(arr,bins,pckwargs['xscale'])
     n, binsOut, patches = plt.hist(arr,bins=bins,label=label,align="left",histtype=u'step',**kwargs)
     PlotConfigs(**pckwargs)
 
@@ -92,17 +95,23 @@ def PlotRootRatio(arrlist,bins=15,hmin=0,hmax=10000,title=None):
 #Utilities
 def CorrectBinsForLog(arr,bins,xscale,rrange=None):
     minx,maxx = (min(arr),max(arr))
-    if rrange!=None:  minx,maxx = (rrange[0],rrange[1])
+    if rrange!=None:  
+        minx,maxx = (rrange[0],rrange[1])
     if minx==maxx: 
         maxx+=1
         bins=1
-    bins = logspace(start=log10(minx), stop=log10(maxx), num=bins)
+    minx=10.0**3
+    maxx=10.0**11
+    bins2 = logspace(start=log10(minx), stop=log10(maxx), num=bins)
     if xscale=='symlog':
-        bins = [0]+list(logspace(start=log10(1), stop=log10(maxx), num=24))
-        if minx>1: bins = logspace(start=log10(minx), stop=log10(maxx), num=25)
+        bins2 = [0]+list(logspace(start=log10(1), stop=log10(maxx), num=24))
+        if minx>1: bins2 = logspace(start=log10(minx), stop=log10(maxx), num=25)
 
-    return bins
+    return bins2
 
+def PlotTrackEfficiency(x,y,label=''):
+    plt.plot(x,y,label=label)
+    plt.scatter(x,y)
 
     
 

Tracker/TrackerRecAlgs/PairVertex/scripts/PlotHelperROOT.py

+from numpy import log10,logspace
+from cmcrameri import cm
+import matplotlib.pyplot as plt
+from ROOT import TEfficiency, TH1F
+
+def PlotConfigs(xscale='linear',yscale='linear',xlim=None,ylim=None,title=None,show=False,xticks=None,xlabel=None,ylabel='#Events',legend=False,legloc='best',save=False):
+    plt.xscale(xscale)
+    plt.yscale(yscale)
+    if xticks: plt.xticks(xticks)
+    if xlabel: plt.xlabel(xlabel)
+    if ylabel: plt.ylabel(ylabel)
+    if legend: plt.legend(loc=legloc)
+    if title:plt.title(title)
+    if show: plt.show()
+    #filename = title.replace(" ", "_").replace("(", "").replace(")", "")
+    #if save: plt.savefig(filename)
+
+
+def configHitMap(arrays,bins=(50,50),pckwargs={'xlabel':"X (mm)",'ylabel':"Y (mm)"},counts=False,kwargs={}):
+    #2D Histogram with colors
+    arrayx=arrays[0]
+    arrayy=arrays[1]
+    plt.hist2d(arrayx, arrayy, bins=bins, cmap=cm.imola,**kwargs)
+    plt.colorbar()
+    if counts: plt.text(50,-185,"#EVENTS: {}".format(len(arrayx)))
+    PlotConfigs(**pckwargs)
+
+def configBasicHist(arr,bins=25,label=None,pckwargs={'xscale':'linear'},kwargs={},rrange=None):
+    #1D Histogram
+    if 'xscale' in pckwargs.keys() and 'log' in pckwargs['xscale']: 
+        if rrange:CorrectBinsForLog(arr,bins,pckwargs['xscale'],rrange=rrange)
+        else: bins=CorrectBinsForLog(arr,bins,pckwargs['xscale'])
+    n, binsOut, patches = plt.hist(arr,bins=bins,label=label,align="left",histtype=u'step',**kwargs)
+    PlotConfigs(**pckwargs)
+
+    return n,binsOut
+
+def configMultiHist(arrlist,bins=25,label=[],pckwargs={'xscale':'linear','legloc':'upper left'},kwargs={}):
+    #1D Histogram for comparison
+    megarr=[]
+    for a in arrlist: megarr+=a
+
+    if 'log' in pckwargs['xscale']: bins=CorrectBinsForLog(megarr,bins,pckwargs['xscale'])
+    for iarr in range(len(arrlist)):
+        arr=arrlist[iarr]
+        labelused=''
+        if len(arr)==0:
+            print(f"Empty array.")
+        else: 
+            if len(label)==len(arrlist): labelused=label[iarr]
+            else: 
+                print("Insufficient labels given.")
+            plt.hist(arr,bins=bins,align="left",histtype=u'step',label=labelused,**kwargs)
+        
+    PlotConfigs(**pckwargs)
+
+def config2DScatter(arrlist,pckwargs=dict(),kwargs={}):
+    x=arrlist[0]
+    y=arrlist[1]
+    # plt.ylim(-25,25)
+    # plt.xlim(-10**-5,10**-5)
+    plt.scatter(x,y,**kwargs)
+    PlotConfigs(**pckwargs)
+
+def PrintBoolPercent(arr,kwargs={}):
+    #For an array of boolians, print out the % of elements that are "true".
+    n=0.0
+    d=len(arr)
+    for x in arr:
+        if x: n+=1.0
+    
+    finalval= (n/d)*100
+    print("Calculating...")
+    print(f"{n}/{d}")
+    print('Percent:'+'{:.2f}'.format(finalval)+'%')
+
+def PlotRootRatio(arrlist,bins=15,hmin=0,hmax=10000,title=None):
+
+    narray,darray=arrlist[0],arrlist[1]
+    h_numerator=TH1F("num","num",bins,hmin,hmax)
+    h_denominator=TH1F("den","den",bins,hmin,hmax)
+    for i in narray: h_numerator.Fill(i)
+    for j in darray: h_denominator.Fill(j)
+
+    h_ratio = TEfficiency(h_numerator,h_denominator)
+    c = h_ratio.CreateGraph()
+    c.GetYaxis().SetLimits(0.0, 1.05)
+    if title: h_ratio.SetTitle(title)
+    h_ratio.Draw("ep")
+    input()
+
+#Utilities
+def CorrectBinsForLog(arr,bins,xscale,rrange=None):
+    minx,maxx = (min(arr),max(arr))
+    if rrange!=None:  minx,maxx = (rrange[0],rrange[1])
+    if minx==maxx: 
+        maxx+=1
+        bins=1
+    bins = logspace(start=log10(minx), stop=log10(maxx), num=bins)
+    if xscale=='symlog':
+        bins = [0]+list(logspace(start=log10(1), stop=log10(maxx), num=24))
+        if minx>1: bins = logspace(start=log10(minx), stop=log10(maxx), num=25)
+
+    return bins
+
+
+    
+

Tracker/TrackerRecAlgs/PairVertex/scripts/Utilities.py

@@ -7,7 +7,7 @@ from operator import lt,eq,ne,gt
 from numpy import nan, seterr, inf,arctan, array
 from numpy import sum as npsum
 import numpy as np
-from math import sqrt
+from math import sqrt,isnan
 print("Other stuff loaded.")
 
 user_input=''
@@ -31,7 +31,11 @@ def ApplyCut(event, cut):
             try:
                 if cut[attr][i]!=None: 
                     value = getattr(event,attr)
-                    bool1 += [tests[i](value,cut[attr][i]) and value != inf and value !=nan and value !=-1*inf]
+                    #print(value)
+                    if 'Veto' in attr: 
+                        if isnan(value): value=0
+                        bool1 += [tests[i](value,cut[attr][i]) and value != inf and value !=-1*inf]
+                    else: bool1 += [tests[i](value,cut[attr][i]) and value != inf and not isnan(value) and value !=-1*inf]
                     if not all(bool1): break
             except AttributeError:
                 if attr not in missingattr: