trainSecMap.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
 
 
from basf2 import *
import sys
import argparse
 
 
# ---------------------------------------------------------------------------------------
# Argument parser for input of training sample file via comandline option.
arg_parser = argparse.ArgumentParser(description='Sector Map Training:\
                                     Trains and stores SecMap from provided data sample.\n\
                                     Usage: basf2 trainSecMap.py -- --train_sample traindata.root --secmap trainedSecMap.root')
 
arg_parser.add_argument('--train_sample', '-i', type=str, action='append',
                        help='List of prepared training data file names which will be used for the training of the SecMap')
arg_parser.add_argument('--secmap', '-s', type=str,
                        help='Inclusion of the root file containing the trained SecMap for the application of the VXDTF2.')
 
arguments = arg_parser.parse_args(sys.argv[1:])
train_data = arguments.train_sample
secmap_name = arguments.secmap
assert len(train_data) > 0, 'No data sample for training provided!'
 
# ---------------------------------------------------------------------------------------
# Logging and Debug Level
set_log_level(LogLevel.ERROR)
log_to_file('logVXDTF2Training.log', append=False)
 
 
# ---------------------------------------------------------------------------------------
path = create_path()
 
# Event Info Setter which is requiered as the SecMap Training is performed in the Initilize
# Phase. Thus, this script does not run over any events at all, but over the data available
# in the prepared training sample root file.
eventinfosetter = register_module('EventInfoSetter')
# default phase3 geometry:
exp_number = 0
# if environment variable is set then phase2 (aka Beast2) geometry will be taken
if os.environ.get('USE_BEAST2_GEOMETRY'):
    exp_number = 1002
eventinfosetter.param("expList", [exp_number])
path.add_module(eventinfosetter)
 
# puts the geometry and gearbox in the path
gearbox = register_module('Gearbox')
path.add_module(gearbox)
# the geometry is loaded from the DB by default now! The correct geometry
# should be pickked according to exp number
geometry = register_module('Geometry')
path.add_module(geometry)
 
 
# SecMapBootStrap Module is requiered, as it holds all  the sector maps and
# for storing the trained SecMap.
secMapBootStrap = register_module('SectorMapBootstrap')
secMapBootStrap.param('ReadSectorMap', False)
secMapBootStrap.param('WriteSectorMap', True)
if secmap_name:
    secMapBootStrap.param('SectorMapsOutputFile', secmap_name)
elif os.environ.get('USE_BEAST2_GEOMETRY'):
    secMapBootStrap.param('SectorMapsOutputFile', 'SectorMaps_Beast2.root')
path.add_module(secMapBootStrap)
 
# Perform SecMap Training on provided data sample
merger = register_module('RawSecMapMerger')
merger.param('rootFileNames', train_data)
path.add_module(merger)
 
print_path(path)
process(path)
print(statistics)