Newer
Older
FairShipTools / newGen / offlineForBarbara_maybeOld.py
  1. from lookAtGeo import *
  2. import tools
  3. import shipunit as u
  4. from ShipGeoConfig import ConfigRegistry
  5. import shipDet_conf
  6.  
  7. from operator import mul, add
  8.  
  9. import sys
  10. sys.path.append('KaterinaLight/')
  11. from StrawHits import StrawHits
  12. ## Use it like:
  13. # f = TFile(fileName)
  14. # t = f.Get("cbmsim")
  15. # sh = offline.StrawHits(t, offline.shipDet_conf.configure(offline.__run, r['ShipGeo']), r['ShipGeo'].straw.resol, 0, None, r['ShipGeo'])
  16. # t.GetEntry(58)
  17. # sh.readEvent()
  18. # sh.FitTracks()
  19.  
  20. #dy = 10.
  21. # init geometry and mag. field
  22. #ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = dy )
  23.  
  24.  
  25. def searchForNodes3_xyz_dict(fGeo, verbose=False):
  26. from tools import findPositionElement, findDimentionBoxElement, findPositionElement2
  27. d = {}
  28. #r = loadGeometry(inputFile)
  29. #fGeo = r['fGeo']
  30. ## Get the top volume
  31. #fGeo = ROOT.gGeoManager
  32. tv = fGeo.GetTopVolume()
  33. topnodes = tv.GetNodes()
  34. for (j,topn) in enumerate(topnodes):
  35. # top volumes
  36. if verbose:
  37. print j, topn.GetName()
  38. print " x: ", findPositionElement(topn)['x'],findDimentionBoxElement(topn)['x']
  39. print " y: ", findPositionElement(topn)['y'],findDimentionBoxElement(topn)['y']
  40. print " z: ", findPositionElement(topn)['z'],findDimentionBoxElement(topn)['z']
  41. d[topn.GetName()] = {'x': {}, 'y':{}, 'z':{}, 'r':{}}
  42. d[topn.GetName()]['x']['pos'] = findPositionElement(topn)['x']
  43. d[topn.GetName()]['x']['dim'] = findDimentionBoxElement(topn)['x']
  44. d[topn.GetName()]['y']['pos'] = findPositionElement(topn)['y']
  45. d[topn.GetName()]['y']['dim'] = findDimentionBoxElement(topn)['y']
  46. d[topn.GetName()]['z']['pos'] = findPositionElement(topn)['z']
  47. d[topn.GetName()]['z']['dim'] = findDimentionBoxElement(topn)['z']
  48. if topn.GetVolume().GetShape().IsCylType():
  49. d[topn.GetName()]['r']['pos'] = findPositionElement(topn)['r']
  50. d[topn.GetName()]['r']['dim'] = findDimentionBoxElement(topn)['r']
  51. else:
  52. d[topn.GetName()]['r']['pos'] = 0.
  53. d[topn.GetName()]['r']['dim'] = 0.
  54. # First children
  55. nodes = topn.GetNodes()
  56. if nodes:
  57. topPos = topn.GetMatrix().GetTranslation()
  58. for (i,n) in enumerate(nodes):
  59. if verbose:
  60. print j, topn.GetName(), i, n.GetName()
  61. print " x: ", findPositionElement2(n,topPos)['x'],findDimentionBoxElement(n)['x']
  62. print " y: ", findPositionElement2(n,topPos)['y'],findDimentionBoxElement(n)['y']
  63. print " z: ", findPositionElement2(n,topPos)['z'],findDimentionBoxElement(n)['z']
  64. d[n.GetName()] = {'x': {}, 'y':{}, 'z':{}, 'r':{}}
  65. d[n.GetName()]['x']['pos'] = findPositionElement2(n,topPos)['x']
  66. d[n.GetName()]['x']['dim'] = findDimentionBoxElement(n)['x']
  67. d[n.GetName()]['y']['pos'] = findPositionElement2(n,topPos)['y']
  68. d[n.GetName()]['y']['dim'] = findDimentionBoxElement(n)['y']
  69. d[n.GetName()]['z']['pos'] = findPositionElement2(n,topPos)['z']
  70. d[n.GetName()]['z']['dim'] = findDimentionBoxElement(n)['z']
  71. if n.GetVolume().GetShape().IsCylType():
  72. d[n.GetName()]['r']['pos'] = findPositionElement2(n,topPos)['r']
  73. d[n.GetName()]['r']['dim'] = findDimentionBoxElement(n)['r']
  74. else:
  75. d[n.GetName()]['r']['pos'] = 0.
  76. d[n.GetName()]['r']['dim'] = 0.
  77. # Second children
  78. cnodes = n.GetNodes()
  79. if cnodes:
  80. localpos = n.GetMatrix().GetTranslation()
  81. localToGlobal = []
  82. for i in xrange(3):
  83. localToGlobal.append(localpos[i]+topPos[i])
  84. for (k,cn) in enumerate(cnodes):
  85. if verbose:
  86. print j, topn.GetName(), i, n.GetName(), k, cn.GetName()
  87. print " x: ", findPositionElement2(cn,localToGlobal)['x'],findDimentionBoxElement(cn)['x']
  88. print " y: ", findPositionElement2(cn,localToGlobal)['y'],findDimentionBoxElement(cn)['y']
  89. print " z: ", findPositionElement2(cn,localToGlobal)['z'],findDimentionBoxElement(cn)['z']
  90. d[cn.GetName()] = {'x': {}, 'y':{}, 'z':{}, 'r':{}}
  91. d[cn.GetName()]['x']['pos'] = findPositionElement2(cn,localToGlobal)['x']
  92. d[cn.GetName()]['x']['dim'] = findDimentionBoxElement(cn)['x']
  93. d[cn.GetName()]['y']['pos'] = findPositionElement2(cn,localToGlobal)['y']
  94. d[cn.GetName()]['y']['dim'] = findDimentionBoxElement(cn)['y']
  95. d[cn.GetName()]['z']['pos'] = findPositionElement2(cn,localToGlobal)['z']
  96. d[cn.GetName()]['z']['dim'] = findDimentionBoxElement(cn)['z']
  97. if cn.GetVolume().GetShape().IsCylType():
  98. d[cn.GetName()]['r']['pos'] = findPositionElement2(cn,localToGlobal)['r']
  99. d[cn.GetName()]['r']['dim'] = findDimentionBoxElement(cn)['r']
  100. else:
  101. d[cn.GetName()]['r']['pos'] = 0.
  102. d[cn.GetName()]['r']['dim'] = 0.
  103. return d
  104.  
  105.  
  106. ff_nu = ROOT.TFile("histoForWeights_nu.root")
  107. h_GioHans_nu = ff_nu.Get("h_Gio")
  108.  
  109. ff_antinu = ROOT.TFile("histoForWeights_antinu.root")
  110. h_GioHans_antinu = ff_antinu.Get("h_Gio")
  111.  
  112. def calcWeightNu(NC, E, w, entries, nuName, ON=True):
  113. # Only for neutrinos and antineutrinos
  114. if not ON:
  115. return 1
  116. if "bar" in nuName:
  117. reduction = 0.5
  118. flux = 1.#6.98e+11 * 2.e+20 / 5.e+13
  119. h_GioHans = h_GioHans_antinu
  120. else:
  121. reduction = 1.
  122. flux = 1.#1.09e+12 * 2.e+20/ 5.e+13
  123. h_GioHans = h_GioHans_nu
  124. crossSec = 6.7e-39*E * reduction
  125. NA = 6.022e+23
  126. binN = h_GioHans.GetXaxis().FindBin(E)
  127. return crossSec * flux * h_GioHans.GetBinContent(binN) * w * NA #/ entries
  128.  
  129.  
  130. def findWeight(sampleType, NC, E, MCTrack, entries, nuName, ON):
  131. if sampleType == 'nuBg':
  132. return calcWeightNu(NC, E, MCTrack.GetWeight(), entries, nuName, ON)
  133. elif sampleType == 'sig':
  134. return MCTrack.GetWeight() # for the acceptance, multiply by normalization
  135. elif sampleType == 'cosmics':
  136. return MCTrack.GetWeight() # multiply by 1.e6 / 200.
  137.  
  138.  
  139.  
  140. def hasStrawStations(event, trackId, listOfWantedStraws):
  141. ok = [False]*len(listOfWantedStraws)
  142. positions = [ (nodes[det]['z']['pos'] - nodes[det]['z']['dim'], nodes[det]['z']['pos'] + nodes[det]['z']['dim'] ) for det in listOfWantedStraws ]
  143. for hit in event.strawtubesPoint:
  144. if hit.GetTrackID() == trackId:
  145. for (i,det) in enumerate(listOfWantedStraws):
  146. if (positions[i][0] < hit.GetZ() < positions[i][1]) and tools.inEllipse(hit.GetX(), hit.GetY(), 250., 500.):
  147. ok[i] = True
  148. return bool(reduce(mul, ok, 1))
  149.  
  150. def hasGoodStrawStations(event, trackId):
  151. #ok = [False]*2
  152. okupstream = [False]*2
  153. okdownstream = [False]*2
  154. upstream = [ (nodes[det]['z']['pos'] - nodes[det]['z']['dim'], nodes[det]['z']['pos'] + nodes[det]['z']['dim'] ) for det in ['Tr1_1', 'Tr2_2'] ]
  155. downstream = [ (nodes[det]['z']['pos'] - nodes[det]['z']['dim'], nodes[det]['z']['pos'] + nodes[det]['z']['dim'] ) for det in ['Tr3_3', 'Tr4_4'] ]
  156. for hit in event.strawtubesPoint:
  157. if hit.GetTrackID() == trackId:
  158. for i in xrange(2):
  159. if (upstream[i][0] < hit.GetZ() < upstream[i][1]) and tools.inEllipse(hit.GetX(), hit.GetY(), 250., 500.):
  160. okupstream[i] = True
  161. if (downstream[i][0] < hit.GetZ() < downstream[i][1]) and tools.inEllipse(hit.GetX(), hit.GetY(), 250., 500.):
  162. okdownstream[i] = True
  163. ok = [ bool(reduce(mul, l, 1)) for l in [okupstream, okdownstream] ]
  164. return bool(reduce(add, ok, 0))
  165.  
  166. def findHNLvertex(event):
  167. for t in event.MCTrack:
  168. if t.GetMotherId() == 1:
  169. return t.GetStartZ()
  170. return False
  171.  
  172. def has_muon_station(event, trackId, station):
  173. zIn = nodes['muondet%s_1'%(station-1)]['z']['pos'] - nodes['muondet%s_1'%(station-1)]['z']['dim']
  174. zOut = nodes['muondet%s_1'%(station-1)]['z']['pos'] + nodes['muondet%s_1'%(station-1)]['z']['dim']
  175. for hit in event.muonPoint:
  176. if hit.GetTrackID() == trackId:
  177. if zIn <= hit.GetZ() <= zOut:
  178. return True
  179. return False
  180.  
  181. def hasEcalDeposit(event, trackId, ELossThreshold):
  182. ELoss = 0.
  183. for hit in event.EcalPoint:
  184. if hit.GetTrackID() == trackId:
  185. ELoss += hit.GetEnergyLoss()
  186. if ELoss >= ELossThreshold:
  187. return True
  188. return False
  189.  
  190. def hasMuons(event, trackId):
  191. m1 = 0
  192. m2 = 0
  193. m3 = 0
  194. m4 = 0
  195. for ahit in event.muonPoint:
  196. if ahit.GetTrackID() == trackId:
  197. detID = ahit.GetDetectorID()
  198. if(detID == 476) :
  199. m1 += 1
  200. if(detID == 477) :
  201. m2 += 1
  202. if(detID == 478) :
  203. m3 += 1
  204. if(detID == 479) :
  205. m4 += 1
  206. return [bool(m1), bool(m2), bool(m3), bool(m4)]
  207.  
  208. def myVertex(t1,t2,PosDir):
  209. # closest distance between two tracks
  210. # d = |pq . u x v|/|u x v|
  211. a = ROOT.TVector3(PosDir[t1][0](0) ,PosDir[t1][0](1), PosDir[t1][0](2))
  212. u = ROOT.TVector3(PosDir[t1][1](0),PosDir[t1][1](1),PosDir[t1][1](2))
  213. c = ROOT.TVector3(PosDir[t2][0](0) ,PosDir[t2][0](1), PosDir[t2][0](2))
  214. v = ROOT.TVector3(PosDir[t2][1](0),PosDir[t2][1](1),PosDir[t2][1](2))
  215. pq = a-c
  216. uCrossv = u.Cross(v)
  217. dist = pq.Dot(uCrossv)/(uCrossv.Mag()+1E-8)
  218. # u.a - u.c + s*|u|**2 - u.v*t = 0
  219. # v.a - v.c + s*v.u - t*|v|**2 = 0
  220. E = u.Dot(a) - u.Dot(c)
  221. F = v.Dot(a) - v.Dot(c)
  222. A,B = u.Mag2(), -u.Dot(v)
  223. C,D = u.Dot(v), -v.Mag2()
  224. t = -(C*E-A*F)/(B*C-A*D)
  225. X = c.x()+v.x()*t
  226. Y = c.y()+v.y()*t
  227. Z = c.z()+v.z()*t
  228. # sT = ROOT.gROOT.FindAnything('cbmsim')
  229. #print 'test2 ',X,Y,Z,dist
  230. #print 'truth',sTree.MCTrack[2].GetStartX(),sTree.MCTrack[2].GetStartY(),sTree.MCTrack[2].GetStartZ()
  231. return X,Y,Z,abs(dist)
  232.  
  233. def addFullInfoToTree(elenaTree):
  234. elenaTree, DaughtersPt = tools.AddVect(elenaTree, 'DaughtersPt', 'float')
  235. elenaTree, DaughtersChi2 = tools.AddVect(elenaTree, 'DaughtersChi2', 'float')
  236. elenaTree, DaughtersNPoints = tools.AddVect(elenaTree, 'DaughtersNPoints', 'int')
  237. elenaTree, DaughtersTruthProdX = tools.AddVect(elenaTree, 'DaughtersTruthProdX', 'float')
  238. elenaTree, DaughtersTruthProdY = tools.AddVect(elenaTree, 'DaughtersTruthProdY', 'float')
  239. elenaTree, DaughtersTruthProdZ = tools.AddVect(elenaTree, 'DaughtersTruthProdZ', 'float')
  240. elenaTree, DaughtersTruthPDG = tools.AddVect(elenaTree, 'DaughtersTruthPDG', 'int')
  241. elenaTree, DaughtersTruthMotherPDG = tools.AddVect(elenaTree, 'DaughtersTruthMotherPDG', 'int')
  242. elenaTree, DaughtersFitConverged = tools.AddVect(elenaTree, 'DaughtersFitConverged', 'int')
  243. elenaTree, straw_x = tools.AddVect(elenaTree, 'straw_x', 'float')
  244. elenaTree, straw_y = tools.AddVect(elenaTree, 'straw_y', 'float')
  245. elenaTree, straw_z = tools.AddVect(elenaTree, 'straw_z', 'float')
  246. elenaTree, muon_x = tools.AddVect(elenaTree, 'muon_x', 'float')
  247. elenaTree, muon_y = tools.AddVect(elenaTree, 'muon_y', 'float')
  248. elenaTree, muon_z = tools.AddVect(elenaTree, 'muon_z', 'float')
  249. elenaTree, ecal_x = tools.AddVect(elenaTree, 'ecal_x', 'float')
  250. elenaTree, ecal_y = tools.AddVect(elenaTree, 'ecal_y', 'float')
  251. elenaTree, ecal_z = tools.AddVect(elenaTree, 'ecal_z', 'float')
  252. elenaTree, hcal_x = tools.AddVect(elenaTree, 'hcal_x', 'float')
  253. elenaTree, hcal_y = tools.AddVect(elenaTree, 'hcal_y', 'float')
  254. elenaTree, hcal_z = tools.AddVect(elenaTree, 'hcal_z', 'float')
  255. elenaTree, veto5_x = tools.AddVect(elenaTree, 'veto5_x', 'float')
  256. elenaTree, veto5_y = tools.AddVect(elenaTree, 'veto5_y', 'float')
  257. elenaTree, veto5_z = tools.AddVect(elenaTree, 'veto5_z', 'float')
  258. elenaTree, liquidscint_x = tools.AddVect(elenaTree, 'liquidscint_x', 'float')
  259. elenaTree, liquidscint_y = tools.AddVect(elenaTree, 'liquidscint_y', 'float')
  260. elenaTree, liquidscint_z = tools.AddVect(elenaTree, 'liquidscint_z', 'float')
  261. elenaTree, DOCA = tools.AddVar(elenaTree, 'DOCA', 'float')
  262. elenaTree, vtxx = tools.AddVar(elenaTree, 'vtxx', 'float')
  263. elenaTree, vtxy = tools.AddVar(elenaTree, 'vtxy', 'float')
  264. elenaTree, vtxz = tools.AddVar(elenaTree, 'vtxz', 'float')
  265. elenaTree, IP0 = tools.AddVar(elenaTree, 'IP0', 'float')
  266. elenaTree, Mass = tools.AddVar(elenaTree, 'Mass', 'float')
  267. elenaTree, Pt = tools.AddVar(elenaTree, 'Pt', 'float')
  268. elenaTree, P = tools.AddVar(elenaTree, 'P', 'float')
  269. elenaTree, NParticles = tools.AddVar(elenaTree, 'NParticles', 'int')
  270. elenaTree, HNLw = tools.AddVar(elenaTree, 'HNLw', 'float')
  271. elenaTree, NuWeight = tools.AddVar(elenaTree, 'NuWeight', 'float')
  272. elenaTree, EventNumber = tools.AddVar(elenaTree, 'EventNumber', 'int')
  273. elenaTree, DaughterMinPt = tools.AddVar(elenaTree, 'DaughterMinPt', 'float')
  274. elenaTree, DaughterMinP = tools.AddVar(elenaTree, 'DaughterMinP', 'float')
  275. elenaTree, DaughtersAlwaysIn = tools.AddVar(elenaTree, 'DaughtersAlwaysIn', 'int')
  276. elenaTree, BadTruthVtx = tools.AddVar(elenaTree, 'BadTruthVtx', 'int')
  277.  
  278. DaughtersFitConverged, DOCA, vtxx, vtxy, vtxz, IP0, HasEcal = None, None, None, None, None, None, None
  279. NoB_DOCA, NoB_vtxx, NoB_vtxy, NoB_vtxz, NoB_IP0 = None, None, None, None, None
  280. DaughtersAlwaysIn, BadTruthVtx, Has1Muon1, Has1Muon2, Has2Muon1, Has2Muon2 = None, None, None, None, None, None
  281. MaxDaughtersRedChi2, MinDaughtersNdf = None, None
  282. NoB_MaxDaughtersRedChi2, NoB_MinDaughtersNdf = None, None
  283. DaughtersMinP, DaughtersMinPt, Mass, P, Pt = None, None, None, None, None
  284. NoB_DaughtersMinP, NoB_DaughtersMinPt, NoB_Mass, NoB_P, NoB_Pt = None, None, None, None, None
  285.  
  286. def addOfflineToTree(elenaTree):
  287. global DaughtersFitConverged, DOCA, vtxx, vtxy, vtxz, IP0, HasEcal
  288. global NoB_DOCA, NoB_vtxx, NoB_vtxy, NoB_vtxz, NoB_IP0
  289. global DaughtersAlwaysIn, BadTruthVtx, Has1Muon1, Has1Muon2, Has2Muon1, Has2Muon2
  290. global MaxDaughtersRedChi2, MinDaughtersNdf, HNLw, NuWeight, NoB_MaxDaughtersRedChi2, NoB_MinDaughtersNdf
  291. global DaughtersMinP, DaughtersMinPt, Mass, P, Pt
  292. global NoB_DaughtersMinP, NoB_DaughtersMinPt, NoB_Mass, NoB_P, NoB_Pt
  293. elenaTree, DaughtersFitConverged = tools.AddVect(elenaTree, 'DaughtersFitConverged', 'int') #
  294. elenaTree, DOCA = tools.AddVect(elenaTree, 'DOCA', 'float') #
  295. elenaTree, NoB_DOCA = tools.AddVect(elenaTree, 'NoB_DOCA', 'float') #
  296. elenaTree, vtxx = tools.AddVect(elenaTree, 'vtxxSqr', 'float') #
  297. elenaTree, vtxy = tools.AddVect(elenaTree, 'vtxySqr', 'float') #
  298. elenaTree, vtxz = tools.AddVect(elenaTree, 'vtxz', 'float') #
  299. elenaTree, NoB_vtxx = tools.AddVect(elenaTree, 'NoB_vtxxSqr', 'float') #
  300. elenaTree, NoB_vtxy = tools.AddVect(elenaTree, 'NoB_vtxySqr', 'float') #
  301. elenaTree, NoB_vtxz = tools.AddVect(elenaTree, 'NoB_vtxz', 'float') #
  302. elenaTree, IP0 = tools.AddVect(elenaTree, 'IP0', 'float') #
  303. elenaTree, NoB_IP0 = tools.AddVect(elenaTree, 'NoB_IP0', 'float') #
  304. #elenaTree, NParticles = tools.AddVar(elenaTree, 'NParticles', 'int') #
  305. elenaTree, DaughtersAlwaysIn = tools.AddVect(elenaTree, 'DaughtersAlwaysIn', 'int') #
  306. elenaTree, BadTruthVtx = tools.AddVect(elenaTree, 'BadTruthVtx', 'int') #
  307. elenaTree, Has1Muon1 = tools.AddVect(elenaTree, 'Has1Muon1', 'int') #
  308. elenaTree, Has1Muon2 = tools.AddVect(elenaTree, 'Has1Muon2', 'int') #
  309. elenaTree, Has2Muon1 = tools.AddVect(elenaTree, 'Has2Muon1', 'int') #
  310. elenaTree, Has2Muon2 = tools.AddVect(elenaTree, 'Has2Muon2', 'int') #
  311. elenaTree, HasEcal = tools.AddVect(elenaTree, 'HasEcal', 'int') #
  312. elenaTree, MaxDaughtersRedChi2 = tools.AddVect(elenaTree, 'MaxDaughtersRedChi2', 'float') #
  313. elenaTree, MinDaughtersNdf = tools.AddVect(elenaTree, 'MinDaughtersNdf', 'int') #
  314. elenaTree, NoB_MaxDaughtersRedChi2 = tools.AddVect(elenaTree, 'NoB_MaxDaughtersRedChi2', 'float') #
  315. elenaTree, NoB_MinDaughtersNdf = tools.AddVect(elenaTree, 'NoB_MinDaughtersNdf', 'int') #
  316. elenaTree, DaughtersMinP = tools.AddVect(elenaTree, 'DaughtersMinP', 'float')
  317. elenaTree, DaughtersMinPt = tools.AddVect(elenaTree, 'DaughtersMinPt', 'float')
  318. elenaTree, P = tools.AddVect(elenaTree, 'P', 'float')
  319. elenaTree, Pt = tools.AddVect(elenaTree, 'Pt', 'float')
  320. elenaTree, Mass = tools.AddVect(elenaTree, 'Mass', 'float')
  321. elenaTree, NoB_DaughtersMinP = tools.AddVect(elenaTree, 'NoB_DaughtersMinP', 'float')
  322. elenaTree, NoB_DaughtersMinPt = tools.AddVect(elenaTree, 'NoB_DaughtersMinPt', 'float')
  323. elenaTree, NoB_P = tools.AddVect(elenaTree, 'NoB_P', 'float')
  324. elenaTree, NoB_Pt = tools.AddVect(elenaTree, 'NoB_Pt', 'float')
  325. elenaTree, NoB_Mass = tools.AddVect(elenaTree, 'NoB_Mass', 'float')
  326. # Add liquid scintillator segmentation
  327. tools.makeLSsegments(nodes, elenaTree)
  328.  
  329. nodes = None
  330. def loadNodes(fGeo):
  331. global nodes
  332. nodes = searchForNodes3_xyz_dict(fGeo)
  333.  
  334. num_bad_z = 0
  335.  
  336. def signalNormalisationZ(tree, datatype, verbose):
  337. # By event
  338. # Uses MC truth!!
  339. global BadTruthVtx, num_bad_z
  340. tools.PutToZero(BadTruthVtx)
  341. z_hnl_vtx = findHNLvertex(tree)
  342. bad_z = False
  343. if not z_hnl_vtx:
  344. if "sig" in datatype:
  345. print 'ERROR: hnl vertex not found!'
  346. ii = 0
  347. for g in tree.MCTrack:
  348. ii +=1
  349. if ("sig" in datatype) and ii < 3:
  350. pass
  351. elif ("sig" in datatype) and ii >= 3:
  352. sys.exit()
  353. if not (nodes['Veto_5']['z']['pos']-nodes['Veto_5']['z']['dim']-500. < z_hnl_vtx < nodes['Tr4_4']['z']['pos']+nodes['Tr4_4']['z']['dim']):
  354. bad_z = True
  355. num_bad_z += 1
  356. if "sig" in datatype:
  357. print z_hnl_vtx
  358. tools.Push(BadTruthVtx, int(bad_z))
  359.  
  360. def nParticles(tree):
  361. # By event
  362. global NParticles
  363. np = 0
  364. for HNL in tree.Particles:
  365. np += 1
  366. tools.PutToZero(NParticles); tools.Push(NParticles, np)
  367.  
  368. def hasEcalAndMuons(tree, particle):
  369. # By particle
  370. global Has1Muon1, Has1Muon2, Has2Muon1
  371. global Has2Muon2, HasEcal
  372. flag2Muon1 = False
  373. flag2Muon2 = False
  374. flag1Muon1 = False
  375. flag1Muon2 = False
  376. flagEcal = False
  377. t1,t2 = tree.fitTrack2MC[particle.GetDaughter(0)], tree.fitTrack2MC[particle.GetDaughter(1)]
  378. # AND or OR?
  379. if ( has_muon_station(tree, t1, 1) and has_muon_station(tree, t2, 1) ):
  380. flag2Muon1 = True
  381. if ( has_muon_station(tree, t1, 2) and has_muon_station(tree, t2, 2) ):
  382. flag2Muon2 = True
  383. if ( has_muon_station(tree, t1, 1) or has_muon_station(tree, t2, 1) ):
  384. flag1Muon1 = True
  385. if ( has_muon_station(tree, t1, 2) or has_muon_station(tree, t2, 2) ):
  386. flag1Muon2 = True
  387. # This also work, but may be slower
  388. #muons1 = hasMuons(tree, t1)
  389. #muons2 = hasMuons(tree, t2)
  390. #if muons1[0] or muons2[0]: flag1Muon1 = True
  391. #if muons1[1] or muons2[1]: flag1Muon2 = True
  392. #if muons1[0] and muons2[0]: flag2Muon1 = True
  393. #if muons1[1] and muons2[1]: flag2Muon2 = True
  394. if ( hasEcalDeposit(tree, t1, 150.*u.MeV) or hasEcalDeposit(tree, t2, 150.*u.MeV) ):
  395. flagEcal = True
  396. tools.Push(Has2Muon1, int(flag2Muon1))
  397. tools.Push(Has2Muon2, int(flag2Muon2))
  398. tools.Push(Has1Muon1, int(flag1Muon1))
  399. tools.Push(Has1Muon2, int(flag1Muon2))
  400. tools.Push(HasEcal, int(flagEcal))
  401.  
  402. def chi2Ndf(tree, particle, ntr, nref):
  403. # By particle
  404. global MaxDaughtersRedChi2, MinDaughtersNdf
  405. t1,t2 = particle.GetDaughter(0),particle.GetDaughter(1)
  406. if ntr>1 and nref==2:#nf>1
  407. t1r,t2r = sh.getReFitTrIDs()[0], sh.getReFitTrIDs()[1]
  408. chi2red_1 = sh.getReFitChi2Ndf(t1r)
  409. ndf_1 = int(round(sh.getReFitNdf(t1r)))
  410. chi2red_2 = sh.getReFitChi2Ndf(t2r)
  411. ndf_2 = int(round(sh.getReFitNdf(t2r)))
  412. reducedChi2 = [chi2red_1, chi2red_2]
  413. ndfs = [ndf_1, ndf_2]
  414. # if the refit didn't work
  415. if (ntr<2) or (nref!=2) or (not ndf_1) or (not ndf_2) or (not chi2red_1) or (not chi2red_2):
  416. reducedChi2 = []
  417. ndfs = []
  418. for tr in [t1,t2]:
  419. x = tree.FitTracks[tr]
  420. ndfs.append( int(round(x.getFitStatus().getNdf())) )
  421. reducedChi2.append( x.getFitStatus().getChi2()/x.getFitStatus().getNdf() )
  422. tools.Push(MaxDaughtersRedChi2, max(reducedChi2))
  423. tools.Push(MinDaughtersNdf, min(ndfs))
  424.  
  425. def NoB_chi2Ndf(tree, particle):
  426. # By particle
  427. global NoB_MaxDaughtersRedChi2, NoB_MinDaughtersNdf, DaughtersFitConverged
  428. t1,t2 = particle.GetDaughter(0),particle.GetDaughter(1)
  429. reducedChi2 = []
  430. ndfs = []
  431. converged = []
  432. for tr in [t1,t2]:
  433. x = tree.FitTracks[tr]
  434. ndfs.append( int(round(x.getFitStatus().getNdf())) )
  435. reducedChi2.append( x.getFitStatus().getChi2()/x.getFitStatus().getNdf() )
  436. converged.append( x.getFitStatus().isFitConverged() )
  437. tools.Push(NoB_MaxDaughtersRedChi2, max(reducedChi2))
  438. tools.Push(NoB_MinDaughtersNdf, min(ndfs))
  439. tools.Push( DaughtersFitConverged, int(converged[0]*converged[1]) )
  440.  
  441. def NoB_kinematics(tree, particle):
  442. global NoB_DaughtersMinP, NoB_DaughtersMinPt, NoB_P, NoB_Pt, NoB_Mass
  443. t1,t2 = particle.GetDaughter(0),particle.GetDaughter(1)
  444. dp, dpt = [], []
  445. for tr in [t1, t2]:
  446. x = tree.FitTracks[tr]
  447. xx = x.getFittedState()
  448. dp.append(xx.getMom().Mag()); dpt.append(xx.getMom().Pt())
  449. tools.Push(NoB_DaughtersMinP, min(dp))
  450. tools.Push(NoB_DaughtersMinPt, min(dpt))
  451. HNLMom = ROOT.TLorentzVector()
  452. particle.Momentum(HNLMom)
  453. tools.Push(NoB_Mass, HNLMom.M())
  454. tools.Push(NoB_Pt, HNLMom.Pt())
  455. tools.Push(NoB_P, HNLMom.P())
  456. def goodBehavedTracks(tree, particle):
  457. # By particle
  458. # Uses MC truth!!
  459. global DaughtersAlwaysIn
  460. t1,t2 = particle.GetDaughter(0),particle.GetDaughter(1)
  461. accFlag = True
  462. for tr in [t1,t2]:
  463. mctrid = tree.fitTrack2MC[tr]
  464. if not hasGoodStrawStations(tree, mctrid):
  465. accFlag = False
  466. tools.Push(DaughtersAlwaysIn, int(accFlag))
  467.  
  468. def NoB_vertexInfo(tree, particle):
  469. # By particle
  470. global NoB_vtxx, NoB_vtxy, NoB_vtxz
  471. global NoB_IP0, NoB_DOCA
  472. HNLPos = ROOT.TLorentzVector()
  473. particle.ProductionVertex(HNLPos)
  474. xv, yv, zv, doca = HNLPos.X(),HNLPos.Y(),HNLPos.Z(),HNLPos.T()
  475. tools.Push(NoB_DOCA, doca)
  476. tools.Push(NoB_vtxx, xv*xv); tools.Push(NoB_vtxy, yv*yv); tools.Push(NoB_vtxz, zv)
  477. # impact parameter to target
  478. HNLMom = ROOT.TLorentzVector()
  479. particle.Momentum(HNLMom)
  480. tr = ROOT.TVector3(0,0,ShipGeo.target.z0)
  481. t = 0
  482. for i in range(3): t += HNLMom(i)/HNLMom.P()*(tr(i)-HNLPos(i))
  483. ip = 0
  484. for i in range(3): ip += (tr(i)-HNLPos(i)-t*HNLMom(i)/HNLMom.P())**2
  485. ip = ROOT.TMath.Sqrt(ip)
  486. tools.Push(NoB_IP0, ip)
  487. """
  488. t1,t2 = particle.GetDaughter(0),particle.GetDaughter(1)
  489. PosDir = {}
  490. for tr in [t1,t2]:
  491. xx = tree.FitTracks[tr].getFittedState()
  492. PosDir[tr] = [xx.getPos(),xx.getDir()]
  493. xv,yv,zv,doca = myVertex(t1,t2,PosDir)
  494. tools.Push(NoB_DOCA, doca)
  495. #tools.Push(NoB_vtxx, xv); tools.Push(NoB_vtxy, yv); tools.Push(NoB_vtxz, zv)
  496. tools.Push(NoB_vtxx, xv*xv); tools.Push(NoB_vtxy, yv*yv); tools.Push(NoB_vtxz, zv)
  497. # impact parameter to target
  498. HNLPos = ROOT.TLorentzVector()
  499. particle.ProductionVertex(HNLPos)
  500. HNLMom = ROOT.TLorentzVector()
  501. particle.Momentum(HNLMom)
  502. tr = ROOT.TVector3(0,0,ShipGeo.target.z0)
  503. t = 0
  504. for i in range(3): t += HNLMom(i)/HNLMom.P()*(tr(i)-HNLPos(i))
  505. ip = 0
  506. for i in range(3): ip += (tr(i)-HNLPos(i)-t*HNLMom(i)/HNLMom.P())**2
  507. ip = ROOT.TMath.Sqrt(ip)
  508. tools.Push(NoB_IP0, ip)
  509. """
  510.  
  511.  
  512. def kinematics(tree, particle, ntr, nref):
  513. global DaughtersMinP, DaughtersMinPt, P, Pt, Mass
  514. t1,t2 = particle.GetDaughter(0),particle.GetDaughter(1)
  515. dminpt, dminp = 0., 0.
  516.  
  517. if ntr>1 and nref==2:
  518. t1r,t2r = sh.getReFitTrIDs()[0], sh.getReFitTrIDs()[1]
  519. Pos1, Dir1, Mom1= sh.getReFitPosDirPval(t1r)
  520. Pos2, Dir2, Mom2= sh.getReFitPosDirPval(t2r)
  521. mass1 = pdg.GetParticle(tree.FitTracks[t1].getFittedState().getPDG()).Mass()
  522. mass2 = pdg.GetParticle(tree.FitTracks[t2].getFittedState().getPDG()).Mass()
  523. LV1 = ROOT.TLorentzVector(Mom1*Dir1, ROOT.TMath.Sqrt( mass1*mass1 + Mom1*Mom1 ))
  524. LV2 = ROOT.TLorentzVector(Mom2*Dir2, ROOT.TMath.Sqrt( mass2*mass2 + Mom2*Mom2 ))
  525. HNLMom = LV1+LV2
  526. if LV1 and LV2:
  527. dminpt = min([LV1.Pt(), LV2.Pt()])
  528. dminp = min([LV1.P(), LV2.P()])
  529.  
  530. if (ntr<2) or (nref!=2) or (not dminp) or (not dminpt) or (not HNLMom):
  531. dp, dpt = [], []
  532. for tr in [t1, t2]:
  533. x = tree.FitTracks[tr]
  534. xx = x.getFittedState()
  535. dp.append(xx.getMom().Mag()); dpt.append(xx.getMom().Pt())
  536. dminpt = min(dpt)
  537. dminp = min(dp)
  538. HNLMom = ROOT.TLorentzVector()
  539. particle.Momentum(HNLMom)
  540. tools.Push(DaughtersMinP, dminp)
  541. tools.Push(DaughtersMinPt, dminpt)
  542. tools.Push(Mass, HNLMom.M())
  543. tools.Push(Pt, HNLMom.Pt())
  544. tools.Push(P, HNLMom.P())
  545. def vertexInfo(tree, particle, ntr, nref):
  546. # By particle
  547. global vtxx, vtxy, vtxz
  548. global IP0, DOCA
  549. t1,t2 = particle.GetDaughter(0),particle.GetDaughter(1)
  550.  
  551. if ntr>1 and nref==2:#nf>1
  552. assert( len(sh.getReFitTrIDs())==2 )
  553. t1r,t2r = sh.getReFitTrIDs()[0], sh.getReFitTrIDs()[1]
  554. #print tree.fitTrack2MC[t1], t1r, tree.fitTrack2MC[t2], t2r
  555. #print ntr, nref, len(sh._StrawHits__docaEval)
  556. doca = sh.getDoca()#sh._StrawHits__docaEval[-1]#getDoca()
  557. v = sh.getReFitVertex()
  558. if v and doca:
  559. xv = v.X(); yv = v.Y(); zv = v.Z()
  560. Pos1, Dir1, Mom1= sh.getReFitPosDirPval(t1r)
  561. Pos2, Dir2, Mom2= sh.getReFitPosDirPval(t2r)
  562. mass1 = pdg.GetParticle(tree.FitTracks[t1].getFittedState().getPDG()).Mass()
  563. mass2 = pdg.GetParticle(tree.FitTracks[t2].getFittedState().getPDG()).Mass()
  564. LV1 = ROOT.TLorentzVector(Mom1*Dir1, ROOT.TMath.Sqrt( mass1*mass1 + Mom1*Mom1 ))
  565. LV2 = ROOT.TLorentzVector(Mom2*Dir2, ROOT.TMath.Sqrt( mass2*mass2 + Mom2*Mom2 ))
  566. HNLMom = LV1+LV2
  567.  
  568. # If something went wrong, take the standard values
  569. if (ntr<2) or (nref!=2) or (not v) or (not doca) or (not HNLMom):#(nf<2)
  570. PosDir = {}
  571. for tr in [t1,t2]:
  572. xx = tree.FitTracks[tr].getFittedState()
  573. PosDir[tr] = [xx.getPos(),xx.getDir()]
  574. xv,yv,zv,doca = myVertex(t1,t2,PosDir)
  575. HNLMom = ROOT.TLorentzVector()
  576. particle.Momentum(HNLMom)
  577.  
  578. tools.Push(DOCA, doca)
  579. #tools.Push(vtxx, xv); tools.Push(vtxy, yv); tools.Push(vtxz, zv)
  580. tools.Push(vtxx, xv*xv); tools.Push(vtxy, yv*yv); tools.Push(vtxz, zv)
  581.  
  582. # impact parameter to target
  583. #HNLPos = ROOT.TLorentzVector()
  584. #particle.ProductionVertex(HNLPos)
  585. HNLPos = ROOT.TVector3(xv, yv, zv)
  586. tr = ROOT.TVector3(0,0,ShipGeo.target.z0)
  587. t = 0
  588. for i in range(3): t += HNLMom(i)/HNLMom.P()*(tr(i)-HNLPos(i))
  589. ip = 0
  590. for i in range(3): ip += (tr(i)-HNLPos(i)-t*HNLMom(i)/HNLMom.P())**2
  591. ip = ROOT.TMath.Sqrt(ip)
  592. tools.Push(IP0, ip)
  593.  
  594.  
  595. def prepareFillingsByParticle():
  596. # By event
  597. global DaughtersAlwaysIn, DaughtersFitConverged, MinDaughtersNdf, MaxDaughtersRedChi2
  598. global NoB_MinDaughtersNdf, NoB_MaxDaughtersRedChi2
  599. global Has1Muon1, Has1Muon2, Has2Muon1, Has2Muon2, HasEcal
  600. global vtxx, vtxy, vtxz, IP0, DOCA
  601. global NoB_vtxx, NoB_vtxy, NoB_vtxz, NoB_IP0, NoB_DOCA
  602. global DaughtersMinP, DaughtersMinPt, Mass, P, Pt
  603. global NoB_DaughtersMinP, NoB_DaughtersMinPt, NoB_Mass, NoB_P, NoB_Pt
  604. tools.PutToZero(DaughtersAlwaysIn)
  605. tools.PutToZero(Has2Muon1); tools.PutToZero(Has2Muon2); tools.PutToZero(HasEcal)
  606. tools.PutToZero(Has1Muon1); tools.PutToZero(Has1Muon2)
  607. tools.PutToZero(DOCA)
  608. tools.PutToZero(vtxx); tools.PutToZero(vtxy); tools.PutToZero(vtxz)
  609. tools.PutToZero(IP0)
  610. tools.PutToZero(NoB_DOCA)
  611. tools.PutToZero(NoB_vtxx); tools.PutToZero(NoB_vtxy); tools.PutToZero(NoB_vtxz)
  612. tools.PutToZero(NoB_IP0)
  613. tools.PutToZero(MinDaughtersNdf); tools.PutToZero(MaxDaughtersRedChi2)
  614. tools.PutToZero(NoB_MinDaughtersNdf); tools.PutToZero(NoB_MaxDaughtersRedChi2)
  615. tools.PutToZero(DaughtersFitConverged)
  616. tools.PutToZero(DaughtersMinP); tools.PutToZero(DaughtersMinPt)
  617. tools.PutToZero(P); tools.PutToZero(Pt); tools.PutToZero(Mass)
  618. tools.PutToZero(NoB_DaughtersMinP); tools.PutToZero(NoB_DaughtersMinPt)
  619. tools.PutToZero(NoB_P); tools.PutToZero(NoB_Pt); tools.PutToZero(NoB_Mass)
  620. ntr = sh.readEvent()
  621. nref = 0
  622. if ntr>1:
  623. nref = sh.FitTracks()
  624. #print ntr, nref
  625. return ntr, nref
  626.  
  627.  
  628. def pushOfflineByEvent(tree, vetoPoints, datatype, verbose, threshold):
  629. # True HNL decay vertex (only for signal normalisation)
  630. signalNormalisationZ(tree, datatype, verbose)
  631. ## Number of particles
  632. #nParticles(tree)
  633. # Empties arrays filled by particle
  634. ntr, nref = prepareFillingsByParticle()
  635. # Liquid scintillator segments
  636. global nodes
  637. tools.hitSegments(vetoPoints, nodes, threshold)
  638. return ntr, nref
  639.  
  640. def pushOfflineByParticle(tree, particle, ntr, nref):
  641. hasEcalAndMuons(tree, particle)
  642. goodBehavedTracks(tree, particle)
  643. NoB_chi2Ndf(tree, particle)
  644. chi2Ndf(tree, particle, ntr, nref)
  645. NoB_vertexInfo(tree, particle)
  646. vertexInfo(tree, particle, ntr, nref)
  647. NoB_kinematics(tree, particle)
  648. kinematics(tree, particle, ntr, nref)
  649.  
  650. fM, tgeom, gMan, geoMat, matEff, modules, run = None, None, None, None, None, None, None
  651.  
  652. def initBField(fileNameGeo):
  653. global fM, tgeom, gMan, geoMat, matEff, modules, run, sh
  654. run = ROOT.FairRunSim()
  655. modules = shipDet_conf.configure(run,ShipGeo)
  656. tgeom = ROOT.TGeoManager("Geometry", "Geane geometry")
  657. gMan = tgeom.Import(fileNameGeo)
  658. geoMat = ROOT.genfit.TGeoMaterialInterface()
  659. matEff = ROOT.genfit.MaterialEffects.getInstance()
  660. matEff.init(geoMat)
  661. bfield = ROOT.genfit.BellField(ShipGeo.Bfield.max, ShipGeo.Bfield.z, 2, ShipGeo.Yheight/2.)
  662. fM = ROOT.genfit.FieldManager.getInstance()
  663. fM.init(bfield)
  664.  
  665. pdg, sh = None, None