extmuid_createPions.py

#!/usr/bin/env python3
 
 
 
 
 
"""
<header>
    <output>pion-ExtMuidValidation.root</output>
    <contact>piilonen@vt.edu</contact>
    <description>Create events with 1 pion track for ext/muid validation.</description>
</header>
"""
 
import glob
import basf2 as b2
import os
from simulation import add_simulation
from reconstruction import add_reconstruction
 
ACTIVE = True
 
 
def run():
    """
    Create pion sample for the ExtMuid validation.
    """
 
    b2.set_random_seed(654321)
 
    output_filename = '../pion-ExtMuidValidation.root'
 
    print(output_filename)
 
    path = b2.create_path()
 
    eventinfosetter = b2.register_module('EventInfoSetter')
    eventinfosetter.param('evtNumList', [1000])
    path.add_module(eventinfosetter)
 
    pgun = b2.register_module('ParticleGun')
    param_pgun = {
        'pdgCodes': [-211, 211],
        'nTracks': 1,
        'varyNTracks': 0,
        'momentumGeneration': 'uniform',
        'momentumParams': [0.5, 5.0],
        'thetaGeneration': 'uniformCos',
        'thetaParams': [15., 150.],
        'phiGeneration': 'uniform',
        'phiParams': [0.0, 360.0],
        'vertexGeneration': 'fixed',
        'xVertexParams': [0.0],
        'yVertexParams': [0.0],
        'zVertexParams': [0.0],
    }
    pgun.param(param_pgun)
    path.add_module(pgun)
 
    # add simulation and reconstruction modules to the path
    if 'BELLE2_BACKGROUND_DIR' in os.environ:
        background_files = glob.glob(os.environ['BELLE2_BACKGROUND_DIR'] + '/*.root')
        add_simulation(path, bkgfiles=background_files)
    else:
        b2.B2FATAL('BELLE2_BACKGROUND_DIR is not set.')
 
    add_reconstruction(path)
 
    output = b2.register_module('RootOutput')
    output.param('outputFileName', output_filename)
    output.param('branchNames', ['MCParticles', 'ExtHits', 'KLMMuidLikelihoods', 'KLMHit2ds'])
    path.add_module(output)
    path.add_module('Progress')
 
    b2.process(path)
    print(b2.statistics)
 
 
if __name__ == '__main__':
    if ACTIVE:
        run()
    else:
        print("This validation deactivated and thus basf2 is not executed.\n"
              "If you want to run this validation, please set the 'ACTIVE' flag above to 'True'.\n"
              "Exiting.")