caf_multiple_collections.py

# This steering file shows off how you can configure a single Calibration to run
# multiple different Collector configurations. Using different database chains,
# different Collector module configs, different input files etc.
 
import basf2 as b2
 
import os
import sys
 
from ROOT.Belle2 import TestCalibrationAlgorithm
 
from caf.framework import Calibration, CAF, Collection
from caf import backends
 
b2.set_log_level(b2.LogLevel.INFO)
 
 
def main(argv):
    if len(argv) == 2:
        data_dir_1 = argv[0]
        data_dir_2 = argv[1]
    else:
        print("Usage: python3 caf_multiple_collections.py <data directory 1> <data directory 2>")
        sys.exit(1)
 
    # In this script we want to use two different sources of input data, and reconstruct them
    # differently before the Collector module runs. So we will need
    #
    # 1) Two sets of input data files
    # 2) Two pre_collector paths for reconstruction
    # 3) Two instances of a Collector module (using the same prefix/name)
    # 4) Two database chains to use
 
    
    input_files_1 = [os.path.join(os.path.abspath(data_dir_1), '*.root')]
    input_files_2 = [os.path.join(os.path.abspath(data_dir_2), '*.root')]
 
    
 
    col_test_1 = b2.register_module('CaTest')
    col_test_1.set_name('Test')   # Sets the prefix of collected data
    col_test_1.param('granularity', 'all')
    col_test_1.param('spread', 1)
 
    col_test_2 = b2.register_module('CaTest')
    col_test_2.set_name('Test')   # Sets the prefix of collected data
    col_test_2.param('granularity', 'all')
    col_test_2.param('spread', 20)
 
    
    rec_path_1 = b2.create_path()
    rec_path_1.add_module('Gearbox')
    rec_path_1.add_module('Geometry', excludedComponents=['SVD', 'PXD', 'ARICH', 'BeamPipe', 'EKLM'])
    # could now add reconstruction modules dependent on the type on input data
 
    # Excluded CDC instead of ARICH
    rec_path_2 = b2.create_path()
    rec_path_2.add_module('Gearbox')
    rec_path_2.add_module('Geometry', excludedComponents=['SVD', 'PXD', 'CDC', 'BeamPipe', 'EKLM'])
    # could now add reconstruction modules dependent on the type on input data
 
    
    alg_test = TestCalibrationAlgorithm()
    # Must have the same prefix as colllector prefixes. By default this algorithm and collectors
    # have the correct prefixes if we hadn't changed anything. So this is unnecessary, but we're just showing you how it's done.
    alg_test.setPrefix('Test')
 
    
    collection_1 = Collection(collector=col_test_1,
                              input_files=input_files_1,
                              pre_collector_path=rec_path_1,
                              max_files_per_collector_job=1
                              # database_chain=[LocalDatabase('path/to/database.txt'), CentralDatabase('global_tag')],
                              # output_patterns=["CollectorOutput.root"],  Does your collector output something else
                              #                                              you want passed to the algorithm?
                              # backend_args={'queue': 's'}  Only necessary for LSF/PBS backends
                              )
    # If you have some local databases or want to override the default global tag for this Collection you can do that
    # with these functions
    # collection_1.reset_database()
    # collection_1.use_local_database("mylocaldb/database.txt")
    # collection_1.use_central_database("BelleII_GlobalTag_Tutorial")
 
    collection_2 = Collection(collector=col_test_2,
                              input_files=input_files_2,
                              pre_collector_path=rec_path_2,
                              max_files_per_collector_job=1
                              )
 
    
    cal_test = Calibration('TestCalibration')
    # Add collections in with unique names
    cal_test.add_collection(name='collection_1', collection=collection_1)
    cal_test.add_collection(name='collection_2', collection=collection_2)
    cal_test.algorithms = alg_test
 
    
 
    
    cal_fw = CAF()
    cal_fw.add_calibration(cal_test)
    cal_fw.backend = backends.Local(max_processes=4)
    cal_fw.run()
    print("End of CAF processing.")
 
 
if __name__ == "__main__":
    main(sys.argv[1:])