caf_multiple_collections.py

 
 
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:])