Source code for vibe.analysis_validation_modes.physics.btojpsikstar_validation_mode

from typing import List
import pandas as pd

import basf2
import modularAnalysis as ma
import vertex as vx
import flavorTagger as ft


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__ = [
    "BtoJpsiKstarValidationMode",
]



@fancy_validation_mode_header
class BtoJpsiKstarValidationMode(ValidationModeBaseClass):
    name = "BtoJpsiKstar"
    latex_str = r"$B \rightarrow J/\psi K^{*0}$"



    def create_basf2_path(self):
        main_path = basf2.Path()

        # fsp selection
        goodTrack = "abs(dz) < 2 and dr < 0.5"
        muons = ("mu-:rec", f"{goodTrack} and p > 0.7 and muonID > 0.9")
        pions = ("pi-:rec", f"{goodTrack} and pionID > 0.5")
        kaons = ("K-:rec", f"{goodTrack} and kaonID > 0.5")
        ma.fillParticleLists(
            decayStringsWithCuts=[muons, pions, kaons],
            path=main_path,
        )

        # J/psi reconstruction
        ma.reconstructDecay(
            decayString="J/psi:mu_rec -> mu-:rec mu+:rec",
            cut="2.9869 < M < 3.2069",
            path=main_path,
        )

        # Kstar reconstruction
        ma.reconstructDecay(
            decayString="K*0:rec -> pi-:rec K+:rec",
            cut="0.79555 < M < 0.99555",
            path=main_path,
        )

        # B meson reconstruction
        ma.reconstructDecay(
            decayString="B0:rec -> J/psi:mu_rec K*0:rec",
            cut="",
            path=main_path,
        )

        # MC matching + vertex + kinematic cuts
        ma.matchMCTruth(
            "B0:rec",
            path=main_path,
        )
        vx.treeFit("B0:rec", 0.001, path=main_path)
        ma.applyCuts("B0:rec", "Mbc>5.27 and abs(deltaE)<0.1", path=main_path)

        # select signal + build ROE
        ma.buildRestOfEvent("B0:rec", path=main_path)

        ft.flavorTagger(particleLists=["B0:rec"], useGNN=True, path=main_path)

        ma.buildEventShape(inputListNames=["B0:rec"], path=main_path)
        ma.buildEventKinematics(inputListNames=["B0:rec"], path=main_path)

        self.variables_to_validation_ntuple(
            decay_str="B0:rec",
            variables=list({"qrGNN", "Mbc", "deltaE"}),
            path=main_path,
        )

        return main_path


    @property
    def analysis_validation_histograms(self) -> List[Histogram]:
        return [
            Histogram(
                name="qrGNN",
                title="",
                hist_variable=HistVariable(
                    df_label=makeROOTCompatible(variable="qrGNN"),
                    label=r"$qrGNN$ (cleaned ROE)",
                    unit=r"unitless",
                    bins=50,
                    scope=(-1, 1),
                ),
                hist_components=[
                    HistComponent(
                        label="All",
                    ),
                ],
            )
        ]



    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