LoopAnalyse.py
Go to the documentation of this file.
import ROOT,os,sys,getopt
import rootUtils as ut
import shipunit as u
import numpy as n
import collections
from ShipGeoConfig import ConfigRegistry
import shipDet_conf
import BookHistos
from MCTrackInfo import MCTrackInfo
from StrawHits import StrawHits
from FitTrackInfo import FitTrackInfo
import RecoSettings
class LoopAnalyse(object):
def __init__(self, tree, pref, geoFile, debug=0,
plotField = False,
plotTracking = False,
plotVertex = False,
plotVertexDetailed = False,
plotMomentum = False,
compareFits = False):
self.tree = tree
self.debug = debug
self.modules = None
self.ShipGeo = None
self.run = None
self.createGeom(pref, geoFile)
self.stat = collections.OrderedDict()
self.hmagnet = {}
self.htrck = {}
self.hvtx = {}
self.hmom = {}
self.hcmfits = {}
self.plotField = False
self.plotTracking = False
self.plotVertex = False
self.plotVertexCh = False
self.plotMomentum = False
self.compareFits = False
if plotField:
self.plotField = True
BookHistos.bookMFieldHistos(self.hmagnet)
if plotTracking:
self.plotTracking = True
BookHistos.bookTrackingHistos(self.htrck)
if plotVertex:
self.plotVertex = True
BookHistos.bookVertexHistos(self.hvtx)
if plotVertexDetailed:
self.plotVertexCh = True
BookHistos.bookVertexHistos(self.hvtx,"mUp")
BookHistos.bookVertexHistos(self.hvtx,"mDn")
if plotMomentum:
self.plotMomentum = True
BookHistos.bookMomentumHistos(self.hmom)
if compareFits:
self.compareFits = True
BookHistos.bookCompareFitHistos(self.hcmfits)
if plotVertex : BookHistos.bookVertexHistos(self.hvtx, "NEW")
if plotMomentum: BookHistos.bookMomentumHistos(self.hmom, "NEW")
self.MCTracks = MCTrackInfo(self.tree, 0)#self.debug)
self.StrHits = StrawHits(self.tree, self.modules, self.ShipGeo.straw.resol, self.debug, self.hmagnet, self.ShipGeo)
self.FitTracks = FitTrackInfo(self.tree, 0)#self.debug)
# FIXneeded - to introduce one more FixTracks!
def createGeom(self,pref, geoFile):
self.ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = pref )
self.run = ROOT.FairRunSim()
self.modules = shipDet_conf.configure(self.run,self.ShipGeo)
tgeom = ROOT.TGeoManager("Geometry", "Geane geometry")
gMan = tgeom.Import(geoFile)
geoMat = ROOT.genfit.TGeoMaterialInterface()
ROOT.genfit.MaterialEffects.getInstance().init(geoMat)
volDict = {}
i=0
for x in ROOT.gGeoManager.GetListOfVolumes():
volDict[i]=x.GetName()
i+=1
bfield = ROOT.genfit.BellField(self.ShipGeo.Bfield.max ,self.ShipGeo.Bfield.z,2, self.ShipGeo.Yheight/2.)
fM = ROOT.genfit.FieldManager.getInstance()
fM.init(bfield)
def CleanStat(self):
self.stat['total'] = 0
self.stat['HNLnTr'] = 0
self.stat['HNLzAcc'] = 0
self.stat['HNLeAcc'] = 0
self.stat['twoFitsO'] = 0
self.stat['twoFitsN'] = 0
self.stat['twoFitsNdfO'] = 0
self.stat['twoFitsNdfN'] = 0
def writeHistos(self):
ut.writeHists(htrck,"LoopA.root")
def hitMCtrackCorellations(self, weight=1):
# FIXneeded - add HNL vertex distribution plots!
if ( not self.plotTracking):
print "hitMCtrackCorellations does nothing since BookHistos.bookTrackingHistos() was not called!"
return
for tid in self.StrHits.getTrIDs(): # only multi-hits>25 tracks
if ( self.StrHits.checkVetoHits(tid) > 0) : continue # remove ones with vetoStattion
vec3h = self.StrHits.getStartHit(tid)
vec3p = self.MCTracks.getTrackPropagation(tid, vec3h.Z())
if(self.debug>0):
print "extr.", tid, "\t", vec3p.X(), vec3p.Y(), vec3p.Z(),
print "\tcharge:", self.MCTracks.getCharge(tid), "\tdx=", vec3p.X()-vec3h.X(), "dy=", vec3p.Y()-vec3h.Y()
dR = ROOT.TMath.Sqrt( (vec3p.X()-vec3h.X())**2 + (vec3p.Y()-vec3h.Y())**2 )
self.htrck['dRvsPz'].Fill(self.MCTracks.getMomentum(tid).Z(), dR, weight)
self.htrck['dRvsZhit'].Fill(vec3h.Z(), dR, weight)
charge = self.MCTracks.getCharge(tid)
theRatio = abs(vec3p.Y())/(vec3h.Z()-self.MCTracks.getHNLdecayVertex().Z())
HNLflag = tid in self.MCTracks.getHNLdecayTrIDs()
if HNLflag:
self.htrck['YMC2dZMCVer'].Fill(theRatio,weight)
if charge<0:
self.htrck['dxdyALLTrHm'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)
if HNLflag:
self.htrck['dxdyHNLTrHm'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)
elif charge>0:
self.htrck['dxdyALLTrHp'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)
if HNLflag:
self.htrck['dxdyHNLTrHp'].Fill(vec3p.X()-vec3h.X(), vec3p.Y()-vec3h.Y(), weight)
def fitMCvertexCorellations(self, weight=1, old=True, cut=False, detailed=False):
if( cut and detailed ) : print "fitMCvertexCorellations can't plot detailed vertex for cut! FIX me!"
vMC = self.MCTracks.getHNLdecayVertex()
if old :
vFt = self.FitTracks.getVertex()
else :
vFt = self.StrHits.getReFitVertex()
if (not vFt or not vMC) : return
if(self.plotVertex or self.plotVertexCh):
distanceToTracker = self.ShipGeo.TrackStation1['z']-vMC.Z()
coord = ['x','y','z']
iz = 1
while ( (iz<11) and (distanceToTracker>iz*500) ):
iz = iz+1
pref=""
prefCut=""
if( cut ) : prefCut='ndf'
if( not old ): pref='NEW'
for i in range(0,3):
delta = vMC(i)-vFt(i)
if( self.plotVertex ):
hname = 'vd'+coord[i]+'2'+prefCut+pref
self.hvtx[hname].Fill(distanceToTracker, delta, weight)
if (not cut):
if (iz<11):
hname='vd'+coord[i]+'2_'+str(iz)+pref
self.hvtx[hname].Fill(delta, weight)
if( self.plotVertexCh and old ):
sign = self.MCTracks.checkVertexUpDown()
if ( sign==0 ): continue
if ( sign<0) :
pref = "mUp"
else:
pref = "mDn"
hname = 'vd'+coord[i]+'2'+pref
self.hvtx[hname].Fill(distanceToTracker, delta, weight)
if( pref=='NEW') :
doca = self.StrHits.getStepDoca(RecoSettings.VertexExtrSteps-1)
if ( doca!=None ) : self.hvtx['doca'+pref].Fill(doca)
doca0 = self.StrHits.getStepDoca(0)
if ( doca0!=None ) : self.hvtx['doca0'+pref].Fill(doca0)
if ( doca0!=None and doca0!=0):
self.hvtx['docaRel'+pref].Fill(1,1.)
for ii in range (1, RecoSettings.VertexExtrSteps):
doca = self.StrHits.getStepDoca(ii)
#print "...", doca, doca/doca0
if ( doca!=None ) :
self.hvtx['docaRel'+pref].Fill(ii+1, doca/doca0)
return vFt
def doCompareFits(self, oldStyle=True, ellCut=False, weight=1):
ntracks = self.StrHits.FitTracks(oldStyle) # does re-fitting
if (self.debug>0) : self.StrHits.PrintNewTracks()
fitOldTrIDs = self.FitTracks.getTrIDs()
fitNewTrIDs = self.StrHits.getReFitTrIDs()
for tid in fitOldTrIDs:
self.hcmfits['Chi2NDFold'].Fill(self.FitTracks.getChi2Ndf(tid))
self.hcmfits['NDFold'].Fill(self.FitTracks.getNdf(tid))
for tid in fitNewTrIDs:
self.hcmfits['Chi2NDFnew'].Fill(self.StrHits.getReFitChi2Ndf(tid))
self.hcmfits['NDFnew'].Fill(self.StrHits.getReFitNdf(tid))
for tid in self.StrHits.getTrIDsALL(): # hists with MCTrID>0
told = ( tid in fitOldTrIDs)
tnew = ( tid in fitNewTrIDs)
if( (not told) and (not tnew) ):
self.hcmfits['FitCorrespondence'].Fill(0)
elif( (not told) and tnew):
self.hcmfits['FitCorrespondence'].Fill(1)
elif( told and (not tnew) ):
self.hcmfits['FitCorrespondence'].Fill(2)
else:
eq = self.StrHits.compareFitTracks(tid, self.FitTracks)
if eq:
self.hcmfits['FitCorrespondence'].Fill(5)
else:
self.hcmfits['FitCorrespondence'].Fill(6)
self.hcmfits['dNFoldFNnew'].Fill(len(fitOldTrIDs), len(fitNewTrIDs))
return ntracks
def doPlotMomentum(self, weight, old=True):
tids = []
pfix = ""
if old:
tids = self.FitTracks.getTrIDs()
else:
tids = self.StrHits.getReFitTrIDs()
pfix = "NEW"
for tid in tids:
r,d,pFit = None, None, None
chi2ndf,ndf = None, None
if old:
r,d,pFit = self.FitTracks.getPosDirPval(tid)
ndf = self.FitTracks.getNdf(tid)
chi2ndf = self.FitTracks.getChi2Ndf(tid)
else:
r,d,pFit = self.StrHits.getReFitPosDirPval(tid)
ndf = self.StrHits.getReFitNdf(tid)
chi2ndf = self.StrHits.getReFitChi2Ndf(tid)
if( not tid in self.MCTracks.getTrIDs() ): continue
pMC = self.MCTracks.getMomentum(tid).Mag()
dPP = (pMC - pFit)/pMC
hname = 'delPOverP0'+pfix
self.hmom[hname].Fill(chi2ndf, dPP, weight)
hname = 'delPOverP1'+pfix
self.hmom[hname].Fill(pMC, dPP, weight)
if( ndf<RecoSettings.trackMinNofHits ): continue
hname = 'delPOverP2'+pfix
self.hmom[hname].Fill(pMC, dPP, weight)
if( not tid in self.MCTracks.getHNLdecayTrIDs() ): continue
hname = 'delPOverP3'+pfix
self.hmom[hname].Fill(pMC, dPP, weight)
iz = 1
while ( (iz<6) and (pMC>iz*10) ):
iz = iz+1
if (iz<6):
hname = 'delPOverP3_'+str(iz)+pfix
self.hmom[hname].Fill(dPP, weight)
def PrintStat(self):
for stage in self.stat:
print "{0:15}{1:10}".format(stage, self.stat[stage])
################################################################################################
def Loop(self, stat=-1):
self.CleanStat()
if(stat==-1):
nEvents = self.tree.GetEntries()
else:
nEvents = min(self.tree.GetEntries(),stat)
# --------------- cut related values ------------------------------------------------------#
# FIXneeded - should be done depending on RecoSettings.trackMinNofStations!
# self.MCTracks.setAccPlaneZ(3380*u.cm) # - after 3rd Station
self.MCTracks.setAccPlaneZ(2500*u.cm) #- before 1st Station
# -----------------------------------------------------------------------------------------#
# -----------------------------------------------------------------------------------------#
for n in range(nEvents):
rc = self.tree.GetEntry(n)
self.stat['total'] +=1
# ------------------- MCTracks -----------------------------------------------------------#
nMCPrimTracks = self.MCTracks.readEvent()
if( self.debug>0) :
print "\n\nevent ", n, "===================================================="
self.MCTracks.PrintHNL()
#=>
if( nMCPrimTracks!=2 ) : continue
self.stat['HNLnTr']+=1
#=>
# FIXneeded - should be setting from ShipGeo here!
if( (self.MCTracks.getHNLdecayVertex().Z()<-1900*u.cm) \
or (self.MCTracks.getHNLdecayVertex().Z()>RecoSettings.VertexMaxZcut) ) : continue
self.stat['HNLzAcc']+=1
#=>
if( (not self.MCTracks.checkEllipticAcc()) ) : continue
self.stat['HNLeAcc'] +=1
if( self.debug>1 ): print "MC event accepted"
weight = self.MCTracks.getEventWeight()
#weight = 1
# ------------------- FitTracks -----------------------------------------------------------#
nFitTracks = self.FitTracks.readEvent()
if(self.plotMomentum):self.doPlotMomentum(weight, old=True)
if(self.debug>0) : self.FitTracks.Print()
printMCVert = False
if (nFitTracks==2):
self.stat['twoFitsO']+=1
v=self.fitMCvertexCorellations(weight, old=True, cut=False, detailed=self.plotVertexCh)
printMCVert = True
print "\n\nOld vertex ", v.X(), v.Y(), v.Z()
ndfok = True
for tid in self.FitTracks.getTrIDs() :
ndfok = ( ndfok and (self.FitTracks.getNdf(tid)>RecoSettings.trackMinNofHits) )
if ( not ndfok ) : break
if ndfok : # FIXneeded - not optimal to call it once more :(
self.stat['twoFitsNdfO']+=1
v=self.fitMCvertexCorellations(weight, old=True, cut=True, detailed=False)
print "the old vertex passes ndf cut!"
# -------------------- StrawHits --------------------------------------------------------------#
if ( self.compareFits or self.plotTracking ) :
self.StrHits.readEvent()
# -------------------- ReFit --------------------------------------------------------------#
nNewFitTracks = 0
if self.compareFits :
nNewFitTracks = self.doCompareFits(oldStyle=True, ellCut=False, weight=weight)
if(self.plotMomentum):self.doPlotMomentum(weight, old=False)
if (nNewFitTracks==2):
self.stat['twoFitsN']+=1
v=self.fitMCvertexCorellations(weight, old=False)
print "New vertex ", v.X(), v.Y(), v.Z()
printMCVert = True
ndfok = True
for tid in self.StrHits.getReFitTrIDs() :
ndfok = ( ndfok and (self.StrHits.getReFitNdf(tid)>RecoSettings.trackMinNofHits) )
if ( not ndfok ) : break
if ndfok : # FIXneeded - not optimal to call it once more :(
self.stat['twoFitsNdfN']+=1
v=self.fitMCvertexCorellations(weight, old=False, cut=True, detailed=False)
print "the new vertex passes ndf cut!"
if printMCVert : self.MCTracks.PrintHNL()
if(self.plotTracking) : self.hitMCtrackCorellations(weight)
# end event loop
#=======================================
print "here -> ", self.plotField
if self.plotField : BookHistos.plotMFieldHistos(self.hmagnet)
if self.plotTracking : BookHistos.plotMCTrackHitCorrelation(self.htrck)
if self.plotVertex : BookHistos.plotVertexHistos(self.hvtx)
if self.plotVertexCh :
BookHistos.plotVertexHistos(self.hvtx,"mUp")
BookHistos.plotVertexHistos(self.hvtx,"mDn")
if self.plotMomentum : BookHistos.plotMomentumHistos(self.hmom)
if self.compareFits :
BookHistos.plotCompareFitHistos(self.hcmfits)
if self.plotVertex : BookHistos.plotVertexHistos(self.hvtx,'NEW')
if self.plotMomentum : BookHistos.plotMomentumHistos(self.hmom, 'NEW')
#=======================================
self.PrintStat()
BookHistos.plotCompareFitHistos
def plotCompareFitHistos
Definition: BookHistos.py:171
LoopAnalyse.LoopAnalyse.plotMomentum
plotMomentum
Definition: LoopAnalyse.py:46
LoopAnalyse.LoopAnalyse.doCompareFits
def doCompareFits
Definition: LoopAnalyse.py:203
LoopAnalyse.LoopAnalyse.plotVertex
plotVertex
Definition: LoopAnalyse.py:44
LoopAnalyse.LoopAnalyse.MCTracks
MCTracks
Definition: LoopAnalyse.py:73
LoopAnalyse.LoopAnalyse.htrck
htrck
Definition: LoopAnalyse.py:38
BookHistos.plotMomentumHistos
def plotMomentumHistos
Definition: BookHistos.py:87
LoopAnalyse.LoopAnalyse.plotField
plotField
Definition: LoopAnalyse.py:42
LoopAnalyse.LoopAnalyse.tree
tree
Definition: LoopAnalyse.py:27
LoopAnalyse.LoopAnalyse.writeHistos
def writeHistos
Definition: LoopAnalyse.py:109
LoopAnalyse.LoopAnalyse.createGeom
def createGeom
Definition: LoopAnalyse.py:79
LoopAnalyse.LoopAnalyse.Loop
def Loop
Definition: LoopAnalyse.py:301
LoopAnalyse.LoopAnalyse.StrHits
StrHits
Definition: LoopAnalyse.py:74
LoopAnalyse.LoopAnalyse.hcmfits
hcmfits
Definition: LoopAnalyse.py:41
LoopAnalyse.LoopAnalyse.CleanStat
def CleanStat
Definition: LoopAnalyse.py:99
BookHistos.plotMCTrackHitCorrelation
def plotMCTrackHitCorrelation
Definition: BookHistos.py:229
BookHistos.plotVertexHistos
def plotVertexHistos
Definition: BookHistos.py:312
LoopAnalyse.LoopAnalyse.hmom
hmom
Definition: LoopAnalyse.py:40
LoopAnalyse.LoopAnalyse.plotTracking
plotTracking
Definition: LoopAnalyse.py:43
LoopAnalyse.LoopAnalyse.PrintStat
def PrintStat
Definition: LoopAnalyse.py:293
LoopAnalyse.LoopAnalyse.compareFits
compareFits
Definition: LoopAnalyse.py:47
LoopAnalyse.LoopAnalyse.__init__
def __init__
Definition: LoopAnalyse.py:26
LoopAnalyse.LoopAnalyse.ShipGeo
ShipGeo
Definition: LoopAnalyse.py:31
LoopAnalyse.LoopAnalyse.hmagnet
hmagnet
Definition: LoopAnalyse.py:37
StrawHits.StrawHits
Definition: StrawHits.py:12
FitTrackInfo.FitTrackInfo
Definition: FitTrackInfo.py:5
LoopAnalyse.LoopAnalyse.run
run
Definition: LoopAnalyse.py:32
LoopAnalyse.LoopAnalyse.fitMCvertexCorellations
def fitMCvertexCorellations
Definition: LoopAnalyse.py:146
LoopAnalyse.LoopAnalyse.debug
debug
Definition: LoopAnalyse.py:28
LoopAnalyse.LoopAnalyse.FitTracks
FitTracks
Definition: LoopAnalyse.py:75
LoopAnalyse.LoopAnalyse.hvtx
hvtx
Definition: LoopAnalyse.py:39
BookHistos.bookVertexHistos
def bookVertexHistos
Definition: BookHistos.py:280
BookHistos.plotMFieldHistos
def plotMFieldHistos
Definition: BookHistos.py:38
LoopAnalyse.LoopAnalyse.plotVertexCh
plotVertexCh
Definition: LoopAnalyse.py:45
LoopAnalyse.LoopAnalyse.stat
stat
Definition: LoopAnalyse.py:35
LoopAnalyse.LoopAnalyse.doPlotMomentum
def doPlotMomentum
Definition: LoopAnalyse.py:238
BookHistos.bookMFieldHistos
def bookMFieldHistos
Definition: BookHistos.py:21
BookHistos.bookTrackingHistos
def bookTrackingHistos
Definition: BookHistos.py:206
MCTrackInfo.MCTrackInfo
For a single event stores HNL, its decay products and several other tracks (if requested) from ShipMC...
Definition: MCTrackInfo.py:10
BookHistos.bookCompareFitHistos
def bookCompareFitHistos
Definition: BookHistos.py:149
LoopAnalyse.LoopAnalyse.hitMCtrackCorellations
def hitMCtrackCorellations
Definition: LoopAnalyse.py:115
LoopAnalyse.LoopAnalyse
Definition: LoopAnalyse.py:19
LoopAnalyse.LoopAnalyse.modules
modules
Definition: LoopAnalyse.py:30
BookHistos.bookMomentumHistos
def bookMomentumHistos
Definition: BookHistos.py:67