import ROOT, argparse, operator, numpy import copy parser=argparse.ArgumentParser(description='Plot stuff along a track') parser.add_argument('recofile',help='the recofile with cluster tree and digi and sample persistence',type=str) parser.add_argument('--folder',help='if the recofile is in an outher folder then mc and digi, input: /folder', type=str, default="not") parser.add_argument('--foldertosave', help='the folder to save the output', type=str, default="opt_clustering/FT_helixc/") parser.add_argument('--ev',help='which event to go to',type=int) parser.add_argument('--pdir',help='the picture print directory',type=str,default="not") parser.add_argument('--filename', help='name of the file', type=str, default="not") args=parser.parse_args() ROOT.gROOT.ProcessLine(".x rootlogon.C") ROOT.gROOT.ProcessLine('gSystem->Load("libPhysics")') ROOT.gROOT.ProcessLine('gStyle->SetPalette(1)') ROOT.gROOT.ProcessLine('gROOT->SetStyle("Plain")') ROOT.gROOT.ProcessLine('gSystem->Load("libGeom")') ROOT.gROOT.LoadMacro("stlPYROOT.h+") Infile=ROOT.TFile(args.recofile,"read") tree=Infile.Get("cbmsim") if tree==None: print "no cbmsim" Infile.Close() exit() mcfile=args.recofile.replace("reco","mc") mcfile=mcfile.replace("_recl","") digifile=args.recofile.replace("reco","raw") digifile=digifile.replace("_recl","") if args.folder != "not": #first_z=mcfile.find('/') first=mcfile.find('/') l=mcfile.find('/',first+1) last=mcfile.find('/',l+1) length=len(mcfile) part_one=mcfile[0:first] part_two=mcfile[last:length] newname=part_one+args.folder+part_two mcfile=newname #first1_z=digifile.find('/') first1=digifile.find('/') l1=digifile.find('/',first1+1) last1=digifile.find('/',l1+1) length1=len(digifile) part_one1=digifile[0:first1] part_two1=digifile[last1:length1] newname1=part_one1+args.folder+part_two1 digifile=newname1 print 'adding',mcfile,'and',digifile,'as friends' tree.AddFriend("cbmsim",mcfile) tree.AddFriend("cbmsim",digifile) tree.SetBranchStatus("*",1) tree.SetBranchStatus("TpcSample.*",1) tree.SetBranchStatus("TpcDigi.*",1) tree.SetBranchStatus("TpcSample2.*",1) tree.SetBranchStatus("TpcDigi2.*",1) tree.SetBranchStatus("TpcCluster.*",1) tree.SetBranchStatus("TpcPreCluster.*",1) tree.SetBranchStatus("TrackPostFit.*",1) #Define Histogramms----------------------------------------------------- hcomp=ROOT.TH1D("hcomp", "completeness Cluser", 101, 0, 101) htrkspl=ROOT.TH1D("htrakspl", "Tracksplitting", 5, 0, 5) hcont=ROOT.TH1D("hcont", "Contamination", 101, 0, 101) #Class------------------------------------------------------------------ class Tracks (object): def __init__(self,prefitcl,fitcl,pdg,numberoftracks): self.PreFitCl = prefitcl self.PreFitComp=[] self.FitCl=fitcl self.FitComp=[] self.Pdg=pdg self.NumberOfTracks=numberoftracks self.Split = 0 self.Splitkoord={} self.NumberOfCl=[] self.NumberOfClPre=[] self.Completeness=[] self.Contamination=[] def GetClCountPerEvent(self): self.ClCountPerEvent=len(self.PreFitCl) return self.ClCountPerEvent def GetPreFitComp(self): if self.NumberOfTracks > 1: counter=0 breaker=False PreFitCompTMP1=[] PreFitClTMP=copy.deepcopy(self.PreFitCl) while breaker == False: PreFitCompTMP2=[] for i in self.PreFitCl: if i == counter: PreFitCompTMP2.append(i) del PreFitClTMP[0] PreFitCompTMP1.append(PreFitCompTMP2) counter +=1 if len(PreFitClTMP)==0: breaker=True for i in range(len(PreFitCompTMP1)): self.PreFitComp.append(PreFitCompTMP1[i]) elif self.NumberOfTracks==1: self.PreFitComp.append(self.PreFitCl) else: self.PreFitComp=[] return self.PreFitComp def GetFitComp(self): if self.NumberOfTracks != 0: #print self.Pdg for i in range(self.NumberOfTracks): self.FitComp.append(FitCl[i]) else: self.FitComp=[] return self.FitComp def GetMap(self): counter=0 for j in self.FitComp: CounterR=0 CounterE=0 for k in j: if k == j[0]: CounterR+=1 else: CounterE+=1 try: self.Splitkoord[j[0]].append(counter) except KeyError: self.Splitkoord[j[0]] = [] self.Splitkoord[j[0]].append(counter) counter+=1 return self.Splitkoord def GetTracksplitting(self): mont=0 if self.NumberOfTracks>0: for i in range(self.NumberOfTracks-1): if self.FitComp[i][0]==self.FitComp[i+1][0]: mont+=1 if mont != 0: self.Split=self.NumberOfTracks/mont elif self.NumberOfTracks == 0: self.Split=-1 else: self.Split=1 return self.Split def GetCompleteness(self): if self.NumberOfTracks!=0: for i in range(len(self.PreFitComp)): if len(self.PreFitComp[i]) != 0: self.NumberOfClPre.append(len(self.PreFitComp[i])) for i in range(len(self.Splitkoord)): l=0 try: for j in range(len(self.Splitkoord[i])): k=self.Splitkoord[i][j] l+=len(self.FitComp[k]) except KeyError: continue self.NumberOfCl.append(l) for i in range(len(self.NumberOfCl)): self.Completeness.append((self.NumberOfCl[i]*100)/(self.NumberOfClPre[i])) else: self.Completeness=[] return self.Completeness def GetContamination(self): cont=[] id_count=0 for i in range(len(self.Splitkoord)): #id_count=0 cont_counter=0 try: for j in range(len(self.Splitkoord[i])): k=self.Splitkoord[i][j] for m in range(len(self.FitComp[k])): if id_count == self.FitComp[k][m]: cont_counter+=1 cont.append(cont_counter) id_count+=1 except KeyError: continue #id_count+=1 for i in range(len(self.NumberOfCl)): self.Contamination.append((cont[i]*100)/(self.NumberOfCl[i])) return self.Contamination #GetArrays-------------------------------------------------------------- evnum = -1 event = -1 print "events:", tree.GetEntries() for e in tree: evnum += 1 event +=1 trackindex = -1 number_of_tracks=e.TrackPostFit.GetEntries() #print "tracks:",number_of_tracks print "EVENT:", event #Pdg pdg_list=[] for mctrk in e.MCTrack: pdg = mctrk.GetPdgCode() pdg_list.append(pdg) #FitCluster FitCl=[] for trk in e.TpcPreTrackPreFit: FitCl_tmp=[] cand = trk.getCand() hitIds = cand.getHitIDs(14) for ihit in hitIds: hit=e.TpcPreSPHit.At(ihit) clid=hit.getClusterID() cl=e.TpcPreCluster.At(clid) McId=cl.mcId() Size=cl.nMcIds() McDomId=McId.DominantID() trackID=McDomId.mctrackID() FitCl_tmp.append(trackID) FitCl.append(FitCl_tmp) #PreFitCluster PreFitCl=[] for cl in e.TpcPreCluster: McPreId=cl.mcId() PreSize=cl.nMcIds() McPreDomId=McPreId.DominantID() trackPreID=McPreDomId.mctrackID() PreFitCl.append(trackPreID) PreFitCl.sort() FitCl.sort() for i in range(len(FitCl)): FitCl[i].sort() #Sorting---------------------------------------------------------------- track=Tracks(PreFitCl,FitCl,pdg_list,number_of_tracks) ClCountPerEvent=track.GetClCountPerEvent() PreFitComp=track.GetPreFitComp() FitComp=track.GetFitComp() #print "Pdg:", pdg_list #print "PreFitComp:", PreFitComp #print "FitComp:", FitComp #Mapping---------------------------------------------------------------- Map=track.GetMap() #print "Map:", Map #Tracksplitting--------------------------------------------------------- split = track.GetTracksplitting() #print "Tracksplitting:", split htrkspl.Fill(split) #Completeness----------------------------------------------------------- Completeness=track.GetCompleteness() #print "Completeness:", Completeness for i in Completeness: hcomp.Fill(i) #Contamination---------------------------------------------------------- Contamination=track.GetContamination() #print "Contamination:", Contamination for i in Contamination: hcont.Fill(i) #SaveDataInFile--------------------------------------------------------- #print "****************************************************************" tracksplitting=htrkspl.GetMean() completeness=hcomp.GetMean() contamination=hcont.GetMean() dic=[1/tracksplitting, completeness/100, contamination/100] if args.filename!="not": data = file(args.foldertosave+"clus_"+args.filename+".txt", "w") for di in dic: data.write(str(di) + "\n") data.close() else: data = file('pics_opt_track_finder_parameters/tmp.txt', 'w') for di in dic: data.write(str(di) + "\n") data.close()