diff --git a/machine_learning_hep/processer.py b/machine_learning_hep/processer.py index 2fa1a77052..4ed72a4573 100755 --- 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 index e068c52128..259e568c06 100644 --- 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