6 from ROOT
import gSystem
7 gSystem.Load(
'libframework')
9 gSystem.Load(
'libtracking')
11 from ROOT
import Belle2
19 from datetime
import datetime
21 from .
import svgdrawing
22 from .svgdrawing
import attributemaps
28 Personal two dimensional event display based on scalable vector graphics
31 def __init__(self, output_folder="/tmp", interactive=True):
38 Target folder for the output
39 interactive : bool, optional
40 Switch to display each event to the user and ask to continue after each event
44 super(CDCSVGDisplayModule, self).
__init__()
263 Property that collects the various names of the draw options to a list
264 that are not related to the CDC cellular automaton track finder.
265 @return list of strings naming the different switches that can be activated.
270 'draw_superlayer_boundaries',
272 'draw_interaction_point',
273 'draw_mcparticle_id',
274 'draw_mcparticle_pdgcode',
275 'draw_mcparticle_primary',
276 'draw_mcparticle_trajectories',
280 'draw_simhit_posflag',
281 'draw_simhit_pdgcode',
282 'draw_simhit_bkgtag',
290 'draw_recotrack_matching',
291 'draw_mcrecotrack_matching',
292 'draw_recotrack_seed_trajectories',
293 'draw_recotrack_fit_trajectories',
319 if not hasattr(self, name):
320 raise NameError(
'%s is not a valid draw option. Fix the misspelling.'
328 Property that collects the all names of the draw options to a list.
329 Note that some draw options only make sense after running the CDC
330 cellular automaton track finder.
331 @return list of strings naming the different switches that can be activated.
336 draw_options = [option
for option
in self.__dict__
337 if option.startswith(
'draw_')]
338 _draw_options = [option[1:]
for option
in self.__dict__
339 if option.startswith(
'_draw_')]
341 result.extend(draw_options)
342 result.extend(_draw_options)
349 Getter for the draw option that indicates if all wires shall be drawn.
350 Since this has some performance impact in animated events the wires
351 are prevented from being drawn in this case.
359 Setter for the option to draw all wires.
366 Initialisation method of the module.
367 Creates the output folder if it is not present yet.
371 if not os.path.exists(output_folder):
372 print(
"CDCSVGDisplay.__init__ : Output folder '", output_folder,
374 answer = input(
'Create it? (y or n)')
376 os.makedirs(output_folder)
391 theCDCWireTopology = \
395 cppplotter.drawWires(theCDCWireTopology)
397 plotter.draw(theCDCWireTopology)
407 segment_relation_filter.initialize()
411 Begin run method of the module. Empty here.
418 Event method of the module. Draws the event into a new svg file.
423 cppplotter = self.prefilled_cppplotter.clone()
425 plotter = self.prefilled_plotter.clone()
427 self.file_number += 1
442 basf2.B2INFO(
"Skip empty event")
449 cppplotter.drawHits(self.cdc_hits_store_array_name,
'ZeroDriftLengthColorMap',
'ZeroDriftLengthStrokeWidthMap')
451 styleDict = {
'stroke': attributemaps.ZeroDriftLengthColorMap(),
452 'stroke-width': attributemaps.ZeroDriftLengthStrokeWidthMap()}
453 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
456 if self.draw_takenflag:
458 cppplotter.drawHits(self.cdc_hits_store_array_name,
'TakenFlagColorMap',
'')
460 styleDict = {
'stroke': attributemaps.TakenFlagColorMap(), }
461 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
464 if self.draw_mcparticle_id:
466 cppplotter.drawHits(self.cdc_hits_store_array_name,
'MCParticleColorMap',
'')
468 styleDict = {
'stroke': attributemaps.MCParticleColorMap()}
469 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
472 if self.draw_mcparticle_pdgcode:
474 cppplotter.drawHits(self.cdc_hits_store_array_name,
'MCPDGCodeColorMap',
'')
476 styleDict = {
'stroke': attributemaps.MCPDGCodeColorMap()}
477 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
480 if self.draw_mcparticle_primary:
482 cppplotter.drawHits(self.cdc_hits_store_array_name,
'MCPrimaryColorMap',
'')
484 styleDict = {
'stroke': attributemaps.MCPrimaryColorMap()}
485 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
488 if self.draw_simhits:
490 cppplotter.drawSimHits(self.cdc_hits_store_array_name,
'',
'.2')
494 simHits_related_to_hits = [hit.getRelated(
'CDCSimHits')
495 for hit
in hit_storearray]
496 styleDict = {
'stroke-width':
'0.2'}
497 plotter.draw_iterable(simHits_related_to_hits, **styleDict)
500 if self.draw_simhit_posflag:
502 cppplotter.drawHits(self.cdc_hits_store_array_name,
'PosFlagColorMap',
'')
504 styleDict = {
'stroke': attributemaps.PosFlagColorMap()}
505 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
511 cppplotter.drawHits(self.cdc_hits_store_array_name,
'RLColorMap',
'')
513 styleDict = {
'stroke': attributemaps.RLColorMap()}
514 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
517 if self.draw_simhit_tof:
519 cppplotter.drawHits(self.cdc_hits_store_array_name,
'TOFColorMap',
'')
521 styleDict = {
'stroke': attributemaps.TOFColorMap()}
522 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
525 if self.draw_simhit_pdgcode:
527 cppplotter.drawHits(self.cdc_hits_store_array_name,
"SimHitPDGCodeColorMap",
"")
529 def color_map(iHit, hit):
530 simHit = hit.getRelated(
'CDCSimHits')
531 pdgCode = simHit.getPDGCode()
533 attributemaps.MCPDGCodeColorMap.color_by_pdgcode.get(pdgCode,
537 styleDict = {
'stroke': color_map}
538 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
541 if self.draw_simhit_bkgtag:
543 cppplotter.drawHits(self.cdc_hits_store_array_name,
'BackgroundTagColorMap',
'')
545 styleDict = {
'stroke': attributemaps.BackgroundTagColorMap()}
546 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
549 if self.draw_simhit_isbkg:
551 cppplotter.drawHits(self.cdc_hits_store_array_name,
'SimHitIsBkgColorMap',
'')
553 def color_map(iHit, hit):
554 simHit = hit.getRelated(
'CDCSimHits')
555 bkgTag = simHit.getBackgroundTag()
556 color = (
'gray' if bkgTag
else 'red')
559 styleDict = {
'stroke': color_map}
560 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
566 cppplotter.drawHits(self.cdc_hits_store_array_name,
'NLoops',
'')
568 print(
'No Python-function defined')
570 if self.draw_connect_tof:
572 cppplotter.drawSimHitsConnectByToF(self.cdc_hits_store_array_name,
"black",
".2")
575 if cdchits_storearray:
576 simhits_related_to_cdchit = [cdchit.getRelated(
'CDCSimHits')
577 for cdchit
in cdchits_storearray]
580 simhits_by_mcparticle = {}
581 for simhit
in simhits_related_to_cdchit:
582 mcparticle = simhit.getRelated(
'MCParticles')
583 if not mcparticle ==
None:
584 mcTrackId = mcparticle.getArrayIndex()
585 simhits_by_mcparticle.setdefault(mcTrackId, [])
586 simhits_by_mcparticle[mcTrackId].append(simhit)
588 for simhits_for_mcparticle
in list(simhits_by_mcparticle.values()):
589 simhits_for_mcparticle.sort(key=
lambda simhit:
590 simhit.getFlightTime())
592 nSimHits = len(simhits_for_mcparticle)
593 for iSimHit
in range(nSimHits - 1):
594 fromSimHit = simhits_for_mcparticle[iSimHit]
595 toSimHit = simhits_for_mcparticle[iSimHit + 1]
597 styleDict = {
'stroke-width':
'0.2',
"stroke":
"black"}
599 fromHit = fromSimHit.getRelated(self.cdc_hits_store_array_name)
600 toHit = toSimHit.getRelated(self.cdc_hits_store_array_name)
615 plotter.draw(tangent, **styleDict)
618 if self.draw_reassigned:
620 cppplotter.drawHits(self.cdc_hits_store_array_name,
'ReassignedSecondaryMap',
'')
622 styleDict = {
'stroke': attributemaps.ReassignedSecondaryMap()}
623 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
626 if self.draw_mcsegments:
629 cppplotter.drawHits(self.cdc_hits_store_array_name,
'MCSegmentIdColorMap',
'')
631 styleDict = {
'stroke': attributemaps.MCSegmentIdColorMap()}
632 plotter.draw_storearray(self.cdc_hits_store_array_name, **styleDict)
635 if self.draw_superclusters:
637 cppplotter.drawClusters(
'CDCWireHitSuperClusterVector',
640 styleDict = {
'stroke': attributemaps.listColors}
641 plotter.draw_storevector(
'CDCWireHitSuperClusterVector', **styleDict)
644 if self.draw_clusters:
646 cppplotter.drawClusters(self.cdc_wire_hit_cluster_store_obj_name,
649 styleDict = {
'stroke': attributemaps.listColors}
650 plotter.draw_storevector(self.cdc_wire_hit_cluster_store_obj_name, **styleDict)
654 if self.draw_segments:
656 cppplotter.drawSegments(self.cdc_segment_vector_store_obj_name,
659 styleDict = {
'stroke': attributemaps.listColors}
660 plotter.draw_storevector(self.cdc_segment_vector_store_obj_name, **styleDict)
662 if self.draw_segment_mctrackids:
665 cppplotter.drawSegments(self.cdc_segment_vector_store_obj_name,
666 "SegmentMCTrackIdColorMap",
"")
668 styleDict = {
'stroke': attributemaps.SegmentMCTrackIdColorMap()}
669 plotter.draw_storevector(self.cdc_segment_vector_store_obj_name, **styleDict)
671 if self.draw_segment_fbinfos:
673 cppplotter.drawSegments(self.cdc_segment_vector_store_obj_name,
674 "SegmentFBInfoColorMap",
"")
676 styleDict = {
'stroke': attributemaps.SegmentFBInfoColorMap()}
677 plotter.draw_storevector(self.cdc_segment_vector_store_obj_name, **styleDict)
679 if self.draw_segment_firstInTrackIds:
681 cppplotter.drawSegments(self.cdc_segment_vector_store_obj_name,
682 "SegmentFirstInTrackIdColorMap",
"")
685 {
'stroke': attributemaps.SegmentFirstInTrackIdColorMap()}
686 plotter.draw_storevector(self.cdc_segment_vector_store_obj_name, **styleDict)
688 if self.draw_segment_lastInTrackIds:
690 cppplotter.drawSegments(self.cdc_segment_vector_store_obj_name,
691 "SegmentLastInTrackIdColorMap",
"")
694 {
'stroke': attributemaps.SegmentLastInTrackIdColorMap()}
695 plotter.draw_storevector(self.cdc_segment_vector_store_obj_name, **styleDict)
697 if self.draw_segment_firstNPassedSuperLayers:
700 cppplotter.drawSegments(self.cdc_segment_vector_store_obj_name,
701 "SegmentFirstNPassedSuperLayersColorMap",
"")
704 {
'stroke': attributemaps.SegmentFirstNPassedSuperLayersColorMap()}
705 plotter.draw_storevector(self.cdc_segment_vector_store_obj_name, **styleDict)
707 if self.draw_segment_lastNPassedSuperLayers:
710 cppplotter.drawSegments(self.cdc_segment_vector_store_obj_name,
711 "SegmentLastNPassedSuperLayersColorMap",
"")
714 {
'stroke': attributemaps.SegmentLastNPassedSuperLayersColorMap()}
715 plotter.draw_storevector(self.cdc_segment_vector_store_obj_name, **styleDict)
717 if self.draw_segmentpairs:
719 cppplotter.drawSegmentPairs(
'CDCSegmentPairs',
'black',
'')
721 styleDict = {
"stroke":
"black"}
722 plotter.draw_storearray(
'CDCSegmentPairs', **styleDict)
725 if self.draw_mcaxialsegmentpairs:
727 cppplotter.drawMCAxialSegmentPairs(self.cdc_segment_vector_store_obj_name,
"black",
'')
730 if segment_storevector:
731 segments = segment_storevector.obj().unwrap()
732 axial_segments = [segment
for segment
in segments
733 if segment.getStereoType() == 0]
735 mc_axial_segment_pair_segment_filter = \
737 axial_segment_pair_relations = \
739 for startSegment
in axial_segments
740 for endSegment
in axial_segments)
742 def is_good_pair(pair):
743 weight = mc_axial_segment_pair_segment_filter(pair)
744 return weight == weight
746 good_axial_segment_pair_relations = [pair
for pair
in
747 axial_segment_pair_relations
if is_good_pair(pair)]
748 styleDict = {
"stroke":
"black"}
749 plotter.draw_iterable(good_axial_segment_pair_relations, **styleDict)
751 if self.draw_mcsegmentpairs:
753 cppplotter.drawMCSegmentPairs(self.cdc_segment_vector_store_obj_name,
"black",
'')
755 print(
'No Python-function defined')
757 if self.draw_mcsegmenttriples:
759 cppplotter.drawMCSegmentTriples(self.cdc_segment_vector_store_obj_name,
'',
'')
762 if segment_storevector:
763 segments = segment_storevector.obj().unwrap()
764 axial_segments = [segment
for segment
in segments
765 if segment.getStereoType() == 0]
767 stereo_segments = [segment
for segment
in segments
768 if segment.getStereoType() != 0]
771 mc_axial_segment_pair_segment_filter = \
773 mc_segment_lookup = \
778 middleSegment, endSegment)
for startSegment
in
779 axial_segments
for middleSegment
in stereo_segments
780 for endSegment
in axial_segments)
782 def is_good_triple(triple):
783 start = triple.getStartSegment()
784 middle = triple.getMiddleSegment()
785 end = triple.getEndSegment()
788 mc_axial_segment_pair_segment_filter(triple)
790 if not pairWeight == pairWeight:
793 startToMiddleFBInfo = \
794 mc_segment_lookup.areAlignedInMCTrack(start, middle)
795 if abs(startToMiddleFBInfo) > 1:
798 middleToEndFBInfo = \
799 mc_segment_lookup.areAlignedInMCTrack(middle, end)
800 if abs(middleToEndFBInfo) > 1:
803 if startToMiddleFBInfo == middleToEndFBInfo:
808 good_segment_triples = [triple
for triple
in segment_triples
809 if is_good_triple(triple)]
811 styleDict = {
"stroke":
"black"}
812 plotter.draw_iterable(good_segment_triples, **styleDict)
815 if self.draw_tangentsegments:
817 print(
'No CPP-function defined')
819 styleDict = {
'stroke-width':
'0.2'}
820 plotter.draw_storearray(
'CDCTangentSegments', **styleDict)
823 if self.draw_axialsegmentpairs:
825 cppplotter.drawAxialSegmentPairs(
'CDCAxialSegmentPairVector',
'',
'')
827 styleDict = {
'stroke': attributemaps.listColors}
828 plotter.draw_storevector(
'CDCAxialSegmentPairVector', **styleDict)
831 if self.draw_segmenttriples:
833 cppplotter.drawSegmentTriples(
'CDCSegmentTripleVector',
'',
'')
835 styleDict = {
'stroke': attributemaps.listColors}
836 plotter.draw_storevector(
'CDCSegmentTriples', **styleDict)
841 cppplotter.drawTracks(
'CDCTrackVector',
'',
'')
843 styleDict = {
'stroke': attributemaps.listColors}
844 plotter.draw_storevector(
'CDCTrackVector', **styleDict)
847 if self.draw_wrong_rl_infos_in_tracks:
850 cppplotter.drawWrongRLHitsInTracks(
'CDCTrackVector')
852 styleDict = {
'stroke': attributemaps.WrongRLColorMap()}
857 wrapped_vector = pystoreobj.obj()
858 vector = wrapped_vector.get()
861 plotter.draw_iterable(list(track.items()), **styleDict)
863 if self.draw_wrong_rl_infos_in_segments:
866 cppplotter.drawWrongRLHitsInSegments(self.cdc_segment_vector_store_obj_name)
868 styleDict = {
'stroke': attributemaps.WrongRLColorMap()}
873 wrapped_vector = pystoreobj.obj()
874 vector = wrapped_vector.get()
877 plotter.draw_iterable(list(track.items()), **styleDict)
880 if self.draw_recotracks:
882 cppplotter.drawRecoTracks(self.reco_tracks_store_array_name,
'ListColors',
'')
884 styleDict = {
'stroke': attributemaps.listColors}
885 plotter.draw_storearray(self.reco_tracks_store_array_name, **styleDict)
888 if self.draw_mcrecotracks:
890 cppplotter.drawRecoTracks(self.mc_reco_tracks_store_array_name,
'ListColors',
'')
892 styleDict = {
'stroke': attributemaps.listColors}
893 plotter.draw_storearray(self.mc_reco_tracks_store_array_name, **styleDict)
896 if self.draw_recotrack_matching:
898 cppplotter.drawRecoTracks(self.reco_tracks_store_array_name,
'MatchingStatus',
'')
900 print(
'No Python-function defined')
903 if self.draw_mcrecotrack_matching:
905 cppplotter.drawRecoTracks(self.mc_reco_tracks_store_array_name,
'MCMatchingStatus',
'')
907 print(
'No Python-function defined')
910 if self.draw_interaction_point:
912 cppplotter.drawInteractionPoint()
914 plotter.draw_interaction_point()
917 if self.draw_superlayer_boundaries:
919 cppplotter.drawSuperLayerBoundaries()
921 plotter.draw_superlayer_boundaries()
926 cppplotter.drawOuterCDCWall(
'black')
927 cppplotter.drawInnerCDCWall(
'black')
929 styleDict = {
'stroke':
'black'}
930 plotter.draw_outer_cdc_wall(**styleDict)
931 plotter.draw_inner_cdc_wall(**styleDict)
934 if self.draw_mcparticle_trajectories:
936 cppplotter.drawMCParticleTrajectories(
"MCParticles",
'black',
'')
938 print(
"Python backend can not draw mc particles")
941 if self.draw_segment_trajectories:
943 cppplotter.drawSegmentTrajectories(self.cdc_segment_vector_store_obj_name,
947 if segment_storevector:
948 segments = segment_storevector.obj().unwrap()
950 iterTrajectories = (segment.getTrajectory2D()
for segment
in segments)
951 plotter.draw_iterable(iterTrajectories)
954 if self.draw_segmenttriple_trajectories:
955 cppplotter.drawSegmentTripleTrajectories(
"CDCSegmentTriples",
960 if segmentTriple_storearray:
961 iterSegmentTriples = iter(segmentTriple_storearray)
962 iterTrajectories = (segmentTriple.getTrajectory2D()
963 for segmentTriple
in iterSegmentTriples)
964 plotter.draw_iterable(iterTrajectories)
967 if self.draw_track_trajectories:
969 cppplotter.drawTrackTrajectories(
"CDCTrackVector",
972 styleDict = {
'stroke': attributemaps.listColors}
974 if track_storevector:
975 tracks = track_storevector.obj().unwrap()
976 iterTrajectories = (cdcTrack.getStartTrajectory3D().getTrajectory2D()
977 for cdcTrack
in tracks)
978 plotter.draw_iterable(iterTrajectories, **styleDict)
981 if self.draw_recotrack_seed_trajectories:
984 if recotrack_storearray:
985 def color_map(iTrajectory, trajectory):
987 return attributemaps.listColors[iTrajectory
988 % len(attributemaps.listColors)]
990 styleDict = {
'stroke': color_map}
993 for recotrack
in recotrack_storearray:
994 tMomentum = recotrack.getMomentumSeed()
995 charge = recotrack.getChargeSeed()
996 tPosition = recotrack.getPositionSeed()
997 time = recotrack.getTimeSeed()
1007 trajectories.append(trajectory)
1009 plotter.draw_iterable(trajectories, **styleDict)
1012 raise NotImplementedError
1014 if self.draw_recotrack_fit_trajectories:
1016 cppplotter.drawRecoTrackTrajectories(self.reco_tracks_store_array_name,
'',
'')
1019 raise NotImplementedError
1021 fileName = self.new_output_filename()
1022 cppfileName = self.new_output_filename()
1025 cppplotter.saveFile(cppfileName)
1027 plotter.saveSVGFile(fileName)
1029 if self.interactive:
1031 print(
" Use the 'display' command to show the svg file", fileName,
1032 'generated for the last event')
1034 print(
" Use the 'display' command to show the svg file", cppfileName,
1035 'generated for the last event with cpp')
1039 procDisplay = subprocess.Popen([
'eog', fileName])
1041 procDisplay = subprocess.Popen([
'eog', cppfileName])
1045 input(
'Hit enter for next event')
1049 endRun methode of the module. Empty here.
1056 teminate methode of the module. Empty here.
1063 Generates a new unique name for the current event without the folder prefix
1066 if self.use_time_in_filename:
1067 output_basename = datetime.now().isoformat() +
'.svg'
1069 output_basename = self.filename_prefix + str(self.file_number).zfill(4) +
'.svg'
1070 return output_basename
1074 Generates a new unique name for the current event with the folder prefix