Source code for vibe.analysis_validation_modes.tracking.prompt_tracks_distributions

from typing import List
import pandas as pd
import numpy as np
import math
import ROOT
from ROOT import Belle2

import basf2 as b2
import modularAnalysis as ma
import variables.collections as vc
from variables import variables as vm
import variables.utils as vu
import vertex as vtx

import argparse
import glob
import sys

from vibe.core.utils.misc import fancy_validation_mode_header
from vibe.core.validation_mode import ValidationModeBaseClass
from vibe.core.helper.histogram_tools import HistVariable, Histogram, HistComponent
from vibe.core.helper.root_helper import makeROOTCompatible    



@fancy_validation_mode_header
class PromptTracksDistributions(ValidationModeBaseClass):
    name = "PromptTracksDistributions"
    


    def create_basf2_path(self):
        main = b2.Path()

        # uncomment to print ERRORS only
        #b2.set_log_level(b2.LogLevel.ERROR)
            
        main.add_module('Progress')
        track_cuts = 'dr < 1.0 and abs(dz) < 2.0'
        
        ma.fillParticleList('pi+:IP', track_cuts, path=main)

        vm.addAlias('nIPTracks',f'nCleanedTracks({track_cuts})')

        ##store if event passed HLT skim
        ##dimuon event
        line_mumu_tight_or_highm = "SoftwareTriggerResult(software_trigger_cut&skim&accept_mumu_tight_or_highm)"
        line_mumutight = "SoftwareTriggerResult(software_trigger_cut&skim&accept_mumutight)"
        line_mumu_2trk = "SoftwareTriggerResult(software_trigger_cut&skim&accept_mumu_2trk)"
        dimuonEventCutString = f"[{line_mumu_tight_or_highm} == 1]"
        dimuonEventCutString += f" or [[isNAN({line_mumu_tight_or_highm}) == 1] and [{line_mumutight} == 1]]"
        dimuonEventCutString += f" or [[isNAN({line_mumutight}) == 1] and [{line_mumu_2trk} == 1]]"
        vm.addAlias('isDimuonEvent',f'passesEventCut({dimuonEventCutString})')

        ## bhabha
        line_bhabha ="SoftwareTriggerResult(software_trigger_cut&skim&accept_bhabha)"
        bhabhaEventCutString = f"[{line_bhabha} == 1]"
        vm.addAlias("isBhabhaEvent",f'passesEventCut({bhabhaEventCutString})')

        ## hadron
        line_hadron ="SoftwareTriggerResult(software_trigger_cut&skim&accept_hadron)"
        hadronEventCutString = f"[{line_hadron} == 1]"
        vm.addAlias("isHadronEvent",f'passesEventCut({hadronEventCutString})')

        
        VARS = ['IPX', 'IPY', 'IPZ', "nTracks", 'nIPTracks']
        VARS += vc.track + vc.kinematics + vc.track_hits + vc.trackfit_parameters 
        VARS += ['cosTheta', 'phi', 'charge', 'trackTime', 'clusterE']
        VARS += ['isTrackFlippedAndRefitted', 'isDimuonEvent','isBhabhaEvent','isHadronEvent']

        VARS = sorted(list(set(VARS)))
        
        
        self.variables_to_validation_ntuple(
            decay_str="pi+:IP",
            variables=VARS,
            path=main,
        )

        return main