Belle II Software development
record.py
1#!/usr/bin/env python3
2
3
10
11# Fix TLS bug
12# ctypes.cdll.LoadLibrary('/space/ofrost/basf2/release/lib/Linux_x86_64/opt/libtracking_trackFindingCDC.so')
13
14
15from ROOT import Belle2 # noqa: make Belle2 namespace available
16from ROOT.Belle2 import TrackFindingCDC as TFCDC
17
18import sys
19import math
20import numpy as np
21
22from tracking.validation.utilities import is_primary
23
24import tracking.harvest.harvesting as harvesting
25import tracking.harvest.refiners as refiners
26from tracking.harvest.run import HarvestingRun
27
28
29import logging
30
31
32def get_logger():
33 return logging.getLogger(__name__)
34
35
36CONTACT = "oliver.frost@desy.de"
37
38
39def getNaiveBetheEnergyLoss(p, pdg_code, s):
40 eloss = s * 0.00008 # Miriam's / Sasha's value
41 eloss = s * 1 * 10**-4 # Miriam's / Sasha's value
42 return eloss
43
44
45def getBetheStoppingPower(p, pdg_code):
46 mass = 0.139 # Hard coded pions
47 energy = np.sqrt(mass * mass + p * p)
48
49 gamma = energy / mass
50 beta = p / energy
51
52 gamma2 = gamma * gamma
53 beta2 = beta * beta
54
55 eMass = 0.511 * Belle2.Unit.MeV
56 m_I = 7.66932e-08
57 m_eDensity = 0.000515726
58
59 Wmax = 2 * eMass * beta2 * gamma2 / (1 + 2 * gamma * eMass / mass)
60 I2 = m_I * m_I
61
62 K = 0.307075 * Belle2.Unit.MeV * Belle2.Unit.cm2
63 dEdx = (
64 K * m_eDensity / beta2 *
65 (1.0 / 2.0 * math.log(2 * eMass * beta2 * gamma2 * Wmax / I2) - beta2)
66 )
67
68 m_I = 4.54e-8 # Miriam's / Sasha's value
69 # I2 = 2.06116e-15
70 I2 = m_I * m_I
71
72 factor = 178.6 * 0.00015e-04 / 2 # Miriam's / Sasha's value - to big...
73 # eDensity = factor / K
74 dEdx = factor / beta2 * (math.log(2 * eMass * beta2 * gamma2 * Wmax / I2) - beta2)
75
76 # Miriam's / Sasha's original
77 # SR = (0.00015e-04*(2*(beta)**2)**(-1))*(math.log((0.0000010404*((beta)**4)*((gamma)**4))/
78 # (2.06116e-15*(1+0.007331*gamma)))-beta**2)
79 # dEdx = 178.6 * SR
80 return dEdx
81
82
83def getBetheEnergyLoss(p, pdg_code, path):
84 dEdx = getBetheStoppingPower(p, pdg_code)
85 eloss = dEdx * path
86 return eloss
87
88
89def getMomentumLossFactor(p, pdg_code, eloss):
90 p_factor = (p - eloss) / p
91 return p_factor
92
93
94def DeltaR(path, particleID, P):
95 eloss = getBetheEnergyLoss(P, particleID, path)
96 return getMomentumLossFactor(P, particleID, eloss)
97
98
100 """Harvester to generate, postprocess and inspect MC events for energy-loss evaluation"""
101
102 n_events = 10000
103
104 generator_module = "eloss_gun"
105
106 detector_setup = "TrackingDetectorConstB"
107
108 @property
110 """Get the output ROOT filename"""
111 return 'eloss.root'
112
113 def harvesting_module(self, path=None):
114 """Harvest and post-process the generated events"""
116 if path:
117 path.add_module(harvesting_module)
118 return harvesting_module
119
120 def create_argument_parser(self, **kwds):
121 """Convert command-line arguments to basf2 argument list"""
122 argument_parser = super().create_argument_parser(**kwds)
123 return argument_parser
124
125 def create_path(self):
126 """
127 Sets up a path that plays back pregenerated events or generates events
128 based on the properties in the base class.
129 """
130 path = super().create_path()
131
132 path.add_module("TFCDC_WireHitPreparer",
133 logLevel=8,
134 flightTimeEstimation="outwards",
135 UseNLoops=1)
136
137 path.add_module('TFCDC_AxialTrackCreatorMCTruth',
138 logLevel=8,
139 useOnlyBeforeTOP=True,
140 fit=True,
141 reconstructedDriftLength=True,
142 reconstructedPositions=True)
143
144 return path
145
146
147class ElossHarvestingModule(harvesting.HarvestingModule):
148
149 """Module to collect information about the generated segments and
150 compose validation plots on terminate."""
151
152 def __init__(self, output_file_name):
153 """Constructor"""
154 super().__init__(foreach='CDCTrackVector',
155 output_file_name=output_file_name)
156
157
158 self.mc_track_lookup = None
159
160 origin_track_fitter = TFCDC.CDCRiemannFitter()
161 origin_track_fitter.setOriginConstrained()
162
163 self.track_fitter = origin_track_fitter
164
165 def initialize(self):
166 """Receive signal at the start of event processing"""
167 super().initialize()
168
169 self.mc_track_lookup = TFCDC.CDCMCTrackLookUp.getInstance()
170
171 self.mc_hit_lookup = TFCDC.CDCMCHitLookUp.getInstance()
172
173 self.eloss_estimator = TFCDC.EnergyLossEstimator.forCDC()
174
175 self.bfield = TFCDC.CDCBFieldUtil
176
177 def prepare(self):
178 """Initialize the MC-hit lookup method"""
179 TFCDC.CDCMCHitLookUp.getInstance().fill()
180
181 def pick(self, track):
182 """Select tracks with at least 4 segments and associated primary MC particle with pt >= 0.25 GeV/c"""
183 if track.size() < 4:
184 return False
185
186 mc_track_lookup = self.mc_track_lookup
187 mc_particle = mc_track_lookup.getMCParticle(track)
188
189 # Check that mc_particle is not a nullptr
190 if mc_particle is None:
191 return False
192
193 if mc_particle.getMomentum().Rho() < 0.250:
194 return False
195
196 return is_primary(mc_particle)
197
198 def peel(self, track):
199 """Aggregate the track and MC information for dE/dx analysis"""
200 mc_track_lookup = self.mc_track_lookup
201
202 # rl_drift_circle = 1
203 # unit_variance = 0
204 # observations2D = TFCDC.CDCObservations2D(rl_drift_circle, unit_variance)
205
206 # for recoHit3D in track:
207 # observations2D.append(recoHit3D)
208
209 # trajectory2D = track_fitter.fit(observations2D)
210 # trajectory2D.setLocalOrigin(TFCDC.Vector2D(0, 0))
211
212 mc_particle = mc_track_lookup.getMCParticle(track)
213 pdg_code = mc_particle.getPDG()
214 t_vertex = mc_particle.getVertex()
215 t_mom = mc_particle.getMomentum()
216 charge = mc_particle.getCharge()
217 mass = mc_particle.getMass()
218 mc_energy = mc_particle.getEnergy()
219
220 mc_vertex2D = TFCDC.Vector2D(t_vertex.XYvector())
221 mc_mom2D = TFCDC.Vector2D(t_mom.XYvector())
222 mc_trajectory2D = TFCDC.CDCTrajectory2D(mc_vertex2D, 0, mc_mom2D, charge)
223 mc_pt = mc_mom2D.norm()
224
225 first_hit = self.mc_track_lookup.getFirstHit(track)
226 first_sim_hit = first_hit.getRelated("CDCSimHits")
227 if first_sim_hit is None:
228 return
229 # Make sure we start the track in the first layer to avoid some confusion
230 if first_sim_hit.getWireID().getICLayer() != 0:
231 return
232
233 last_hit = self.mc_track_lookup.getLastHit(track)
234 last_sim_hit = last_hit.getRelated("CDCSimHits")
235 if last_sim_hit is None:
236 return
237 # Make sure we start the track in the last layer to avoid some confusion
238 # if last_sim_hit.getWireID().getICLayer() != 55: return
239 # last_sim_mom3D = TFCDC.Vector3D(last_sim_hit.getMomentum())
240
241 first_sim_pos3D = TFCDC.Vector3D(first_sim_hit.getPosTrack())
242 first_sim_mom3D = TFCDC.Vector3D(first_sim_hit.getMomentum())
243 first_sim_tof = first_sim_hit.getFlightTime()
244 first_sim_energy = np.sqrt(first_sim_mom3D.norm() ** 2 + mass ** 2)
245 first_sim_pt = first_sim_mom3D.cylindricalR()
246 first_sim_pz = first_sim_mom3D.z()
247
248 first_sim_mom2D = first_sim_mom3D.xy()
249
250 # first_sim_mom2D.normalizeTo(last_sim_mom3D.cylindricalR())
251 first_sim_trajectory2D = TFCDC.CDCTrajectory2D(first_sim_pos3D.xy(), first_sim_tof, first_sim_mom2D, charge)
252
253 # mc_trajectory3D = mc_track_lookup.getTrajectory3D(track)
254 # mc_trajectory2D = mc_trajectory3D.getTrajectory2D()
255
256 for reco_hit3D in track:
257 sim_hit = self.mc_hit_lookup.getSimHit(reco_hit3D.getWireHit().getHit())
258 if sim_hit is None:
259 continue
260
261 sim_mom3D = TFCDC.Vector3D(sim_hit.getMomentum())
262 sim_energy = np.sqrt(sim_mom3D.norm() ** 2 + mass ** 2)
263 sim_pt = sim_mom3D.cylindricalR()
264 sim_pz = sim_mom3D.z()
265
266 mc_eloss_truth = mc_energy - sim_energy
267 first_eloss_truth = first_sim_energy - sim_energy
268
269 sim_pos3D = TFCDC.Vector3D(sim_hit.getPosTrack())
270 sim_pos2D = sim_pos3D.xy()
271
272 layer_cid = reco_hit3D.getWire().getICLayer()
273 bz = self.bfield.getBFieldZ(sim_pos3D)
274 r = sim_pos3D.cylindricalR()
275
276 # recoPos3D = reco_hit3D.getRecoPos3D()
277 # recoPos3D = TFCDC.Vector3D(sim_hit.getPosTrack())
278 # recoPos2D = recoPos3D.xy()
279 # refPos2D = reco_hit3D.getRefPos2D()
280 # l = reco_hit3D.getSignedRecoDriftLength()
281
282 mc_s2D = mc_trajectory2D.calcArcLength2D(sim_pos2D)
283 first_s2D = first_sim_trajectory2D.calcArcLength2D(sim_pos2D)
284
285 if mc_s2D < 0:
286 # mc_s2D += mc_trajectory2D.getArcLength2DPeriod()
287 continue
288
289 mc_eloss_estimate = self.eloss_estimator.getEnergyLoss(mc_pt, pdg_code, mc_s2D)
290 first_eloss_estimate = self.eloss_estimator.getEnergyLoss(first_sim_pt, pdg_code, first_s2D)
291 first_ploss_factor = self.eloss_estimator.getMomentumLossFactor(first_sim_pt, pdg_code, first_s2D)
292
293 sasha_eloss_estimate = getBetheEnergyLoss(first_sim_pt, pdg_code, first_s2D)
294 sasha_ploss_factor = DeltaR(first_s2D, pdg_code, first_sim_pt)
295
296 # first_residual2D = first_trajectory2D.getDist2D(refPos2D) - l
297 first_residual2D = first_sim_trajectory2D.getDist2D(sim_pos2D)
298 first_disp2D = charge * first_residual2D
299
300 first_loss_disp2D_estimate = abs(self.eloss_estimator.getLossDist2D(first_sim_pt, pdg_code, first_s2D))
301 first_delossed_disp2D = first_disp2D - first_loss_disp2D_estimate
302
303 # mc_residual2D = mc_trajectory2D.getDist2D(refPos2D) - l
304 mc_residual2D = mc_trajectory2D.getDist2D(sim_pos2D)
305 mc_disp2D = charge * mc_residual2D
306
307 mc_loss_disp2D_estimate = abs(self.eloss_estimator.getLossDist2D(mc_pt, pdg_code, mc_s2D))
308 mc_delossed_disp2D = mc_disp2D - mc_loss_disp2D_estimate
309
310 # s2D = trajectory2D.calcArcLength2D(recoPos2D)
311 # if s2D < 0: s2D += trajectory2D.getArcLength2DPeriod()
312
313 # residual2D = trajectory2D.getDist2D(recoPos2D)
314 # disp2D = charge * residual2D
315
316 # # Original approach
317 # center2D = trajectory2D.getGlobalCircle().center()
318 # radius2D = recoPos2D - center2D;
319
320 # eLoss = self.eloss_estimator.getEnergyLoss(mc_pt, pdg_code, mc_s2D)
321 # dx = self.eloss_estimator.getMomentumLossFactor(mc_pt, pdg_code, mc_s2D)
322 # deloss_radius2D = recoPos2D - center2D;
323 # deloss_radius2D.scale (1.0 / dx)
324 # loss_disp2D_estimate2 = (deloss_radius2D - radius2D).norm()
325
326 # loss_disp2D_estimate3 = radius2D.norm() * (1.0 /dx - 1.0)
327 # loss_disp2D_estimate3 = radius2D.norm() * ((1.0 - dx) /dx)
328 # loss_disp2D_estimate3 = radius2D.norm() * (eLoss / (mc_pt - eLoss))
329
330 if abs(mc_residual2D) > 6:
331 continue
332
333 if abs(first_residual2D) > 6:
334 continue
335
336 yield dict(
337 layer_cid=layer_cid,
338 r=r,
339 bz=bz,
340 # pt=trajectory2D.getAbsMom2D(),
341 pdg_code=abs(pdg_code),
342 mass=mass,
343 charge=charge,
344 mc_pt=mc_pt,
345 first_sim_pt=first_sim_pt,
346 sim_pt=sim_pt,
347
348 diff_pt=first_sim_pt - sim_pt,
349 diff_pz=first_sim_pz - sim_pz,
350
351 mc_eloss_truth=mc_eloss_truth,
352 mc_eloss_estimate=mc_eloss_estimate,
353
354 first_eloss_truth=first_eloss_truth,
355
356 first_eloss_estimate=first_eloss_estimate,
357 sasha_eloss_estimate=sasha_eloss_estimate,
358
359 first_ploss_factor=first_ploss_factor,
360 sasha_ploss_factor=sasha_ploss_factor,
361
362 # l=l,
363
364 first_s2D=first_s2D,
365 first_residual2D=first_residual2D,
366 first_disp2D=first_disp2D,
367 first_loss_disp2D_estimate=first_loss_disp2D_estimate,
368 first_delossed_disp2D=first_delossed_disp2D,
369
370 mc_s2D=mc_s2D,
371 mc_residual2D=mc_residual2D,
372 mc_disp2D=mc_disp2D,
373 mc_loss_disp2D_estimate=mc_loss_disp2D_estimate,
374 mc_delossed_disp2D=mc_delossed_disp2D,
375
376 # s2D=s2D,
377 # disp2D=disp2D,
378 # residual2D=residual2D,
379
380 # loss_disp2D_estimate2=loss_disp2D_estimate2,
381 # loss_disp2D_estimate3=loss_disp2D_estimate3,
382 )
383
384
386 save_tree = refiners.save_tree()
387
388 save_histograms = refiners.save_histograms(
389 outlier_z_score=5.0,
390 allow_discrete=True,
391 )
392
393
394 save_histograms_stackby_charge = refiners.save_histograms(
395 select=[
396 # "mc_disp2D",
397 "first_disp2D",
398 "charge",
399 ],
400 outlier_z_score=5.0,
401 allow_discrete=True,
402 fit="gaus",
403 fit_z_score=1,
404 groupby="charge",
405 )
406
407
408 save_scatter = refiners.save_scatters(
409 x=['mc_s2D'],
410 y=[
411 # 'mc_disp2D',
412 'disp2D',
413 ],
414 groupby=[None, "charge"],
415 filter=lambda x: x == 211,
416 filter_on="pdg_code",
417 )
418
419
420 save_profiles = refiners.save_profiles(
421 x=['mc_s2D'],
422 y=[
423 # 'mc_disp2D',
424 'first_disp2D',
425 'disp2D',
426 ],
427 groupby=[None, "charge"],
428 )
429
430
431 save_bz_profiles = refiners.save_profiles(
432 x='r',
433 y='bz',
434 )
435
436 # Loss displacement #
437 # ################# #
438
439
440 save_cid_histogram = refiners.save_histograms(
441 select=[
442 # 'mc_disp2D',
443 # 'mc_delossed_disp2D',
444 'first_disp2D',
445 'first_delossed_disp2D',
446 'bz',
447 # 'disp2D',
448 ],
449 outlier_z_score=5.0,
450 groupby=["layer_cid"],
451 # stackby="pdg_code",
452 )
453
454
455 save_cid_profiles = refiners.save_profiles(
456 x=["mc_pt"],
457 y=[
458 # 'mc_disp2D',
459 # 'mc_delossed_disp2D',
460 'first_disp2D',
461 'first_delossed_disp2D',
462 ],
463 outlier_z_score=5.0,
464 groupby=["layer_cid"],
465 stackby="pdg_code",
466 )
467
468
469 save_cid_scatters = refiners.save_scatters(
470 x=["mc_pt"],
471 y=[
472 # 'mc_disp2D',
473 # 'mc_delossed_disp2D',
474 'first_disp2D',
475 'first_delossed_disp2D',
476 ],
477 outlier_z_score=5.0,
478 groupby=["layer_cid"],
479 stackby="pdg_code",
480 )
481
482 # Energy loss #
483 # ########### #
484
485 save_energy_cid_histogram = refiners.save_histograms(
486 select=[
487 'pdg_code',
488 'first_eloss_estimate',
489 'first_eloss_truth',
490 ],
491 outlier_z_score=5.0,
492 groupby=["layer_cid"],
493 stackby="pdg_code",
494 folder_name='energy/{groupby_addition}',
495 )
496
497
498 save_energy_cid_profiles = refiners.save_profiles(
499 x=["mc_pt"],
500 y=[
501 'first_eloss_truth',
502 'first_eloss_estimate',
503 ],
504 outlier_z_score=5.0,
505 groupby=["layer_cid"],
506 stackby="pdg_code",
507 folder_name='energy/{groupby_addition}',
508 )
509
510
511 save_energy_cid_scatters = refiners.save_scatters(
512 x=["mc_pt"],
513 y=[
514 'first_eloss_truth',
515 'first_eloss_estimate',
516 ],
517 outlier_z_score=5.0,
518 groupby=["layer_cid"],
519 stackby="pdg_code",
520 folder_name='energy/{groupby_addition}',
521 )
522
523
524def main():
525 run = ElossHarvestingRun()
526 run.configure_and_execute_from_commandline()
527
528
529if __name__ == "__main__":
530 logging.basicConfig(stream=sys.stdout, level=logging.INFO, format='%(levelname)s:%(message)s')
531 main()
mc_track_lookup
by default, there is no method to find matching MC tracks
Definition: record.py:158
mc_hit_lookup
Method to find matching MC hits.
Definition: record.py:171
track_fitter
Use the CDCReimannFitter with a constrained origin for track fitting.
Definition: record.py:163
def peel(self, track)
Definition: record.py:198
bfield
Method to interrogate the magnetic field values.
Definition: record.py:175
def pick(self, track)
Definition: record.py:181
eloss_estimator
Method to estimate dE/dx in the CDC.
Definition: record.py:173
def __init__(self, output_file_name)
Definition: record.py:152
def create_argument_parser(self, **kwds)
Definition: record.py:120
def harvesting_module(self, path=None)
Definition: record.py:113
def output_file_name(self)
Definition: record.py:109
None output_file_name
Disable the writing of an output ROOT file.
Definition: run.py:20
None output_file_name
There is no default for the name of the output TFile.
Definition: mixins.py:60
Definition: main.py:1