Modify the processer to delete pkl Reco files

alisw · Dec 20, 2023 · 9e05e7e · 9e05e7e
1 parent 98b2e11
commit 9e05e7e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 3 deletions.
diff --git a/machine_learning_hep/processer.py b/machine_learning_hep/processer.py
@@ -28,7 +28,7 @@
 from machine_learning_hep.selectionutils import selectfidacc
 from machine_learning_hep.bitwise import tag_bit_df #, filter_bit_df
 from machine_learning_hep.utilities import selectdfquery, merge_method, mask_df
-from machine_learning_hep.utilities import list_folders, createlist, appendmainfoldertolist
+from machine_learning_hep.utilities import list_folders, is_file_indir, createlist, appendmainfoldertolist
 from machine_learning_hep.utilities import create_folder_struc, seldf_singlevar, openfile
 from machine_learning_hep.utilities import mergerootfiles, count_df_length_pkl
 from machine_learning_hep.utilities import get_timestamp_string
@@ -181,9 +181,14 @@ def __init__(self, case, datap, run_param, mcordata, p_maxfiles, # pylint: disab
         if os.path.isdir(self.d_root):
             self.l_path = list_folders(self.d_root, self.n_root, self.p_maxfiles,
                                        self.select_jobs)
-        else:
+        elif is_file_indir(self.d_pkl, self.n_reco):
             self.l_path = list_folders(self.d_pkl, self.n_reco, self.p_maxfiles,
                                        self.select_jobs)
+        else:
+            self.n_sk = self.n_reco.replace(".pkl", "_%s%d_%d.pkl" % \
+                          (self.v_var_binning, self.lpt_anbinmin[0], self.lpt_anbinmax[0]))
+            self.l_path = list_folders(self.d_pklsk, self.n_sk, self.p_maxfiles,
+                                       self.select_jobs)
 
         self.l_root = createlist(self.d_root, self.l_path, self.n_root)
         self.l_reco = createlist(self.d_pkl, self.l_path, self.n_reco)

diff --git a/machine_learning_hep/utilities.py b/machine_learning_hep/utilities.py
@@ -165,7 +165,7 @@ def list_folders(main_dir, filenameinput, maxfiles, select=None): # pylint: disa
             for subdir1 in list_subdir1:
                 subdir1full = os.path.join(subdir0full, subdir1)
                 if os.path.isdir(subdir1full):
-                    if os.listdir(subdir1full)[0] == filenameinput:
+                    if filenameinput in os.listdir(subdir1full):
                         list_files_ = os.listdir(subdir1full)
                         for myfile in list_files_:
                             filefull = os.path.join(subdir1full, myfile)
@@ -194,6 +194,15 @@ def list_folders(main_dir, filenameinput, maxfiles, select=None): # pylint: disa
         listfolders = listfolders[:maxfiles]
     return  listfolders
 
+def is_file_indir(main_dir, filenameinput):
+    """
+    Check if a file is contained in a directory or all its subdirectories
+    """
+    for root, dirs, files in os.walk(main_dir):
+        if filenameinput in files:
+            return True
+    return False
+
 def create_folder_struc(maindir, listpath):
     """
     Reproduce the folder structure as input