Source code for vibe.analysis_validation_modes.physics.mc_compare

from typing import List
import pandas as pd

import basf2
import modularAnalysis as ma
import vertex as vx
from variables import variables as vm

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

__all__ = [
    "MCCompareMode",
]



@fancy_validation_mode_header
class MCCompareMode(ValidationModeBaseClass):
    name = "mc_compare"
    
    


    def create_basf2_path(self):
        
        main_path = basf2.Path()
        treename = 'B+:feiSL'
        
        ma.matchMCTruth(
            treename,
            path=main_path,
        )

        # select signal + build ROE
        ma.buildRestOfEvent(treename, path=main_path)

        # define ROE mask cuts
        tracksCut = "abs(d0) < 10.0 and abs(z0) < 20.0"
        clustersCut1 = "[clusterE1E9 > 0.4 or E > 0.075] and [[E > 0.062 and abs(clusterTiming)<18 and clusterReg==1] or "
        clustersCut2 = (
            "[E > 0.060 and abs(clusterTiming)<20 and clusterReg==2]"
            " or [E > 0.056 and abs(clusterTiming)<44 and clusterReg==3]]"
        )
        clustersCut = clustersCut1 + clustersCut2

        # define and append ROE masks
        m1 = ("m1", tracksCut, clustersCut)
        ma.appendROEMasks(treename, [m1], path=main_path)
        
        ma.rankByHighest(
            treename,
            "extraInfo(sigProb)",
            #allowMultiRank=True,
            outputVariable="sigProbRank",
            path=main_path,
        )
        
        ma.buildEventShape(inputListNames=[treename], path=main_path)
        ma.buildEventKinematics(inputListNames=[treename], path=main_path)

        vm.addAlias('sigProbRank', 'extraInfo(sigProbRank)')
        vm.addAlias('cosThetaBY', 'cosThetaBetweenParticleAndNominalB')
        
        self.variables_to_validation_ntuple(
            decay_str=treename,
            variables=list(
                {
                    "cosThetaBY",
                    "isSignalAcceptMissingNeutrino",
                    "sigProbRank",
                }
            ),
            path=main_path,
        )

        return main_path


    @property
    def analysis_validation_histograms(self) -> List[Histogram]:
        return [
            Histogram(
                name='cosThetaBY',
                title=r"cos$\theta_{BY}$: ccbar vs charged rank 1",
                hist_variable=HistVariable(
                    df_label=makeROOTCompatible(variable="cosThetaBY"),
                    label=r"cos$\theta_{BY}$",
                    unit=r"",
                    bins=50,
                    scope=(-5, 5),
                ),
                hist_components = [
                    HistComponent(
                        label="rank1 B's",
                        additional_cut_str= "sigProbRank==1",
                    ),
                ]
            ),
            Histogram(
                name='cosThetaBY_allVSrank1',
                title=r"cos$\theta_{BY}$: ccbar vs charged",
                hist_variable=HistVariable(
                    df_label=makeROOTCompatible(variable="cosThetaBY"),
                    label=r"cos$\theta_{BY}$",
                    unit=r"",
                    bins=50,
                    scope=(-5, 5),
                ),
                hist_components = [
                    HistComponent(
                        label='all ranks'
                    ),
                    HistComponent(
                        label="rank1 B's",
                        additional_cut_str= "sigProbRank==1",
                    ),
                ]
            ),
        ]



    def offline_df_manipulation(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.sample(frac=1.0).groupby(by=["__event__"]).head(1)  # Applying rand BCS offline
        return df




    def get_number_of_signal_for_efficiency(self, df: pd.DataFrame) -> float:
        return df["isSignalAcceptMissingNeutrino"].sum()