Source code for vibe.analysis_validation_modes.svd.SVDResolution_validation_mode

import basf2

try:
    import rawdata
    import tracking
    import svd
except ImportError:
    pass
import modularAnalysis as ma
import numpy as np

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

__all__ = [
    "SVDResolution",
]




class SVDResolution(ValidationModeBaseClass):
    name = "SVDResolution"



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

        main_path.add_module("Gearbox")
        main_path.add_module("Geometry", useDB=True)

        rawdata.add_unpackers(main_path)

        # now do reconstruction
        tracking.add_tracking_reconstruction(
            main_path,
            mcTrackFinding=False,
            trackFitHypotheses=[211],
            append_full_grid_cdc_eventt0=True,
            skip_full_grid_cdc_eventt0_if_svd_time_present=False,
        )

        # Reconstruct strips
        svd.add_svd_create_recodigits(main_path)

        ma.fillParticleLists(decayStringsWithCuts=[("pi+:all", "")], path=main_path)

        self.variables_to_validation_ntuple(
            decay_str="pi+:all",
            variables=[
                "SVDClusterCharge",
                "SVDResidual",
                "SVDTrackPrime",
                "SVDLayer",
                "SVDLadder",
                "SVDSensor",
                "SVDSide",
            ],
            path=main_path,
        )

        return main_path


    @property
    def analysis_validation_histograms(self):
        return [
            Histogram(
                name="SVDClusterCharge",
                title="",
                hist_variable=HistVariable(
                    df_label=makeROOTCompatible(variable="SVDClusterCharge"),
                    label=r"$Cluster charge$",
                    unit=r"[e-]$",
                    bins=50,
                    scope=(0, 100e3),
                ),
                hist_components=[
                    HistComponent(
                        label="All",
                    ),
                ],
                particle_list="pi+:all",
            )
        ]

    @property
    def plotting_strategies(self):
        return ["svd_resolution"]



    def get_ntuple_particle_lists(self, basf2_kwargs=None):
        if basf2_kwargs is None:
            basf2_kwargs = {}
        basf2_path = self.create_basf2_path(**basf2_kwargs)

        return {
            p.values
            for m in basf2_path.modules()
            if m.type() == "SVDVariablesToStorage"
            for p in m.available_params()
            if p.name == "particleListName"
        }




    def get_ntuple_trees(self, basf2_kwargs=None):
        if basf2_kwargs is None:
            basf2_kwargs = {}
        basf2_path = self.create_basf2_path(**basf2_kwargs)

        return {
            p.values
            for m in basf2_path.modules()
            if m.type() == "SVDVariablesToStorage"
            for p in m.available_params()
            if p.name == "containerName"
        }




    def variables_to_validation_ntuple(self, decay_str, variables, path):
        path.add_module(
            "SVDVariablesToStorage",
            outputFileName=self.output_filename,
            containerName=self.get_tree_name(decay_str),
            particleListName=decay_str,
            variablesToNtuple=variables,
        )

        last_module = path.modules()[-1]
        last_module.set_name(f"{last_module.type()}_{self.name}")




    def offline_df_manipulation(self, df):
        # Set this as they are not set by the SVDVariablesToStorage module and are required for plotting
        df[["__experiment__", "__run__", "__event__"]] = np.nan
        df["__weight__"] = 1.0
        return df