Belle II Software  release-08-01-10
trainSecMap.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
11 
12 
31 
32 import basf2 as b2
33 import sys
34 import argparse
35 import os
36 
37 
38 # ---------------------------------------------------------------------------------------
39 # Argument parser for input of training sample file via comandline option.
40 arg_parser = argparse.ArgumentParser(description='Sector Map Training:\
41  Trains and stores SecMap from provided data sample.\n\
42  Usage: basf2 trainSecMap.py -- --train_sample traindata.root --secmap trainedSecMap.root')
43 
44 arg_parser.add_argument('--train_sample', '-i', type=str, nargs='*',
45  help='List of prepared training data file names which will be used for the training of the SecMap')
46 arg_parser.add_argument('--secmap', '-s', type=str,
47  help='Inclusion of the root file containing the trained SecMap for the application of the VXDTF2.')
48 arg_parser.add_argument(
49  '--threshold',
50  '-t',
51  type=int,
52  default=70,
53  help='Relative threshold (in %) used to prune the sector maps. Will remove X % of the least used subgraphs.')
54 
55 arguments = arg_parser.parse_args(sys.argv[1:])
56 train_data = arguments.train_sample
57 secmap_name = arguments.secmap
58 relThreshold = arguments.threshold
59 assert len(train_data) > 0, 'No data sample for training provided!'
60 
61 # ---------------------------------------------------------------------------------------
62 # Logging and Debug Level
63 b2.set_log_level(b2.LogLevel.ERROR)
64 b2.log_to_file('logVXDTF2Training.log', append=False)
65 
66 
67 # ---------------------------------------------------------------------------------------
68 path = b2.create_path()
69 
70 # Event Info Setter which is requiered as the SecMap Training is performed in the Initilize
71 # Phase. Thus, this script does not run over any events at all, but over the data available
72 # in the prepared training sample root file.
73 eventinfosetter = b2.register_module('EventInfoSetter')
74 # default phase3 geometry:
75 exp_number = 0
76 # if environment variable is set then phase2 (aka Beast2) geometry will be taken
77 if os.environ.get('USE_BEAST2_GEOMETRY'):
78  exp_number = 1002
79 eventinfosetter.param("expList", [exp_number])
80 path.add_module(eventinfosetter)
81 
82 # puts the geometry and gearbox in the path
83 gearbox = b2.register_module('Gearbox')
84 path.add_module(gearbox)
85 # the geometry is loaded from the DB by default now! The correct geometry
86 # should be pickked according to exp number
87 geometry = b2.register_module('Geometry')
88 path.add_module(geometry)
89 
90 
91 # SecMapBootStrap Module is requiered, as it holds all the sector maps and
92 # for storing the trained SecMap.
93 secMapBootStrap = b2.register_module('SectorMapBootstrap')
94 secMapBootStrap.param('ReadSectorMap', False)
95 secMapBootStrap.param('WriteSectorMap', True)
96 if secmap_name:
97  secMapBootStrap.param('SectorMapsOutputFile', secmap_name)
98 elif os.environ.get('USE_BEAST2_GEOMETRY'):
99  secMapBootStrap.param('SectorMapsOutputFile', 'SectorMaps_Beast2.root')
100 path.add_module(secMapBootStrap)
101 
102 # Perform SecMap Training on provided data sample
103 merger = b2.register_module('RawSecMapMerger')
104 merger.param('rootFileNames', train_data)
105 merger.param('threshold', relThreshold)
106 path.add_module(merger)
107 
108 b2.print_path(path)
109 b2.process(path)
110 print(b2.statistics)