B2A711-DeepContinuumSuppression_Input.py

#!/usr/bin/env python3
 
 
 
 
 
import basf2
import modularAnalysis as ma
from stdV0s import stdKshorts
from stdPi0s import stdPi0s
from vertex import TagV
import sys
import variables as v
import uproot
 
basf2.set_log_level(basf2.LogLevel.ERROR)
 
# --I/O----------------------------------------------------------------------------------------
step = 'train'
 
if len(sys.argv) >= 2:
    if sys.argv[1] not in ['train', 'test', 'apply_signal', 'apply_qqbar']:
        sys.exit("usage:\n\tbasf2 B2A711-DeepContinuumSuppression_Input.py <train,test,apply_signal,apply_qqbar>")
    else:
        step = str(sys.argv[1])
 
if step == 'train':
    input_file_list = [basf2.find_file('ccbar_sample_to_train.root', 'examples', False),
                       basf2.find_file('Bd2K0spi0_to_train.root', 'examples', False)]
elif step == 'test':
    input_file_list = [basf2.find_file('ccbar_sample_to_test.root', 'examples', False),
                       basf2.find_file('Bd2K0spi0_to_test.root', 'examples', False)]
elif step == 'apply_signal':
    input_file_list = [basf2.find_file('Bd2K0spi0_to_test.root', 'examples', False)]
elif step == 'apply_qqbar':
    input_file_list = [basf2.find_file('ccbar_sample_to_test.root', 'examples', False)]
else:
    sys.exit('Step does not match any of the available samples: `train`, `test`, `apply_signal`or `apply_qqbar`')
 
outfile = 'DNN_' + step + '.root'
# ---------------------------------------------------------------------------------------------
 
# Perform analysis.
firstpath = basf2.Path()
 
ma.inputMdstList(input_file_list, path=firstpath)
 
firstpath.add_module('ProgressBar')
 
# Build B candidate like in B2A701-ContinuumSuppression_Input.py
stdKshorts(path=firstpath)
stdPi0s('eff40_May2020', path=firstpath)
ma.reconstructDecay('B0 -> K_S0:merged pi0:eff40_May2020', '5.2 < Mbc < 5.3 and -0.3 < deltaE < 0.2', path=firstpath)
 
ma.matchMCTruth('B0', path=firstpath)
ma.buildRestOfEvent('B0', path=firstpath)
 
cleanMask = ('cleanMask', 'nCDCHits > 0 and useCMSFrame(p)<=3.2', 'p >= 0.05 and useCMSFrame(p)<=3.2')
ma.appendROEMasks('B0', [cleanMask], path=firstpath)
 
ma.buildContinuumSuppression('B0', 'cleanMask', path=firstpath)
 
# Accept only correctly reconstructed B candidates as signal
ma.applyCuts('B0', 'isSignal or isContinuumEvent', path=firstpath)
 
# Tag B candidate for Vertex information
TagV('B0', path=firstpath)
 
# Loop over each possible ROE (1 for every B candidate) in every event
roe_path = basf2.create_path()
 
deadEndPath = basf2.create_path()
 
ma.signalSideParticleFilter('B0', '', roe_path, deadEndPath)
 
# Build particle lists for low level variables
ma.fillParticleList('gamma:roe', 'isInRestOfEvent == 1 and goodBelleGamma == 1', path=roe_path)
ma.fillParticleList('gamma:signal', 'isInRestOfEvent == 0 and goodBelleGamma == 1', path=roe_path)
ma.fillParticleList('pi+:chargedProe', 'isInRestOfEvent == 1', path=roe_path)
ma.fillParticleList('pi+:chargedPsignal', 'isInRestOfEvent == 0', path=roe_path)
ma.fillParticleList('pi-:chargedMroe', 'isInRestOfEvent == 1', path=roe_path)
ma.fillParticleList('pi-:chargedMsignal', 'isInRestOfEvent == 0', path=roe_path)
 
v.variables.addAlias('cmsp', 'useCMSFrame(p)')
 
ma.rankByHighest('gamma:roe', 'cmsp', path=roe_path)
ma.rankByHighest('gamma:signal', 'cmsp', path=roe_path)
ma.rankByHighest('pi+:chargedProe', 'cmsp', path=roe_path)
ma.rankByHighest('pi+:chargedPsignal', 'cmsp', path=roe_path)
ma.rankByHighest('pi-:chargedMroe', 'cmsp', path=roe_path)
ma.rankByHighest('pi-:chargedMsignal', 'cmsp', path=roe_path)
 
# Define traditional Continuum Suppression Variables
contVars = [
    'R2',
    'thrustBm',
    'thrustOm',
    'cosTBTO',
    'cosTBz',
    'KSFWVariables(et)',
    'KSFWVariables(mm2)',
    'KSFWVariables(hso00)',
    'KSFWVariables(hso02)',
    'KSFWVariables(hso04)',
    'KSFWVariables(hso10)',
    'KSFWVariables(hso12)',
    'KSFWVariables(hso14)',
    'KSFWVariables(hso20)',
    'KSFWVariables(hso22)',
    'KSFWVariables(hso24)',
    'KSFWVariables(hoo0)',
    'KSFWVariables(hoo1)',
    'KSFWVariables(hoo2)',
    'KSFWVariables(hoo3)',
    'KSFWVariables(hoo4)',
    'CleoConeCS(1)',
    'CleoConeCS(2)',
    'CleoConeCS(3)',
    'CleoConeCS(4)',
    'CleoConeCS(5)',
    'CleoConeCS(6)',
    'CleoConeCS(7)',
    'CleoConeCS(8)',
    'CleoConeCS(9)'
]
 
# Define additional low level variables
basic_variables = ['p', 'phi', 'cosTheta', 'pErr', 'phiErr', 'cosThetaErr']
vertex_variables = ['distance', 'dphi', 'dcosTheta']
cluster_specific_variables = ['clusterNHits', 'clusterTiming', 'clusterE9E25', 'clusterReg', 'isInRestOfEvent']
track_specific_variables = ['kaonID', 'electronID', 'muonID', 'protonID', 'pValue', 'nCDCHits', 'isInRestOfEvent', 'charge']
 
# Aliases from normal coordinates to thrustframe coordinates
for variablename in basic_variables + vertex_variables:
    v.variables.addAlias('thrustsig' + variablename, 'useBThrustFrame(' + variablename + ',Signal)')
 
cluster_variables = cluster_specific_variables[:]
for variablename in basic_variables:
    cluster_variables.append('thrustsig' + variablename)
 
track_variables = track_specific_variables
for variablename in basic_variables + vertex_variables:
    track_variables.append('thrustsig' + variablename)
 
# General variables and training targets, which are nice to have in the Ntuple
variables = ['isContinuumEvent', 'isNotContinuumEvent', 'isSignal', 'M', 'p', 'Mbc', 'DeltaZ',
             'deltaE', 'daughter(0, M)', 'daughter(0, p)', 'daughter(1, M)', 'daughter(1, p)']
 
# Aliases for variable ranks created by rankByHighest function
for rank in range(10):
    for shortcut, particlelist in [('Croe', 'gamma:roe'), ('Csig', 'gamma:signal')]:
        for variable in cluster_variables:
            v.variables.addAlias(
                f'{variable}_{shortcut}{rank}', f'getVariableByRank({particlelist}, cmsp, {variable}, {rank + 1})')
            variables.append(f'{variable}_{shortcut}{rank}')
 
for rank in range(5):
    for shortcut, particlelist in [('TProe', 'pi+:chargedProe'), ('TPsig', 'pi+:chargedPsignal'),
                                   ('TMroe', 'pi+:chargedMroe'), ('TMsig', 'pi+:chargedMsignal')]:
        for variable in track_variables:
            v.variables.addAlias(
                f'{variable}_{shortcut}{rank}', f'getVariableByRank({particlelist}, cmsp, {variable}, {rank + 1})')
            variables.append(f'{variable}_{shortcut}{rank}')
 
# Create output file.
ma.variablesToNtuple('B0', variables + contVars, treename='tree', filename=outfile, path=roe_path)
 
# Loop over each possible ROE (1 for every B candidate) in every event
firstpath.for_each('RestOfEvent', 'RestOfEvents', roe_path)
 
basf2.process(firstpath)
print(basf2.statistics)
 
# Shuffle Data. Use only if enough Ram is available
try:
    with uproot.open(outfile) as outf:
        df = outf['tree'].arrays(library='pd')
    df = df.sample(frac=1)
    with uproot.recreate(outfile) as outf:
        outf['tree'] = df
except OSError as e:
    print(e)