Source code for background

#!/usr/bin/env python3

##########################################################################
# basf2 (Belle II Analysis Software Framework)                           #
# Author: The Belle II Collaboration                                     #
#                                                                        #
# See git log for contributors and copyright holders.                    #
# This file is licensed under LGPL-3.0, see LICENSE.md.                  #
##########################################################################

from basf2 import create_path, B2ERROR, B2INFO
from basf2.utils import pretty_print_table
import os
import glob


[docs]def get_background_files(folder=None, output_file_info=True): """ Loads the location of the background files from the environmant variable BELLE2_BACKGROUND_DIR which is set on the validation server and ensures that background files exist and returns the list of background files which can be directly used with add_simulation() : >>> add_simulation(main, bkgfiles=background.get_background_files()) Will fail with an assert if no background folder set or if no background file was found in the set folder. Parameters: folder (str): A specific folder to search for background files can be given as an optional parameter output_file_info (str): If true, a list of the found background files and there size will be printed This is useful to understand later which background campaign has been used to simulate events. """ env_name = 'BELLE2_BACKGROUND_DIR' bg = None if folder is None: if env_name not in os.environ: raise RuntimeError("Environment variable {} for backgound files not set. Terminanting this script.".format(env_name)) folder = os.environ[env_name] bg = glob.glob(folder + '/*.root') if len(bg) == 0: raise RuntimeError("No background files found in folder {} . Terminating this script.".format(folder)) B2INFO("Background files loaded from folder {}".format(folder)) # sort for easier comparison bg = sorted(bg) if output_file_info: bg_sizes = [os.path.getsize(f) for f in bg] # reformat to work with pretty_print_table table_rows = [list(entry) for entry in zip(bg, bg_sizes)] table_rows.insert(0, ["- Background file name -", "- file size -"]) pretty_print_table(table_rows, [0, 0]) return bg
[docs]def add_output(path, bgType, realTime, sampleType, phase=3, fileName='output.root', excludeBranches=None): ''' A function to be used for output of BG simulation. @param path path name @param bgType background type, to get available types: basf2 -m BeamBkgTagSetter @param realTime equivalent time of superKEKB running in [ns] @param sampleType 'study' (for BG studies) or 'usual', 'PXD', 'ECL' (for BG mixer) @param specify the Phase, 1 for Phase 1, 2 for Phase 2, and 3 for Physics Run or Phase 3 @param fileName optional file name, can be overridden by basf2 -o ''' if excludeBranches is None: excludeBranches = [] if sampleType == 'study': madeFor = '' branches = [] elif sampleType == 'usual' and phase == 3: madeFor = '' branches = [ 'PXDSimHits', 'SVDSimHits', 'CDCSimHits', 'TOPSimHits', 'ARICHSimHits', 'ECLHits', 'KLMSimHits', ] elif sampleType == 'usual' and phase == 2: madeFor = '' branches = [ 'PXDSimHits', 'SVDSimHits', 'CDCSimHits', 'TOPSimHits', 'ARICHSimHits', 'ECLHits', 'KLMSimHits', 'CLAWSSimHits', 'FANGSSimHits', 'PlumeSimHits', 'BeamabortSimHits', 'PindiodeSimHits', 'QcsmonitorSimHits', 'He3tubeSimHits', 'MicrotpcSimHits', ] elif sampleType == 'usual' and phase == 1: madeFor = '' branches = [ 'ClawSimHits', 'BeamabortSimHits', 'PindiodeSimHits', 'QcsmonitorSimHits', 'He3tubeSimHits', 'MicrotpcSimHits', 'BgoSimHits', 'CsiSimHits', ] elif sampleType == 'ECL': madeFor = 'ECL' branches = ['ECLHits'] elif sampleType == 'PXD': madeFor = 'PXD' branches = ['PXDSimHits'] else: madeFor = '' branches = [] B2ERROR('add_output - invalid value of argument sampleType: %s' % sampleType) # Set background tag in SimHits and add BackgroundMetaData into persistent tree tagSetter = path.add_module('BeamBkgTagSetter', backgroundType=bgType, realTime=realTime, specialFor=madeFor, Phase=phase) # Write out only non-empty events when producing samples for BG mixer if sampleType != 'study': emptyPath = create_path() tagSetter.if_false(emptyPath) # Output to file. We don't need a TTreeIndex for background files and memory # consumption can be improved by setting a lower autoFlushSize so that # fewer and or smaller amounts of data have to be read for each GetEntry() path.add_module('RootOutput', outputFileName=fileName, branchNames=branches, excludeBranchNames=excludeBranches, buildIndex=False, autoFlushSize=-500000)