Belle II Software  release-05-01-25
background.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 from basf2 import create_path, B2ERROR, B2INFO
5 from basf2.utils import pretty_print_table
6 import os
7 import glob
8 
9 
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() :
15 
16  >>> add_simulation(main, bkgfiles=background.get_background_files())
17 
18  Will fail with an assert if no background folder set or if no background file was
19  found in the set folder.
20 
21  Parameters:
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
25  to simulate events.
26  """
27 
28  env_name = 'BELLE2_BACKGROUND_DIR'
29  bg = None
30 
31  if folder is None:
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]
35 
36  bg = glob.glob(folder + '/*.root')
37 
38  if len(bg) == 0:
39  raise RuntimeError("No background files found in folder {} . Terminating this script.".format(folder))
40 
41  B2INFO("Background files loaded from folder {}".format(folder))
42 
43  # sort for easier comparison
44  bg = sorted(bg)
45 
46  if output_file_info:
47  bg_sizes = [os.path.getsize(f) for f in bg]
48  # reformat to work with pretty_print_table
49  table_rows = [list(entry) for entry in zip(bg, bg_sizes)]
50  table_rows.insert(0, ["- Background file name -", "- file size -"])
51 
52  pretty_print_table(table_rows, [0, 0])
53 
54  return bg
55 
56 
57 def add_output(path, bgType, realTime, sampleType, phase=3, fileName='output.root', excludeBranches=[]):
58  '''
59  A function to be used for output of BG simulation.
60  @param path path name
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
66  '''
67 
68  if sampleType == 'study':
69  madeFor = ''
70  branches = []
71  elif sampleType == 'usual' and phase == 3:
72  madeFor = ''
73  branches = [
74  'PXDSimHits',
75  'SVDSimHits',
76  'CDCSimHits',
77  'TOPSimHits',
78  'ARICHSimHits',
79  'ECLHits',
80  'BKLMSimHits',
81  'EKLMSimHits',
82  ]
83  elif sampleType == 'usual' and phase == 2:
84  madeFor = ''
85  branches = [
86  'PXDSimHits',
87  'SVDSimHits',
88  'CDCSimHits',
89  'TOPSimHits',
90  'ARICHSimHits',
91  'ECLHits',
92  'BKLMSimHits',
93  'EKLMSimHits',
94  'CLAWSSimHits',
95  'FANGSSimHits',
96  'PlumeSimHits',
97  'BeamabortSimHits',
98  'PindiodeSimHits',
99  'QcsmonitorSimHits',
100  'He3tubeSimHits',
101  'MicrotpcSimHits',
102  ]
103  elif sampleType == 'usual' and phase == 1:
104  madeFor = ''
105  branches = [
106  'ClawSimHits',
107  'BeamabortSimHits',
108  'PindiodeSimHits',
109  'QcsmonitorSimHits',
110  'He3tubeSimHits',
111  'MicrotpcSimHits',
112  'BgoSimHits',
113  'CsiSimHits',
114  ]
115  elif sampleType == 'ECL':
116  madeFor = 'ECL'
117  branches = ['ECLHits']
118  elif sampleType == 'PXD':
119  madeFor = 'PXD'
120  branches = ['PXDSimHits']
121  else:
122  madeFor = ''
123  branches = []
124  B2ERROR('add_output - invalid value of argument sampleType: %s'
125  % sampleType)
126 
127  # Set background tag in SimHits and add BackgroundMetaData into persistent tree
128  tagSetter = path.add_module('BeamBkgTagSetter', backgroundType=bgType, realTime=realTime,
129  specialFor=madeFor, Phase=phase)
130 
131  # Write out only non-empty events when producing samples for BG mixer
132  if sampleType != 'study':
133  emptyPath = create_path()
134  tagSetter.if_false(emptyPath)
135 
136  # Output to file. We don't need a TTreeIndex for background files and memory
137  # consumption can be improved by setting a lower autoFlushSize so that
138  # fewer and or smaller amounts of data have to be read for each GetEntry()
139  path.add_module('RootOutput', outputFileName=fileName, branchNames=branches, excludeBranchNames=excludeBranches,
140  buildIndex=False, autoFlushSize=-500000)
basf2.utils
Definition: utils.py:1