Belle II Software  release-08-01-10
testVXDTF2.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
11 
12 
21 
22 
23 import basf2 as b2
24 import argparse
25 # Import custom module chain for VXDTF2
26 from setup_modules import setup_VXDTF2
27 import sys
28 
29 
30 # ---------------------------------------------------------------------------------------
31 # Argument parser for input of trained Sector Map.
32 arg_parser = argparse.ArgumentParser(description='VXDTF2 application:\
33  Applies trained SecMap on provided data and returns RecoTrackCands.\n\
34  Usage: basf2 testVXDTF2.py -i <inputFileName> -o <outputFileName> -- --secmap <secmapFile>')
35 
36 arg_parser.add_argument('--secmap', '-s', type=str,
37  help='Inclusion of the root file containing the trained SecMap for the application of the VXDTF2.')
38 
39 arguments = arg_parser.parse_args(sys.argv[1:])
40 secmap_name = arguments.secmap
41 
42 
43 # ---------------------------------------------------------------------------------------
44 # Settings
45 usePXD = False
46 useDisplay = False
47 
48 performFit = False
49 generateTimeSeedAfterFit = False
50 
51 # Logging and Debug Levels
52 b2.set_log_level(b2.LogLevel.ERROR)
53 b2.log_to_file('logVXDTF2Execution.log', append=False)
54 
55 
56 # ---------------------------------------------------------------------------------------
57 path = b2.create_path()
58 
59 # Input
60 rootInput = b2.register_module('RootInput')
61 path.add_module(rootInput)
62 
63 # Event Info Module
64 eventinfoprinter = b2.register_module('EventInfoPrinter')
65 path.add_module(eventinfoprinter)
66 
67 # Gearbox
68 gearbox = b2.register_module('Gearbox')
69 path.add_module(gearbox)
70 
71 # Geometry
72 geometry = b2.register_module('Geometry')
73 geometry.param('components', ['BeamPipe',
74  'MagneticFieldConstant4LimitedRSVD',
75  'PXD',
76  'SVD'])
77 path.add_module(geometry)
78 
79 
80 # VXDTF2: Including actual VXDTF2 Modul Chain
81 setup_VXDTF2(path=path,
82  use_pxd=usePXD,
83  sec_map_file=secmap_name,
84  overlap_filter='hopfield',
85  quality_estimator='circleFit')
86 
87 
88 if performFit:
89  # This is required for RecoFitter
90  genFitExtrapolation = b2.register_module('SetupGenfitExtrapolation')
91  path.add_module(genFitExtrapolation)
92 
93  if not generateTimeSeedAfterFit:
94  timeSeed = b2.register_module('IPTrackTimeEstimator')
95  timeSeed.param('useFittedInformation', False)
96  path.add_module(timeSeed)
97 
98  fitter = b2.register_module('DAFRecoFitter')
99  path.add_module(fitter)
100 
101  if generateTimeSeedAfterFit:
102  timeSeedAfterFit = b2.register_module('IPTrackTimeEstimator')
103  timeSeedAfterFit.param('useFittedInformation', True)
104  path.add_module(timeSeedAfterFit)
105 
106 
107 output = b2.register_module('RootOutput')
108 path.add_module(output)
109 
110 
111 if useDisplay:
112  display = b2.register_module('Display')
113  display.param('showAllPrimaries', True)
114  path.add_module(display)
115 
116 path.add_module('Progress')
117 b2.process(path)
118 print(b2.statistics)