11from datetime
import datetime
14import matplotlib.pyplot
as plt
17from ROOT
import Belle2
19from ROOT
import gSystem
21gSystem.Load(
'libtracking')
22gSystem.Load(
'libtracking_trackFindingCDC')
25class QuadTreePlotter(basf2.Module):
27 This Module is able to draw the content coming from a QuadTreeImplementation with debugOutput = True.
30 def __init__(self, queue):
32 Do not forget to set the ranges! Otherwise you will end up with an empty plot..
34 basf2.Module.__init__(self)
36 self.file_name_of_quad_tree_content =
"output.root"
38 self.draw_quad_tree_content =
True
53 def plot_quad_tree_content(self):
55 Draw the quad tree content coming from the root file if enabled.
60 if not self.draw_quad_tree_content:
63 input_file = ROOT.TFile(self.file_name_of_quad_tree_content)
65 hist = input_file.Get(
"histUsed")
67 xAxis = hist.GetXaxis()
68 yAxis = hist.GetYaxis()
70 x_edges = np.array([xAxis.GetBinLowEdge(iX)
for iX
in range(1, xAxis.GetNbins() + 2)])
71 y_edges = np.array([yAxis.GetBinLowEdge(iY)
for iY
in range(1, yAxis.GetNbins() + 2)])
73 arr_l = np.array([[hist.GetBinContent(iX, iY)
for iY
in range(1, yAxis.GetNbins() + 1)]
74 for iX
in range(1, xAxis.GetNbins() + 1)])
76 hist = input_file.Get(
"histUnused")
78 xAxis = hist.GetXaxis()
79 yAxis = hist.GetYaxis()
81 x_edges = np.array([xAxis.GetBinLowEdge(iX)
for iX
in range(1, xAxis.GetNbins() + 2)])
82 y_edges = np.array([yAxis.GetBinLowEdge(iY)
for iY
in range(1, yAxis.GetNbins() + 2)])
84 l2 = np.array([[hist.GetBinContent(iX, iY)
for iY
in range(1, yAxis.GetNbins() + 1)]
85 for iX
in range(1, xAxis.GetNbins() + 1)])
87 cmap = sb.cubehelix_palette(8, start=2, rot=0, dark=0, light=1, reverse=
False, as_cmap=
True)
89 plt.gca().pcolorfast(x_edges, y_edges, (arr_l + l2).T, cmap=cmap)
91 x_labels = [f
"{x:0.{int(not float(x).is_integer())}f}" if i % 4 == 0
else "" for i, x
in enumerate(x_edges)]
92 plt.xticks(x_edges, x_labels)
93 y_labels = [f
"{y:0.{int(not float(y).is_integer())}f}" if i % 4 == 0
else "" for i, y
in enumerate(y_edges)]
94 plt.yticks(y_edges, y_labels)
96 def save_and_show_file(self):
98 Save the plot to a svg and show it (maybe a png would be better?)
100 fileName = tempfile.gettempdir() +
"/" + datetime.now().isoformat() +
'.svg'
101 plt.savefig(fileName)
102 self.file_names.append(fileName)
104 def init_plotting(self):
106 Initialize the figure with the plot ranges
107 We need to implement axes labels later!
110 plt.xlim(self.range_x_min, self.range_x_max)
111 plt.ylim(self.range_y_min, self.range_y_max)
118 self.plot_quad_tree_content()
119 self.save_and_show_file()
122 """Termination signal at the end of the event processing"""
123 self.queue.put(
"quadTree", self.file_names)
126class SegmentQuadTreePlotter(QuadTreePlotter):
129 Implementation of a quad tree plotter for SegmentQuadTrees
133 draw_segment_intersection =
True and False
135 draw_segment =
True and False
137 draw_segment_averaged =
True and False
139 draw_segment_fitted =
True and False
141 draw_mc_information =
True and False
143 draw_mc_hits =
True and False
146 theta_shifted =
False
148 maximum_theta = np.pi
150 def calculateIntersectionInQuadTreePicture(self, first, second):
152 Calculate the point where the two given hits intersect
159 positionFront = first.getRecoPos2D().conformalTransformed()
160 positionBack = second.getRecoPos2D().conformalTransformed()
162 theta_cut = np.arctan2((positionBack - positionFront).x(), (positionFront - positionBack).y())
164 if self.theta_shifted:
165 while theta_cut < - self.maximum_theta / 2:
166 theta_cut += self.maximum_theta
169 theta_cut += self.maximum_theta
171 r_cut = positionFront.x() * np.cos(theta_cut) + positionFront.y() * np.sin(theta_cut)
173 return theta_cut, r_cut
175 def calculatePositionInQuadTreePicture(self, position):
177 Transform a given normal coordinate position to a legendre position (conformal transformed)
181 position: TrackFindingCDC.Vector2D
183 position = position.conformalTransformed()
185 theta = np.linspace(self.range_x_min, self.range_x_max, 100)
186 r = position.x() * np.cos(theta) + position.y() * np.sin(theta)
190 def forAllAxialSegments(self, f):
192 Loop over all segments and execute a function
199 wrapped_vector = items.obj()
200 vector = wrapped_vector.get()
202 for quad_tree_item
in vector:
203 if quad_tree_item.getStereoType() == 0:
206 def convertToQuadTreePicture(self, phi, mag, charge):
208 Convert given track parameters into a point in the legendre space
212 phi: phi of the track
214 charge: charge of the track
216 theta = phi + np.pi / 2
217 r = 1 / mag * 1.5 * 0.00299792458 * charge
218 if self.theta_shifted:
219 if theta > self.maximum_theta / 2
or theta < -self.maximum_theta / 2:
220 theta = theta % self.maximum_theta - self.maximum_theta / 2
224 if theta > self.maximum_theta
or theta < 0:
225 theta = theta % self.maximum_theta
232 Draw everything according to the given options
236 draw_segment_intersection
238 draw_segment_averaged
243 if self.theta_shifted:
245 self.range_x_min = -self.maximum_theta / 2
247 self.range_x_max = self.maximum_theta / 2
250 self.range_x_max = self.maximum_theta
253 self.range_y_min = -0.08
255 self.range_y_max = 0.08
260 if self.draw_segment_intersection:
261 map = attributemaps.SegmentMCTrackIdColorMap()
264 theta, r = self.calculateIntersectionInQuadTreePicture(segment.front(), segment.back())
265 plt.plot(theta, r, color=list(
map(0, segment)), marker=
"o")
267 self.forAllAxialSegments(f)
269 if self.draw_segment:
270 map = attributemaps.SegmentMCTrackIdColorMap()
273 theta, r = self.calculatePositionInQuadTreePicture(segment.front().getRecoPos2D())
274 plt.plot(theta, r, color=list(
map(0, segment)), marker=
"", ls=
"-")
276 self.forAllAxialSegments(f)
278 if self.draw_segment_averaged:
279 map = attributemaps.SegmentMCTrackIdColorMap()
282 middle_index = int(np.round(segment.size() / 2.0))
283 middle_point = list(segment.items())[middle_index]
284 theta_front, r_front = self.calculateIntersectionInQuadTreePicture(segment.front(), middle_point)
285 theta_back, r_back = self.calculateIntersectionInQuadTreePicture(middle_point, segment.back())
287 plt.plot([theta_front, theta_back], [r_front, r_back], color=list(
map(0, segment)), marker=
"o", ls=
"-")
289 self.forAllAxialSegments(f)
291 if self.draw_segment_fitted:
292 map = attributemaps.SegmentMCTrackIdColorMap()
296 trajectory = fitter.fit(segment)
298 theta, r = self.convertToQuadTreePicture(momentum.phi(), momentum.norm(), trajectory.getChargeSign())
299 plt.plot(theta, r, color=list(
map(0, segment)), marker=
"o")
301 self.forAllAxialSegments(f)
306 wireHits = storedWireHits.obj().get()
309 cdc_hits = [cdcHits[i]
for track
in array
for i
in track.getHitIDs()]
311 for cdcHit
in cdcHits:
312 if cdcHit
in cdc_hits:
314 wireHit = wireHits.at(bisect.bisect_left(wireHits, cdcHit))
315 theta, r = self.calculatePositionInQuadTreePicture(wireHit.getRefPos2D())
317 plt.plot(theta, r, marker=
"", color=
"black", ls=
"-", alpha=0.8)
319 if self.draw_mc_hits:
321 wireHits = storedWireHits.obj().get()
323 map = attributemaps.listColors
328 mcTrackID = track.getMcTrackId()
330 for cdcHitID
in track.getHitIDs(Belle2.Const.CDC):
331 cdcHit = cdcHits[cdcHitID]
332 wireHit = wireHits.at(bisect.bisect_left(wireHits, cdcHit))
334 theta, r = self.calculatePositionInQuadTreePicture(wireHit.getRefPos2D())
336 plt.plot(theta, r, marker=
"", color=map[mcTrackID % len(map)], ls=
"-", alpha=0.2)
338 if self.draw_mc_information:
339 map = attributemaps.listColors
343 momentum = track.getMomSeed()
346 theta, r = self.convertToQuadTreePicture(momentum.Phi(), momentum.Mag(), track.getChargeSeed())
347 mcTrackID = track.getMcTrackId()
349 plt.plot(theta, r, marker=
"o", color=
"black", ms=10)
350 plt.plot(theta, r, marker=
"o", color=map[mcTrackID % len(map)], ms=5)
352 self.save_and_show_file()
355class StereoQuadTreePlotter(QuadTreePlotter):
358 Implementation of a quad tree plotter for StereoHitAssignment
364 draw_mc_tracks =
False
366 draw_track_hits =
False
368 draw_last_track =
True
370 delete_bad_hits =
False
372 def create_trajectory_from_track(self, track):
374 Convert a genfit::TrackCand into a TrackFindingCDC.CDCTrajectory3D
378 track: genfit::TrackCand
383 position = Vector3D(track.getPosSeed())
384 momentum = Vector3D(track.getMomSeed())
385 charge = track.getChargeSeed()
387 return Trajectory3D(position, momentum, charge)
389 def create_reco_hit3D(self, cdcHit, trajectory3D, rlInfo):
391 Use a cdc hit and a trajectory to reconstruct a CDCRecoHit3D
396 trajectory3D: TrackFindingCDC.CDCTrajectory3D
397 rlInfo: RightLeftInfo ( = short)
400 wireHits = storedWireHits.obj().get()
403 wireHit = wireHits.at(bisect.bisect_left(wireHits, cdcHit))
405 if rightLeftWireHit.getStereoType() != 0:
406 recoHit3D = CDCRecoHit3D.reconstruct(rightLeftWireHit, trajectory3D.getTrajectory2D())
411 def get_plottable_line(self, recoHit3D):
413 Minim the task of the StereoQuadTree by showing the line of quadtree nodes
416 z0 = [self.range_y_min, self.range_y_max]
417 arr_l = np.array((np.array(recoHit3D.getRecoPos3D().z()) - z0) / recoHit3D.getArcLength2D())
420 def plot_hit_line(self, recoHit3D, color):
425 if recoHit3D.getStereoType() == 0:
428 arr_l, z0 = self.get_plottable_line(recoHit3D)
429 plt.plot(arr_l, z0, marker=
"", ls=
"-", alpha=0.4, color=color)
433 Draw the hit content according to the attributes
444 self.draw_quad_tree_content =
True
447 self.range_x_min = -2 - np.sqrt(3)
449 self.range_x_max = 2 + np.sqrt(3)
452 self.range_y_min = -100
454 self.range_y_max = -self.range_y_min
457 self.plot_quad_tree_content()
459 map = attributemaps.listColors
463 wrapped_vector = items.obj()
464 track_vector = wrapped_vector.get()
470 wireHits = storedWireHits.obj().get()
472 if self.draw_mc_hits:
475 for track
in mc_track_cands:
476 mcTrackID = track.getMcTrackId()
477 trajectory = self.create_trajectory_from_track(track)
479 for cdcHitID
in track.getHitIDs(Belle2.Const.CDC):
480 cdcHit = cdcHits[cdcHitID]
482 leftRecoHit3D = self.create_reco_hit3D(cdcHit, trajectory, -1)
483 rightRecoHit3D = self.create_reco_hit3D(cdcHit, trajectory, 1)
485 self.plot_hit_line(leftRecoHit3D, color=map[mcTrackID % len(map)])
486 self.plot_hit_line(rightRecoHit3D, color=map[mcTrackID % len(map)])
488 if self.draw_mc_tracks:
491 for track
in mc_track_cands:
492 mcTrackID = track.getMcTrackId()
493 trajectory = self.create_trajectory_from_track(track)
494 z0 = trajectory.getTrajectorySZ().getZ0()
496 for cdcHitID
in track.getHitIDs(Belle2.Const.CDC):
497 cdcHit = cdcHits[cdcHitID]
498 recoHit3D = self.create_reco_hit3D(cdcHit, trajectory, mcHitLookUp.getRLInfo(cdcHit))
501 arr_l = (recoHit3D.getRecoPos3D().z() - z0) / recoHit3D.getArcLength2D()
502 plt.plot(arr_l, z0, marker=
"o", color=map[mcTrackID % len(map)], ls=
"", alpha=0.2)
504 if self.draw_track_hits:
505 for id, track
in enumerate(track_vector):
506 for recoHit3D
in list(track.items()):
507 self.plot_hit_line(recoHit3D, color=map[id % len(map)])
509 if self.draw_last_track
and len(track_vector) != 0:
511 last_track = track_vector[-1]
512 trajectory = last_track.getStartTrajectory3D().getTrajectory2D()
514 for wireHit
in wireHits:
515 for rlInfo
in (-1, 1):
518 if (self.delete_bad_hits
and
519 (wireHit.getRLInfo() != mcHitLookUp.getRLInfo(wireHit.getWireHit().getHit())
or
520 not recoHit3D.isInCellZBounds())):
523 if recoHit3D
in list(last_track.items()):
524 color = map[len(track_vector) % len(map)]
526 if wireHit.getRLInfo() == 1:
530 self.plot_hit_line(recoHit3D, color)
532 plt.xlabel(
r"$\tan \ \lambda$")
534 self.save_and_show_file()
537class QueueStereoQuadTreePlotter(StereoQuadTreePlotter):
540 A wrapper around the svg drawer in the tracking package that
541 writes its output files as a list to the queue
544 def __init__(self, queue, label, *args, **kwargs):
545 """ The same as the base class, except:
550 queue: The queue to write to
551 label: The key name in the queue
557 StereoQuadTreePlotter.__init__(self, *args, **kwargs)
563 """ Overwrite the terminate to put the list to the queue"""
564 StereoQuadTreePlotter.terminate(self)
565 self.queue.put(self.label, self.file_list)
567 def save_and_show_file(self):
568 """ Overwrite the function to listen for every new filename """
570 from datetime
import datetime
571 from matplotlib
import pyplot
as plt
573 fileName = tempfile.gettempdir() +
"/" + datetime.now().isoformat() +
'.svg'
574 plt.savefig(fileName)
575 self.file_list.append(fileName)
A (simplified) python wrapper for StoreArray.
a (simplified) python wrapper for StoreObjPtr.
Interface class to the Monte Carlo information for individual hits.
Class representing an oriented hit wire including a hypotheses whether the causing track passes left ...
Class representing a three dimensional reconstructed hit.
static CDCRecoHit3D reconstruct(const CDCRecoHit2D &recoHit2D, const CDCTrajectory2D &trajectory2D)
Reconstructs the three dimensional hit from the two dimensional and the two dimensional trajectory.
static const CDCRiemannFitter & getOriginCircleFitter()
Static getter for an origin circle fitter.
Particle full three dimensional trajectory.
A two dimensional vector which is equipped with functions for correct handling of orientation relate...
A three dimensional vector.