Belle II Software development
run.py
1
8
9import argparse
10import pickle
11
12from tracking.run.event_generation import StandardEventGenerationRun
13from tracking.run.mixins import BrowseTFileOnTerminateRunMixin, PostProcessingRunMixin
14# @cond internal_test
15
16
17class HarvestingRunMixin(BrowseTFileOnTerminateRunMixin, PostProcessingRunMixin):
18 """Harvester to select crops, postprocess, and inspect"""
19
20
21 output_file_name = None
22
23 def harvesting_module(self, path=None):
24 """This virtual method must be overridden by the inheriting class"""
25 raise RuntimeError("Override the harvesting_module method")
26
27 def create_argument_parser(self, **kwds):
28 """Parse the arguments and append them to the harvester's list"""
29 argument_parser = super().create_argument_parser(**kwds)
30 harvesting_argument_group = argument_parser.add_argument_group("Harvest arguments")
31
32 harvesting_argument_group.add_argument(
33 "-o",
34 "--output",
35 dest="output_file_name",
36 default=argparse.SUPPRESS,
37 help="File name for the harvest products"
38 )
39
40 return argument_parser
41
42 def pickle_crops(self, harvesting_module, crops, **kwds):
43 """Save the raw crops as a pickle file"""
44 with open(self.output_file_name + ".pickle", "wb") as pickle_file:
45 pickle.dump(crops, pickle_file)
46
47 def unpickle_crops(self):
48 """Load the raw crops from a pickle file"""
49 with open(self.output_file_name + ".pickle", "rb") as pickle_file:
50 return pickle.load(pickle_file)
51
52 def postprocess(self):
53 """Post-process the crops"""
54 if self.postprocess_only:
55 harvesting_module = self.harvesting_module()
56 if self.output_file_name:
57 harvesting_module.output_file_name = self.output_file_name
58 try:
59 crops = self.unpickle_crops()
60 except FileNotFoundError:
61 print("Crops pickle file not found. Create it now.")
62 else:
63 harvesting_module.refine(crops)
64
65 super().postprocess()
66
67 def adjust_path(self, path):
68 """Add the harvester to the basf2 path"""
69 super().adjust_path(path)
70 harvesting_module = self.harvesting_module()
71 if self.output_file_name:
72 harvesting_module.output_file_name = self.output_file_name
73 harvesting_module.refiners.append(self.pickle_crops)
74 path.add_module(harvesting_module)
75 return path
76
77
78class HarvestingRun(HarvestingRunMixin, StandardEventGenerationRun):
79 """Harvester to generate MC events followed by crop selection, postprocessing, inspection"""
80
81# @endcond