Belle II Software development
display.py
1#!/usr/bin/env python3
2
3
10
11from tracking.run.tracked_event_generation import ReadOrGenerateTrackedEventsRun
12import trackfindingcdc.cdcdisplay as cdcdisplay
13import tempfile
14
15# @cond internal_test
16
17
18class CDCDisplayRun(ReadOrGenerateTrackedEventsRun):
19 """Prepare and execute a basf2 job to read generated events or generate new events then display the CDC tracks"""
20
21
22 output_folder = tempfile.gettempdir()
23
24
25 iteractive = True
26
27
28 show_all_drawoptions = False
29
30
31 filename_prefix = ""
32
33 def __init__(self):
34 """Constructor"""
35 super().__init__()
36 ## Use the CDCSVGDisplay module to draw the CDC and tracks/hits
37 self._cdc_display_module = cdcdisplay.CDCSVGDisplayModule(self.output_folder)
38
39 @property
40 def cdc_display_module(self):
41 """Get the display module"""
42 cdc_display_module = self._cdc_display_module
43 return cdc_display_module
44
45 def create_argument_parser(self, **kwds):
46 """Translate the command-lne arguments to basf2-job parameter"""
47 argument_parser = super().create_argument_parser(**kwds)
48
49 argument_parser.add_argument(
50 '-o',
51 '--output-folder',
52 dest='output_folder',
53 default=self.output_folder,
54 help='Folder where the output files are written to. If the folder does not exist create it. '
55 )
56
57 argument_parser.add_argument(
58 '--non-interactive',
59 dest='interactive',
60 action='store_false',
61 help='Run in batch mode and do not show each event.'
62 )
63
64 argument_parser.add_argument(
65 '--use-python',
66 dest='use_python',
67 action='store_true',
68 help='Swtich to activate the legacy implementation written in python'
69 )
70
71 argument_parser.add_argument(
72 "--use_time_in_filename",
73 action='store_true',
74 help='Use the current time in the names of the generated files'
75 )
76
77 argument_parser.add_argument(
78 "-pf",
79 '--filename_prefix',
80 default=self.filename_prefix,
81 help='Prefix to the names of the generated files'
82 )
83
84 argument_parser.add_argument(
85 '-m',
86 '--mc-tracks',
87 action='store_const',
88 dest='finder_module',
89 const='TrackFinderMCTruthRecoTracks',
90 default=self.finder_module,
91 help='Generate the mc tracks using the TrackFinderMCTruthRecoTracks. Short hand for -f TrackFinderMCTruthRecoTracks'
92 )
93
94 subparser_description = \
95 """
96Various options to configure what shall be drawn in the display.
97Note that some options are only relevant, if the cellular automaton finder in the CDC has been run before.
98"""
99 draw_argument_group = argument_parser.add_argument_group(
100 title='Draw options',
101 description=subparser_description
102 )
103
104 cdc_display_module = self.cdc_display_module
105
106 if self.show_all_drawoptions:
107 drawoptions = cdc_display_module.all_drawoptions
108 else:
109 drawoptions = cdc_display_module.drawoptions
110
111 for option in sorted(drawoptions):
112 options_flag = f"--{option.replace('_', '-')} "
113
114 draw_argument_group.add_argument(
115 options_flag,
116 dest=option,
117 action='store_true',
118 default=getattr(cdc_display_module, option)
119 )
120
121 return argument_parser
122
123 def configure(self, arguments):
124 """Configure the basf2 job script using the translated command-line arguments"""
125 super().configure(arguments)
126
127 cdc_display_module = self.cdc_display_module
128
129 cdc_display_module.output_folder = arguments.output_folder
130 cdc_display_module.interactive = arguments.interactive
131
132 cdc_display_module.use_python = arguments.use_python
133 cdc_display_module.use_cpp = not arguments.use_python
134 cdc_display_module.use_time_in_filename = arguments.use_time_in_filename
135 cdc_display_module.filename_prefix = arguments.filename_prefix
136
137 if self.show_all_drawoptions:
138 drawoptions = cdc_display_module.all_drawoptions
139 else:
140 drawoptions = cdc_display_module.drawoptions
141
142 for option in drawoptions:
143 try:
144 is_active_option = getattr(arguments, option)
145 except AttributeError:
146 continue
147 else:
148 print('Setting', option, 'to', is_active_option)
149 setattr(cdc_display_module, option, is_active_option)
150
151 def create_path(self):
152 """Create the basf2 path"""
153 main_path = super().create_path()
154 main_path.add_module(self.cdc_display_module)
155 return main_path
156
157# @endcond