implementation for the plotting + some bugfixing

rdstar/studies/low_level_cut_study/decay_channels_cut_study.py

@@ -3,10 +3,24 @@ This file updates the decay channels cut infos for the low level cut study.
 
 Patrick Ecker 2020
 """
+import copy
+from rdstar.online_analysis.decay_channels import D_decay_label_to_cut_info_list as _original_D_decay_label_to_cut_info_list
+from rdstar.online_analysis.decay_channels import Dstar_decay_label_to_cut_info_list as _original_Dstar_decay_label_to_cut_info_list
+from rdstar.online_analysis.decay_channels import decay_channels_utils, channels_with_neutrinos, \
+    generator_level_only_channels
+
+__all__ = [
+    "decay_channel_dict",
+    "D_decay_label_to_cut_info_list",
+    "D_decay_label_to_cut_list",
+    "Dstar_decay_label_to_cut_info_list",
+    "Dstar_decay_label_to_cut_list"
+]
+
 
-from rdstar.online_analysis.decay_channels import *
 
 # Update the D cuts
+D_decay_label_to_cut_info_list = copy.copy(_original_D_decay_label_to_cut_info_list)
 for channel, cut_infos in D_decay_label_to_cut_info_list.items():
     D_decay_label_to_cut_info_list[channel] = [cut_info.scale_cut(percentage=0.1) for cut_info in cut_infos]
 
@@ -15,6 +29,7 @@ for channel, cut_infos in D_decay_label_to_cut_info_list.items():
     D_decay_label_to_cut_list[channel] = [cut_info.cut_str for cut_info in cut_infos]
 
 # Update the D* cuts
+Dstar_decay_label_to_cut_info_list = copy.copy(_original_Dstar_decay_label_to_cut_info_list)
 for channel, cut_infos in Dstar_decay_label_to_cut_info_list.items():
     Dstar_decay_label_to_cut_info_list[channel] = [cut_info.scale_cut(percentage=0.1) for cut_info in cut_infos]
 

rdstar/studies/low_level_cut_study/low_level_cut_defintions.py

@@ -7,10 +7,10 @@ from rdstar.online_analysis.final_state_particles import fsp_cut_sets
 
 from rdstar.studies.low_level_cut_study.decay_channels_cut_study import D_decay_label_to_cut_info_list as D_loose_dict, \
     Dstar_decay_label_to_cut_info_list as Dstar_loose_dict
-from rdstar.online_analysis.decay_channels import *
+from rdstar.online_analysis.decay_channels import D_decay_label_to_cut_info_list, Dstar_decay_label_to_cut_info_list, \
+    D_decay_label_to_decay_id, Dstar_decay_label_to_decay_id
 
 
-# todo: from analysistools.processing.observables_utils import makeROOTCompatible, invert_make_ROOT_compatible
 
 ntuple_labels_dict = {
     "e_ntuple": "e",
@@ -65,7 +65,7 @@ for n_tuple, cut_lists in ntuple_d_dict.items():
     for channel in cut_lists["cuts"].keys():
         D_cut_dict.update(
             {
-                n_tuple : {
+                f"{n_tuple}{channel}" : {
                     "cuts": cut_lists["cuts"][channel],
                     "cuts_loose": cut_lists["cuts_loose"][channel],
                     "channel": channel
@@ -73,13 +73,23 @@ for n_tuple, cut_lists in ntuple_d_dict.items():
             }
         )
 
+combined_decay_mode_id_dict = {}
+combined_decay_mode_id_dict.update(D_decay_label_to_decay_id)
+combined_decay_mode_id_dict.update(Dstar_decay_label_to_decay_id)
+
 if __name__ == '__main__':
-    #for tuple, info in fsp_cut_dict.items():
-    #    print(tuple)
-    #    for item in info["cuts"]:
-    #        print(item)
-    #
-    #    for item in info["cuts_loose"]:
-    #        print(item)
+    for channel, item in Dstar_decay_label_to_cut_info_list.items():
+        for c in item:
+            print(c)
+
+    for tuple, info in D_cut_dict.items():
+        print(tuple)
+        for item in info["cuts"]:
+            print("cuts")
+            print(item)
+        for item in info["cuts_loose"]:
+            print(item)
 
-    print(D_cut_dict)
\ No newline at end of file
+    print(fsp_cut_dict)
+    print(D_cut_dict)
+    print(combined_decay_mode_id_dict)
\ No newline at end of file

rdstar/studies/low_level_cut_study/low_level_cut_study.py

@@ -16,6 +16,7 @@ from rdstar.offline_analysis.tasks.logistics_tasks import DataCopyTask
 
 from rdstar.offline_analysis.tasks import utilities
 from rdstar.studies.low_level_cut_study.utilities import get_config_base_path
+from rdstar.studies.low_level_cut_study.plotting import LowLevelCutPlotter
 
 logger = logging.getLogger()
 logger.setLevel(logging.DEBUG)
@@ -117,11 +118,18 @@ class LowLevelCutsEvaluationTask(utilities.BaseOfflineTask):
 
     def run_evaluation(self, dry_run: bool=False) -> Dict[str, os.fspath]:
         output_dir = os.path.join(self.cut_evaluation_dir, self.n_tuple)
-
+        output = {}
         if not dry_run:
-            l=1
-
-        return {"test": os.path.join(output_dir, "test")}
+            df = pd.read_hdf(self.input_dict[f"{self.n_tuple}_total"])
+            assert isinstance(df, pd.DataFrame)
+            llcp = LowLevelCutPlotter(
+                n_tuple=self.n_tuple,
+                df=df
+            )
+            plot_dict = llcp.plot_on(target_path=output_dir)
+            output.update(plot_dict)
+
+        return output
 
     def run(self) -> None:
         utilities.host_check(hosts=utilities.GlobalParameters().work_host_list)

rdstar/studies/low_level_cut_study/plotting.py

+import matplotlib.pyplot as plt
+
+from typing import Any, Union, Optional, List, Tuple, Dict, Generator
+import pandas as pd
+import copy
+import os
+
+from rdstar.offline_analysis.selection.selection_info_containers import CutInfo
+
+from templatefitter.plotter.plot_style import set_matplotlibrc_params, KITColors
+
+from rdstar.studies.low_level_cut_study.low_level_cut_defintions import fsp_cut_dict, D_cut_dict, combined_decay_mode_id_dict
+
+
+class LowLevelCutPlotter:
+
+    def __init__(self, n_tuple: str, df: pd.DataFrame) -> None:
+        self.n_tuple = n_tuple
+        self.df = df
+        self.cut_dicts = fsp_cut_dict
+        self.cut_dicts.update(D_cut_dict)
+        print(self.cut_dicts)
+
+    def apply_cuts(self, pivot: CutInfo, cut_infos: List[CutInfo]) -> pd.DataFrame:
+        df_temp = copy.copy(self.df)
+        cut_infos = [cut_info.make_root_compatible() for cut_info in cut_infos if cut_info != pivot]
+        cut_strings = [cut_info.cut_str for cut_info in cut_infos]
+        for cut in cut_strings:
+            df_temp.query(cut, inplace=True)
+
+        return df_temp
+
+    def plot_on(self, target_path: os.fspath) -> Dict[str, os.fspath]:
+        output = {}
+        relevant_keys = [key for key in self.cut_dicts if key.startswith(self.n_tuple)]
+        for key in relevant_keys:
+            channel = self.cut_dicts[key]["channel"]
+            channel_cut = f"extraInfo__bodecayModeID__bc == {combined_decay_mode_id_dict[channel]}" if channel else None
+            for cut_info in self.cut_dicts[key]["cuts"]:
+                set_matplotlibrc_params()
+                temp_cut_infos = copy.copy(self.cut_dicts[key]["cuts"])
+
+                new_cut_info = cut_info.make_root_compatible()
+                cut_var = new_cut_info.cut_var
+                df = self.apply_cuts(pivot=cut_info, cut_infos=temp_cut_infos)
+                if channel_cut:
+                    df.query(channel_cut, inplace=True)
+
+                if len(df) == 0:
+                    continue
+
+                var_range = (df[cut_var].min(),df[cut_var].max())
+
+                fig, ax = plt.subplots()
+                cuts = list(new_cut_info.cut_limits) if isinstance(new_cut_info.cut_limits, Tuple) else [new_cut_info.cut_limits]
+                for cut in cuts:
+                    ax.axvline(cut, color=KITColors.kit_black)
+                df_bkg = df.query("isSignal == 0")
+                df_bkg[cut_var].hist(
+                    bins=50,
+                    range=var_range,
+                    histtype="stepfilled",
+                    lw=1.5,
+                    color=KITColors.kit_green,
+                    label="Background"
+                )
+                df = df.query("isSignal == 1")
+                df[cut_var].hist(
+                    bins=50,
+                    range=var_range,
+                    histtype="step",
+                    lw=2,
+                    color=KITColors.kit_red,
+                    label="Signal"
+                )
+
+                plt.grid(False)
+                plt.title(cut_info.cut_str, loc="center")
+                plt.title(channel, loc="right")
+                plt.legend(loc="best")
+                save_path = os.path.join(target_path, f"{key}_{new_cut_info.cut_name}_evaluation.pdf")
+                output.update({f"{key}_{new_cut_info.cut_name}_evaluation": save_path})
+                plt.show()
+
+        return output
+
+if __name__ == '__main__':
+    #df = pd.read_hdf("/ceph/pecker/rdstar/low_level_cut_study/CombinedNTuples/e_ntuple_total.h5")
+    df = pd.read_hdf("/ceph/pecker/rdstar/low_level_cut_study/CombinedNTuples/dcharged_ntuple_total.h5")
+    llcp = LowLevelCutPlotter(n_tuple="dcharged_ntuple", df=df, cut_infos=D_cut_dict)
+    #llcp.apply_cuts(pivot=fsp_cut_dict["e_ntuple"]["cuts"][0], cut_infos=fsp_cut_dict["e_ntuple"]["cuts"])
+    llcp.plot_on(target_path="")
\ No newline at end of file