Belle II Software  release-05-01-25
trainingPreparation.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
17 
18 
19 from basf2 import *
20 from setup_modules import setup_RTCtoSPTCConverters
21 import argparse
22 import os
23 
24 # ---------------------------------------------------------------------------------------
25 # Argument parser to enable training sample selection via comandline option.
26 parser = argparse.ArgumentParser(description='Training sample preparation:\
27  Prepare a data sample to train the sector map.\n\
28  Usage: basf2 traininPreparation.py -i dataSample.root -- --enable_selection boolean')
29 parser.add_argument(
30  '--enable_selection',
31  dest='use_NoKick',
32  action='store_const',
33  const=True,
34  default=False,
35  help='enable the selection of training sample based on track parameters')
36 
37 parser.add_argument(
38  '--disable_checkFit',
39  dest='checkFit',
40  action='store_const',
41  const=False,
42  default=True,
43  help="By default only RecoTracks with valid fit are taken for training. Using this option will disable that. ")
44 
45 arguments = parser.parse_args()
46 use_noKick = arguments.use_NoKick
47 
48 
49 # ---------------------------------------------------------------------------------------
50 # Settings
51 
52 # Logging and Debug Level
53 # TODO: Remove logLevel, as it can be set via basf2 option -l
54 set_log_level(LogLevel.ERROR)
55 log_to_file('logVXDTF2Preparation.log', append=False)
56 # if false PXD hits will be ignored in the trainings data collection
57 # Currently we dont do PXD tracking with vxdtf2 (as of 15.02.2018)
58 usePXD = False
59 
60 # ---------------------------------------------------------------------------------------
61 # Create paths
62 path = create_path()
63 
64 
65 # Input Module
66 rootInputM = register_module('RootInput')
67 path.add_module(rootInputM)
68 
69 # Event Info Module
70 eventinfoprinter = register_module('EventInfoPrinter')
71 path.add_module(eventinfoprinter)
72 
73 path.add_module("PrintCollections", printForEvent=1)
74 
75 
76 # puts the geometry and gearbox in the path
77 gearbox = register_module('Gearbox')
78 path.add_module(gearbox)
79 # the geometry is loaded from the DB by default now! The correct geometry
80 # should be pickked according to exp number for the generated events.
81 geometry = register_module('Geometry')
82 path.add_module(geometry)
83 
84 # Event counter
85 # eventCounter = register_module('EventCounter')
86 # path.add_module(eventCounter)
87 
88 
89 # put PXD and SVD SpacePoints into the same StoreArray
90 if usePXD:
91  spCreatorPXD = register_module('PXDSpacePointCreator')
92  spCreatorPXD.param('NameOfInstance', 'PXDSpacePointCreator')
93  spCreatorPXD.param('SpacePoints', 'PXDSpacePoints')
94  path.add_module(spCreatorPXD)
95 
96 spCreatorSVD = register_module('SVDSpacePointCreator')
97 spCreatorSVD.param('OnlySingleClusterSpacePoints', False)
98 spCreatorSVD.param('NameOfInstance', 'SVDSpacePointCreator')
99 spCreatorSVD.param('SpacePoints', 'SVDSpacePoints')
100 path.add_module(spCreatorSVD)
101 
102 
103 # Converts GenFit track candidates and checks them, with respect to the SecMap settings
104 # Produces SpacePoint TrackCand which is used in VXDTFTrainingDataCollector.
105 setup_RTCtoSPTCConverters(path=path,
106  SVDSPscollection='SVDSpacePoints',
107  PXDSPscollection='PXDSpacePoints',
108  RTCinput='MCRecoTracks',
109  sptcOutput='checkedSPTCs',
110  usePXD=usePXD,
111  logLevel=LogLevel.ERROR,
112  useNoKick=use_noKick,
113  useOnlyFittedTracks=True) # train on fitted tracks only
114 
115 
116 # SecMap BootStrap
117 # Module to fetch SecMap Config and store or load SecMap Training.
118 # Config is defined in /tracking/modules/vxdtfRedesing/src/SectorMapBootstrapModule.cc
119 # and must be available for the training of the SecMap
120 # Double False only fetches config.
121 secMapBootStrap = register_module('SectorMapBootstrap')
122 secMapBootStrap.param('ReadSectorMap', False)
123 secMapBootStrap.param('WriteSectorMap', False)
124 path.add_module(secMapBootStrap)
125 
126 
127 # Module for generation of train sample for SecMap Training
128 nameTag = 'Belle2'
129 if os.environ.get('USE_BEAST2_GEOMETRY'):
130  nameTag = 'Beast2'
131 
132 if usePXD:
133  nameTag += '_VXD'
134 else:
135  nameTag += '_SVDOnly'
136 
137 #
138 SecMapTrainerBase = register_module('VXDTFTrainingDataCollector')
139 SecMapTrainerBase.param('NameTag', nameTag)
140 SecMapTrainerBase.param('SpacePointTrackCandsName', 'checkedSPTCs')
141 # SecMapTrainerBase.logging.log_level = LogLevel.DEBUG
142 path.add_module(SecMapTrainerBase)
143 
144 # this can take quite long so it is good to know if it is still running
145 path.add_module('Progress')
146 
147 path.add_module("PrintCollections", printForEvent=1)
148 
149 process(path)
150 
151 # to show the settings of all modules (only those differing from default)
152 print_path(path)
153 
154 print(statistics)
tracking.validation.run.TrackingValidationRun.create_path
def create_path(self)
Definition: run.py:114