4 from basf2
import create_path, B2ERROR, B2INFO
10 def get_background_files(folder=None, output_file_info=True):
11 """ Loads the location of the background files from the environmant variable
12 BELLE2_BACKGROUND_DIR which is set on the validation server and ensures that background
13 files exist and returns the list of background files which
14 can be directly used with add_simulation() :
16 >>> add_simulation(main, bkgfiles=background.get_background_files())
18 Will fail with an assert if no background folder set or if no background file was
19 found in the set folder.
22 folder (str): A specific folder to search for background files can be given as an optional parameter
23 output_file_info (str): If true, a list of the found background files and there size will be printed
24 This is useful to understand later which background campaign has been used
28 env_name =
'BELLE2_BACKGROUND_DIR'
32 if env_name
not in os.environ:
33 raise RuntimeError(
"Environment variable {} for backgound files not set. Terminanting this script.".format(env_name))
34 folder = os.environ[env_name]
36 bg = glob.glob(folder +
'/*.root')
39 raise RuntimeError(
"No background files found in folder {} . Terminating this script.".format(folder))
41 B2INFO(
"Background files loaded from folder {}".format(folder))
47 bg_sizes = [os.path.getsize(f)
for f
in bg]
49 table_rows = [list(entry)
for entry
in zip(bg, bg_sizes)]
50 table_rows.insert(0, [
"- Background file name -",
"- file size -"])
52 pretty_print_table(table_rows, [0, 0])
57 def add_output(path, bgType, realTime, sampleType, phase=3, fileName='output.root', excludeBranches=[]):
59 A function to be used for output of BG simulation.
61 @param bgType background type, to get available types: basf2 -m BeamBkgTagSetter
62 @param realTime equivalent time of superKEKB running in [ns]
63 @param sampleType 'study' (for BG studies) or 'usual', 'PXD', 'ECL' (for BG mixer)
64 @param specify the Phase, 1 for Phase 1, 2 for Phase 2, and 3 for Physics Run or Phase 3
65 @param fileName optional file name, can be overridden by basf2 -o
68 if sampleType ==
'study':
71 elif sampleType ==
'usual' and phase == 3:
83 elif sampleType ==
'usual' and phase == 2:
103 elif sampleType ==
'usual' and phase == 1:
115 elif sampleType ==
'ECL':
117 branches = [
'ECLHits']
118 elif sampleType ==
'PXD':
120 branches = [
'PXDSimHits']
124 B2ERROR(
'add_output - invalid value of argument sampleType: %s'
128 tagSetter = path.add_module(
'BeamBkgTagSetter', backgroundType=bgType, realTime=realTime,
129 specialFor=madeFor, Phase=phase)
132 if sampleType !=
'study':
133 emptyPath = create_path()
134 tagSetter.if_false(emptyPath)
139 path.add_module(
'RootOutput', outputFileName=fileName, branchNames=branches, excludeBranchNames=excludeBranches,
140 buildIndex=
False, autoFlushSize=-500000)