fixed an issue with the output

rdstar/studies/low_level_cut_study/low_level_cut_study.py

@@ -121,16 +121,17 @@ 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 = {}
+        llcp = LowLevelCutPlotter(
+            n_tuple=self.n_tuple,
+            target_path=output_dir
+        )
+        output = llcp.generate_output_dict()
         if not dry_run:
             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)
+            llcp.get_df(df=df)
+            plot_dict = llcp.plot_on()
+            assert plot_dict == llcp.generate_output_dict()
 
         return output
 

rdstar/studies/low_level_cut_study/plotting.py

@@ -14,12 +14,15 @@ from rdstar.studies.low_level_cut_study.low_level_cut_defintions import fsp_cut_
 
 class LowLevelCutPlotter:
 
-    def __init__(self, n_tuple: str, df: pd.DataFrame) -> None:
+    def __init__(self, n_tuple: str, target_path: os.fspath) -> None:
         self.n_tuple = n_tuple
-        self.df = df
+        self.target_path = target_path
+        self.df = None
         self.cut_dicts = fsp_cut_dict
         self.cut_dicts.update(D_cut_dict)
-        print(self.cut_dicts)
+
+    def get_df(self, df: pd.DataFrame) -> None:
+        self.df = df
 
     def apply_cuts(self, pivot: CutInfo, cut_infos: List[CutInfo]) -> pd.DataFrame:
         df_temp = copy.copy(self.df)
@@ -30,7 +33,17 @@ class LowLevelCutPlotter:
 
         return df_temp
 
-    def plot_on(self, target_path: os.fspath) -> Dict[str, os.fspath]:
+    def generate_output_dict(self) -> 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:
+            for cut_info in self.cut_dicts[key]["cuts"]:
+                save_path = os.path.join(self.target_path, f"{key}_{cut_info.cut_name}_evaluation.pdf")
+                output.update({f"{key}_{cut_info.cut_name}_evaluation": save_path})
+
+        return output
+
+    def plot_on(self) -> 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:
@@ -47,6 +60,9 @@ class LowLevelCutPlotter:
                     df.query(channel_cut, inplace=True)
 
                 if len(df) == 0:
+                    save_path = os.path.join(self.target_path, f"{key}_{new_cut_info.cut_name}_evaluation.pdf")
+                    output.update({f"{key}_{new_cut_info.cut_name}_evaluation": save_path})
+                    plt.savefig(save_path)
                     continue
 
                 fig, ax = plt.subplots()
@@ -71,7 +87,7 @@ class LowLevelCutPlotter:
                 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")
+                save_path = os.path.join(self.target_path, f"{key}_{new_cut_info.cut_name}_evaluation.pdf")
                 output.update({f"{key}_{new_cut_info.cut_name}_evaluation": save_path})
                 plt.savefig(save_path)
                 #plt.show()
@@ -81,6 +97,8 @@ class LowLevelCutPlotter:
 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)
+    llcp = LowLevelCutPlotter(n_tuple="dcharged_ntuple", target_path="")
+    print(llcp.generate_output_dict())
+    llcp.get_df(df=df)
     #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
+    llcp.plot_on()
\ No newline at end of file