Belle II Software  release-05-01-25
cdst_timeResoNtuple.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 # ---------------------------------------------------------------------------------------
5 # Makes ntuple w/ variable size arrays to study single photon time resolution with di-muon data
6 # by comparing photon times w.r.t leading PDF peak position.
7 # Provides also full information of a PDF peak assocoated with a detected photon.
8 # The following draw command must show a peak that represents an overall time resolution:
9 # tree->Draw("time-t0>>h(100,-1,2)", "t0<10")
10 #
11 # Usage: basf2 cdst_timeResoNtuple.py -i <cdst_file-dimuon_skim.root> [mc]
12 # ---------------------------------------------------------------------------------------
13 
14 from basf2 import *
15 from ROOT import Belle2
16 from ROOT import TH1F, TH2F, TFile
17 from ROOT import gROOT, AddressOf, gRandom
18 from ROOT import TLorentzVector
19 import math
20 import ROOT
21 import glob
22 import sys
23 
24 MC = False
25 if len(sys.argv) > 1:
26  MC = True
27  file_num = sys.argv[1]
28 
29 
30 gROOT.ProcessLine('struct TreeStruct {\
31  int run; /* run number */ \
32  int itrk; /* track counter */ \
33  float offset; /* current bunch offset */ \
34  int usedTrk; /* number of tracks used for bunch reconstruction */ \
35  int slot; /* slot ID */ \
36  float p; /* extHit momentum */ \
37  float cth; /* extHit local cos(theta) (global must be pretty much the same) */ \
38  float phi; /* extHit local phi */ \
39  float z; /* extHit local z */ \
40  float x; /* extHit local x */ \
41  float tof; /* extHit time-of-flight */ \
42  int charge; /* particle charge */ \
43  float poca_x; /* POCA x */ \
44  float poca_y; /* POCA y */ \
45  float poca_z; /* POCA z */ \
46  int hitsCDC; /* number of CDC hits */ \
47  float dEcms; /* delta Ecms for muon */ \
48  float rec_t0; /* t0 of this track determined with bunch finder */ \
49  int valid_t0; /* reconstruction status of t0 */ \
50  int nfot; /* number of photons */ \
51  int channel[1000]; /* channel */ \
52  int pixel[1000] ; /* pixel ID */ \
53  float time[1000]; /* photon time */ \
54  float timeErr[1000]; /* photon time uncertainty */ \
55  int pulseHeight[1000]; /* photon pulse height */ \
56  float pulseWidth[1000]; /* photon pulse width */ \
57  int sample[1000]; /* sample number modulo 256 */ \
58  int status[1000]; /* calibration status bits */ \
59  float t0[1000]; /* leading PDF peak: position */ \
60  float wid0[1000]; /* leading PDF peak: width w/o TTS */ \
61  float t1[1000]; /* next to leading PDF peak: position */ \
62  float t_pdf[1000]; /* associated pdf peak time */ \
63  int type[1000]; /* 0 bkg, 1 direct, 2 reflected */ \
64  int nx[1000]; /* total number of reflections in x */ \
65  int ny[1000]; /* total number of reflections in y */ \
66  int nxm[1000]; /* number of reflections in x before mirror */ \
67  int nym[1000]; /* number of reflections in y before mirror */ \
68  int nxe[1000]; /* number of reflections in x in prism */ \
69  int nye[1000]; /* number of reflections in y in prism */ \
70  int nys[1000]; /* number of reflections on slanted surface of prism */ \
71  float xd[1000]; /* unfolded x coordinate of a pixel */ \
72  float yd[1000]; /* unfolded y coordinate of a pixel */ \
73  float xm[1000]; /* unfolded x coordinate of a reconstructed point on mirror */ \
74  float kx[1000]; /* reconstructed photon direction in x at emission */ \
75  float ky[1000]; /* reconstructed photon direction in y at emission */ \
76  float alpha[1000]; /* impact angle on photo-cathode [degrees] */ \
77 };')
78 
79 int_arrays = ['channel', 'pixel', 'pulseHeight', 'sample', 'status', 'type', 'nx', 'ny',
80  'nxm', 'nym', 'nxe', 'nye', 'nys']
81 float_arrays = ['time', 'timeErr', 'pulseWidth', 't0', 'wid0', 't1', 't_pdf', 'alpha',
82  'xd', 'yd', 'xm', 'kx', 'ky']
83 
84 from ROOT import TreeStruct
85 
86 
87 class Ntuple(Module):
88  ''' Makes a flat ntuple '''
89 
90 
91  nhisto = 0
92 
93  def initialize(self):
94  ''' initialize: open root file, construct ntuple '''
95 
96  evtMetaData = Belle2.PyStoreObj('EventMetaData')
97  expNo = evtMetaData.obj().getExperiment()
98  runNo = evtMetaData.obj().getRun()
99  exp_run = '-e' + '{:0=4d}'.format(expNo) + '-r' + '{:0=5d}'.format(runNo)
100  if not MC:
101  outName = 'out_data/timeResoNtuple' + exp_run + '.root'
102  else:
103  outName = 'out_mc/timeResoNtuple' + exp_run + '-' + file_num + '.root'
104 
105 
106  self.file = ROOT.TFile(outName, 'recreate')
107 
108 
109  self.bunchOffset = TH1F("bunchOffset", "bunch offset", 100, -1.0, 1.0)
110  self.bunchOffset.SetXTitle("bunch offset [ns]")
111 
112  self.h_cth_vs_p = TH2F("cth_vs_p", "local cos #theta vs. p", 100, 0.0, 10.0,
113  100, -1.0, 1.0)
114  self.h_cth_vs_p.SetXTitle("p [GeV/c]")
115  self.h_cth_vs_p.SetYTitle("cos #theta")
116 
117  self.h_momentum = TH1F("momentum", "momentum", 100, 0.0, 10.0)
118  self.h_momentum.SetXTitle("p [GeV/c]")
119 
120  self.h_cth = TH1F("cos_theta", "local cos #theta", 100, -1.0, 1.0)
121  self.h_cth.SetXTitle("cos #theta")
122 
123  self.h_phi = TH1F("local_phi", "local phi", 100, -math.pi, math.pi)
124  self.h_phi.SetXTitle("local #phi")
125 
126  self.h_z = TH1F("local_z", "local z", 100, -140.0, 140.0)
127  self.h_z.SetXTitle("local z [cm]")
128 
129  self.h_x = TH1F("local_x", "local x", 100, -23.0, 23.0)
130  self.h_x.SetXTitle("local x [cm]")
131 
132  self.h_charge = TH1F("charge", "charge", 3, -1.5, 1.5)
133  self.h_charge.SetXTitle("charge")
134 
135  self.h_poca_xy = TH2F("poca_xy", "POCA distribution in x-y", 100, -1.0, 1.0,
136  100, -1.0, 1.0)
137  self.h_poca_xy.SetXTitle("x [cm]")
138  self.h_poca_xy.SetYTitle("y [cm]")
139 
140  self.h_poca_z = TH1F("poca_z", "POCA distribution in z", 100, -2.0, 2.0)
141  self.h_poca_z.SetXTitle("z [cm]")
142 
143  self.h_hitsCDC = TH1F("cdc_hits", "CDC hits", 100, 0.0, 100.0)
144  self.h_hitsCDC.SetXTitle("number of CDC hits")
145 
146  self.h_Ecms = TH1F("Ecms", "c.m.s. energy of muon", 300, 4.5, 6.0)
147  self.h_Ecms.SetXTitle("E_{cm} [GeV/c]")
148 
149 
150  self.tree = ROOT.TTree('tree', 'time resolution')
151 
152  self.data = TreeStruct()
153  self.data.itrk = 0
154 
155  for key in TreeStruct.__dict__.keys():
156  if '__' not in key:
157  formstring = '/F'
158  if isinstance(self.data.__getattribute__(key), int):
159  formstring = '/I'
160  if key in int_arrays:
161  formstring = '[nfot]/I'
162  elif key in float_arrays:
163  formstring = '[nfot]/F'
164  self.tree.Branch(key, AddressOf(self.data, key), key + formstring)
165 
166  def sortPeaks(self, unsortedPeaks):
167  ''' sort PDF peaks according to their positions '''
168 
169  py_list = [x for x in unsortedPeaks]
170  return sorted(py_list, key=lambda x: (x.mean))
171 
172  def pdfHistogram(self, pdf):
173  ''' make histogram of PDF peak positions for the first 20 tracks '''
174 
175  self.nhisto += 1
176  if self.nhisto > 20:
177  return
178  h = TH2F('pdf' + str(self.nhisto), 'muon PDF, itrk = ' + str(self.data.itrk),
179  512, 0.0, 512.0, 1000, 0.0, 75.0)
180  h.SetXTitle('pixelID - 1')
181  h.SetYTitle('peak positions [ns]')
182  for x in range(512):
183  peaks = pdf[x]
184  for peak in peaks:
185  h.Fill(x, peak.mean)
186  h.Write()
187 
188  def getEcms(self, tfit, chargedStable):
189  ''' returns c.m.s. energy of a particle '''
190  mom = tfit.getMomentum()
191  lorentzLab = TLorentzVector()
192  mass = chargedStable.getMass()
193  lorentzLab.SetXYZM(mom.X(), mom.Y(), mom.Z(), mass)
194  T = Belle2.PCmsLabTransform()
195  lorentzCms = T.labToCms(lorentzLab)
196  return [lorentzCms.Energy(), T.getCMSEnergy()/2]
197 
198  def trackSelected(self):
199  ''' returns true if track fulfills selection criteria '''
200 
201  if abs(self.data.dEcms) > 0.1:
202  return False
203  if abs(self.data.poca_x) > 0.2:
204  return False
205  if abs(self.data.poca_y) > 0.2:
206  return False
207  if abs(self.data.poca_z) > 0.5:
208  return False
209  return True
210 
211  def getTOF(self, track, pdg, slot):
212  ''' returns timo-of-flight of extrapolated track '''
213 
214  extHits = track.getRelationsWith('ExtHits')
215  extEnter = None
216  extExit = None
217  for extHit in extHits:
218  if abs(extHit.getPdgCode()) != pdg:
219  continue
220  if extHit.getDetectorID() != Belle2.Const.TOP:
221  continue
222  if extHit.getCopyID() != slot:
223  continue
224  if extHit.getStatus() == Belle2.EXT_ENTER:
225  extEnter = extHit
226  elif extHit.getStatus() == Belle2.EXT_EXIT:
227  extExit = extHit
228  if extEnter:
229  return (extEnter.getTOF() + extExit.getTOF()) / 2
230  return 0
231 
232  def event(self):
233  ''' event processing '''
234 
235  recBunch = Belle2.PyStoreObj('TOPRecBunch')
236  if not recBunch:
237  B2ERROR('no TOPRecBunch')
238  return
239  if not recBunch.isReconstructed():
240  return
241  self.bunchOffset.Fill(recBunch.getCurrentOffset())
242 
243  evtMetaData = Belle2.PyStoreObj('EventMetaData')
244  self.data.run = evtMetaData.getRun()
245  self.data.offset = recBunch.getCurrentOffset()
246  self.data.usedTrk = recBunch.getUsedTracks()
247 
248  for track in Belle2.PyStoreArray('Tracks'):
249  pdfs = track.getRelated('TOPPDFCollections')
250  if not pdfs:
251  continue
252  self.data.slot = pdfs.getModuleID()
253  momentum = pdfs.getAssociatedLocalMomentum()
254  position = pdfs.getAssociatedLocalHit()
255  self.data.p = momentum.Mag()
256  self.data.cth = momentum.CosTheta()
257  self.data.phi = momentum.Phi()
258  self.data.z = position.Z()
259  self.data.x = position.X()
260  self.data.tof = self.getTOF(track, 13, self.data.slot)
261  try:
262  tfit = track.getTrackFitResultWithClosestMass(Belle2.Const.muon)
263  except:
264  B2ERROR("No trackFitResult available")
265  continue
266  self.data.charge = tfit.getChargeSign()
267  pocaPosition = tfit.getPosition()
268  self.data.poca_x = pocaPosition.X()
269  self.data.poca_y = pocaPosition.Y()
270  self.data.poca_z = pocaPosition.Z()
271  self.data.hitsCDC = tfit.getHitPatternCDC().getNHits()
272  Ecms = self.getEcms(tfit, Belle2.Const.muon)
273  self.data.dEcms = Ecms[0] - Ecms[1]
274  if not self.trackSelected():
275  continue
276  try:
277  topll = track.getRelated('TOPLikelihoods')
278  extHit = topll.getRelated('ExtHits')
279  timeZero = extHit.getRelated('TOPTimeZeros')
280  self.data.rec_t0 = timeZero.getTime()
281  self.data.valid_t0 = timeZero.isValid()
282  except:
283  self.data.rec_t0 = 0
284  self.data.valid_t0 = 0
285 
286  self.h_cth_vs_p.Fill(self.data.p, self.data.cth)
287  self.h_momentum.Fill(self.data.p)
288  self.h_cth.Fill(self.data.cth)
289  self.h_phi.Fill(self.data.phi)
290  self.h_z.Fill(self.data.z)
291  self.h_x.Fill(self.data.x)
292  self.h_charge.Fill(self.data.charge)
293  self.h_poca_xy.Fill(self.data.poca_x, self.data.poca_y)
294  self.h_poca_z.Fill(self.data.poca_z)
295  self.h_hitsCDC.Fill(self.data.hitsCDC)
296  self.h_Ecms.Fill(Ecms[0])
297  try:
298  pdf = pdfs.getHypothesisPDF(13)
299  except:
300  B2ERROR("No PDF available for PDG = 13")
301  continue
302  self.data.itrk += 1
303  self.pdfHistogram(pdf)
304  x0 = position.X() - momentum.X() / momentum.Y() * position.Y() # emission
305  z0 = position.Z() - momentum.Z() / momentum.Y() * position.Y() # emission
306  self.data.nfot = 0
307  for digit in Belle2.PyStoreArray('TOPDigits'):
308  if digit.getModuleID() == self.data.slot and digit.getHitQuality() == 1:
309  k = self.data.nfot
310  if k >= 1000:
311  continue
312  self.data.nfot += 1
313  self.data.channel[k] = digit.getChannel()
314  self.data.pixel[k] = digit.getPixelID()
315  self.data.time[k] = digit.getTime()
316  self.data.timeErr[k] = digit.getTimeError()
317  self.data.pulseHeight[k] = digit.getPulseHeight()
318  self.data.pulseWidth[k] = digit.getPulseWidth()
319  self.data.sample[k] = digit.getModulo256Sample()
320  self.data.status[k] = digit.getStatus()
321  peaks = pdf[digit.getPixelID() - 1]
322  if peaks.empty():
323  self.data.t0[k] = 0
324  self.data.wid0[k] = 0
325  self.data.t1[k] = 0
326  else:
327  sorted_peaks = self.sortPeaks(peaks)
328  self.data.t0[k] = sorted_peaks[0].mean
329  self.data.wid0[k] = sorted_peaks[0].width
330  self.data.t1[k] = self.data.t0[k] + 100
331  if peaks.size() > 1:
332  self.data.t1[k] = sorted_peaks[1].mean
333  self.data.t_pdf[k] = 0
334  self.data.type[k] = 0
335  self.data.nx[k] = 0
336  self.data.ny[k] = 0
337  self.data.nxm[k] = 0
338  self.data.nym[k] = 0
339  self.data.nxe[k] = 0
340  self.data.nye[k] = 0
341  self.data.nys[k] = 0
342  self.data.xd[k] = 0
343  self.data.yd[k] = 0
344  self.data.xm[k] = 0
345  self.data.kx[k] = 0
346  self.data.ky[k] = 0
347  self.data.alpha[k] = 0
348  assocPDF = digit.getRelated('TOPAssociatedPDFs')
349  if assocPDF:
350  pik = assocPDF.getSinglePeak()
351  if pik:
352  self.data.t_pdf[k] = pik.position
353  self.data.type[k] = pik.type
354  self.data.nx[k] = pik.nx
355  self.data.ny[k] = pik.ny
356  self.data.nxm[k] = pik.nxm
357  self.data.nym[k] = pik.nym
358  self.data.nxe[k] = pik.nxe
359  self.data.nye[k] = pik.nye
360  if pik.kyd > 0:
361  self.data.nys[k] = int(pik.nye / 2)
362  else:
363  self.data.nys[k] = int((pik.nye + 1) / 2)
364  self.data.xd[k] = pik.xd
365  self.data.yd[k] = pik.yd
366  if pik.kze > 0:
367  self.data.xm[k] = x0 + pik.kxe / pik.kze * (130.0 - z0)
368  self.data.kx[k] = pik.kxe
369  self.data.ky[k] = pik.kye
370  self.data.alpha[k] = math.degrees(math.acos(abs(pik.kzd)))
371 
372  self.tree.Fill()
373 
374  def terminate(self):
375  ''' terminate: close root file '''
376 
377  self.file.cd()
378  self.file.Write()
379  self.file.Close()
380 
381 
382 # Create path
383 main = create_path()
384 
385 # Input: cdst file(s), use -i option
386 main.add_module('RootInput')
387 
388 # Initialize TOP geometry parameters (creation of Geant geometry is not needed)
389 main.add_module('TOPGeometryParInitializer')
390 
391 # Channel masking
392 main.add_module('TOPChannelMasker')
393 
394 # Make a muon PDF available at datastore
395 main.add_module('TOPPDFDebugger', pdgCodes=[13]) # default
396 
397 # Write ntuple
398 main.add_module(Ntuple())
399 
400 # Print progress
401 progress = register_module('Progress')
402 main.add_module(progress)
403 
404 # Process events
405 process(main)
406 
407 # Print call statistics
408 print(statistics)
cdst_timeResoNtuple.Ntuple.bunchOffset
bunchOffset
histogram of bunch offset
Definition: cdst_timeResoNtuple.py:109
cdst_timeResoNtuple.Ntuple.h_poca_xy
h_poca_xy
histogram of POCA x-y
Definition: cdst_timeResoNtuple.py:135
cdst_timeResoNtuple.Ntuple.h_cth_vs_p
h_cth_vs_p
histogram cos(theta) vs, momentum
Definition: cdst_timeResoNtuple.py:112
cdst_timeResoNtuple.Ntuple.h_cth
h_cth
histogram of cos(theta)
Definition: cdst_timeResoNtuple.py:120
Belle2::PyStoreObj
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:69
cdst_timeResoNtuple.Ntuple.getTOF
def getTOF(self, track, pdg, slot)
Definition: cdst_timeResoNtuple.py:211
cdst_timeResoNtuple.Ntuple.initialize
def initialize(self)
Definition: cdst_timeResoNtuple.py:93
cdst_timeResoNtuple.Ntuple.pdfHistogram
def pdfHistogram(self, pdf)
Definition: cdst_timeResoNtuple.py:172
cdst_timeResoNtuple.Ntuple.h_charge
h_charge
histogram of charge
Definition: cdst_timeResoNtuple.py:132
cdst_timeResoNtuple.Ntuple.nhisto
int nhisto
histogram counter
Definition: cdst_timeResoNtuple.py:91
cdst_timeResoNtuple.Ntuple.h_phi
h_phi
histogram of local phi
Definition: cdst_timeResoNtuple.py:123
cdst_timeResoNtuple.Ntuple.h_poca_z
h_poca_z
histogram of POCA z
Definition: cdst_timeResoNtuple.py:140
cdst_timeResoNtuple.Ntuple.data
data
data structure
Definition: cdst_timeResoNtuple.py:152
Belle2::getRun
static ExpRun getRun(map< ExpRun, pair< double, double >> runs, double t)
Get exp number + run number from time.
Definition: Splitter.cc:262
cdst_timeResoNtuple.Ntuple.getEcms
def getEcms(self, tfit, chargedStable)
Definition: cdst_timeResoNtuple.py:188
cdst_timeResoNtuple.Ntuple.trackSelected
def trackSelected(self)
Definition: cdst_timeResoNtuple.py:198
cdst_timeResoNtuple.Ntuple.file
file
file object
Definition: cdst_timeResoNtuple.py:106
cdst_timeResoNtuple.Ntuple.h_momentum
h_momentum
histogram of momentum
Definition: cdst_timeResoNtuple.py:117
cdst_timeResoNtuple.Ntuple.h_Ecms
h_Ecms
histogram of Ecms
Definition: cdst_timeResoNtuple.py:146
cdst_timeResoNtuple.Ntuple.h_z
h_z
histogram of local z
Definition: cdst_timeResoNtuple.py:126
cdst_timeResoNtuple.Ntuple.event
def event(self)
Definition: cdst_timeResoNtuple.py:232
cdst_timeResoNtuple.Ntuple.terminate
def terminate(self)
Definition: cdst_timeResoNtuple.py:374
Belle2::PyStoreArray
a (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:58
Belle2::PCmsLabTransform
Class to hold Lorentz transformations from/to CMS and boost vector.
Definition: PCmsLabTransform.h:37
cdst_timeResoNtuple.Ntuple.sortPeaks
def sortPeaks(self, unsortedPeaks)
Definition: cdst_timeResoNtuple.py:166
cdst_timeResoNtuple.Ntuple
Definition: cdst_timeResoNtuple.py:87
cdst_timeResoNtuple.Ntuple.h_hitsCDC
h_hitsCDC
histogram of number of hits in CDC
Definition: cdst_timeResoNtuple.py:143
cdst_timeResoNtuple.Ntuple.tree
tree
tree object
Definition: cdst_timeResoNtuple.py:150
cdst_timeResoNtuple.Ntuple.h_x
h_x
histogram of local x
Definition: cdst_timeResoNtuple.py:129