test/kalmanTools/createGains.py

0001 import ROOT,itertools,math
0002 from array import array
0003 from DataFormats.FWLite import Events, Handle
0004 ROOT.FWLiteEnabler.enable()
0005
0006 def getBit(q,i):
0007     return ((1<<i) & q)>>i
0008
0009 def fetchKMTF(event,etaMax=0.83,chi2=800000,dxyCut=100000):
0010     kmtfH  = Handle('BXVector<L1MuKBMTrack>')
0011     event.getByLabel('simKBmtfDigis',kmtfH)
0012     kmtf=kmtfH.product()
0013     out=[]
0014     for bx in [0]:
0015         for j in range(0,kmtf.size(bx)):
0016             mu =  kmtf.at(bx,j)
0017             if abs(mu.eta())<etaMax and mu.approxChi2()<chi2 and abs(mu.dxy())<dxyCut:
0018                 out.append(mu)
0019     return sorted(out,key=lambda x: x.pt(),reverse=True)
0020
0021 ####Save the Kalman Gains for LUTs
0022 kalmanGain={}
0023 kalmanGain2H={}
0024 kalmanGain2L={}
0025 kalmanGain2={}
0026 for track in [3,5,6,7,9,10,11,12,13,14,15]:
0027     for station in [3,2,1]:
0028         if getBit(track,station-1)==0:
0029             continue
0030         partialMask = track & (15<<station)
0031         if partialMask==0:
0032             continue;
0033
0034         if track in [3,5,6,9,10,12]:
0035             if not (partialMask in kalmanGain2.keys()):
0036                 kalmanGain2[partialMask] = {}
0037             if not (partialMask in kalmanGain2H.keys()):
0038                 kalmanGain2H[partialMask] = {}
0039             if not (partialMask in kalmanGain2L.keys()):
0040                 kalmanGain2L[partialMask] = {}
0041             kalmanGain2[partialMask][station]={}
0042             for q1 in ['H', 'L']:
0043                 kalmanGain2[partialMask][station][q1]={}
0044                 for q2 in ['H', 'L']:
0045                     kalmanGain2[partialMask][station][q1][q2]={}
0046
0047                     kalmanGain2[partialMask][station][q1][q2][0]=ROOT.TH2D("gain2_{track}_{station}_0_{q1}{q2}".format(track=partialMask,station=station,q1=q1,q2=q2),"h",64,0,512,256*2,-100*2,100*2)
0048                     kalmanGain2[partialMask][station][q1][q2][1]=ROOT.TH2D("gain2_{track}_{station}_1_{q1}{q2}".format(track=partialMask,station=station,q1=q1,q2=q2),"h",64,0,512,256*4,-8*4,8*4)
0049                     kalmanGain2[partialMask][station][q1][q2][4]=ROOT.TH2D("gain2_{track}_{station}_4_{q1}{q2}".format(track=partialMask,station=station,q1=q1,q2=q2),"h",64,0,512,256*2,-15*2,0)
0050                     kalmanGain2[partialMask][station][q1][q2][5]=ROOT.TH2D("gain2_{track}_{station}_5_{q1}{q2}".format(track=partialMask,station=station,q1=q1,q2=q2),"h",64,0,512,256*2,0,1*2)
0051
0052         else:
0053             if not (partialMask in kalmanGain.keys()):
0054                 kalmanGain[partialMask] = {}
0055             kalmanGain[partialMask][station] = {}
0056             kalmanGain[partialMask][station][0]=ROOT.TH2D("gain_{track}_{station}_0".format(track=partialMask,station=station),"h",64,0,1024,256,-100,100)
0057             kalmanGain[partialMask][station][4]=ROOT.TH2D("gain_{track}_{station}_4".format(track=partialMask,station=station),"h",64,0,1024,256,-15,0)
0058
0059
0060     for station in [0]:
0061         if not (track in kalmanGain.keys()):
0062             kalmanGain[track] = {}
0063         kalmanGain[track][0]={}
0064         kalmanGain[track][0][0]=ROOT.TH2D("gain_{track}_0_0".format(track=track),"h",64,0,1024,128,-5,5)
0065         kalmanGain[track][0][1]=ROOT.TH2D("gain_{track}_0_1".format(track=track),"h",64,0,1024,128,-5,5)
0066
0067
0068 ##############################
0069
0070
0071
0072
0073 for p in [3,5,6,7,9,10,11,12,13,14,15]:
0074     events=Events(['singleMu0_{}.root'.format(p)]) # Run KBMTF with only 1 pattern at a time to increase statistics
0075     counter=-1
0076     for event in events:
0077         counter=counter+1
0078         #fetch stubs
0079         kmtf=[]
0080         kmtf = fetchKMTF(event,1.5,1000000,1000000)
0081         ##Fill histograms and rates
0082         for track in kmtf:
0083             mask = track.hitPattern()
0084             qual = {}
0085             q1 = 'L'
0086             for stub in track.stubs():
0087                 qual[stub.stNum()] = stub.quality()
0088             if qual[max(qual)] >= 4:
0089                 q1='H'
0090             for station in [3,2,1]:
0091                 if not getBit(mask,station-1):
0092                     continue
0093                 gain = track.kalmanGain(station)
0094                 partialMask = mask & (15<<station)
0095                 if partialMask==0:
0096                     continue
0097                 q2 = 'L'
0098                 if qual[station] >= 4:
0099                     q2 = 'H'
0100                 if mask in [3,5,6,9,10,12]:
0101                     for element in [0,1,4,5]:
0102                         kalmanGain2[partialMask][station][q1][q2][element].Fill(gain[0]/8,gain[element+1])
0103                 else:
0104                     for element in [0,4]:
0105                         kalmanGain[partialMask][station][element].Fill(gain[0]/4,gain[element+1])
0106
0107             for station in [0]:
0108                 gain = track.kalmanGain(station)
0109                 kalmanGain[mask][station][0].Fill(gain[0]/2,gain[1])
0110                 kalmanGain[mask][station][1].Fill(gain[0]/2,gain[2])
0111
0112
0113 f=ROOT.TFile("gains.root","RECREATE")
0114
0115
0116 for k in kalmanGain.keys():
0117     for s in kalmanGain[k].keys():
0118         for e in kalmanGain[k][s].keys():
0119             kalmanGain[k][s][e].Write()
0120
0121 for k in kalmanGain2.keys():
0122     for s in kalmanGain2[k].keys():
0123         for q1 in kalmanGain2[k][s].keys():
0124             for q2 in kalmanGain2[k][s][q1].keys():
0125                 for e in kalmanGain2[k][s][q1][q2].keys():
0126                     kalmanGain2[k][s][q1][q2][e].Write()
0127
0128 f.Close()
0129
0130
0131
0132
0133
0134
0135
0136