diff --git a/data/QUGS/convert.py b/data/QUGS/convert.py index 2982a03..d495b0d 100644 --- a/data/QUGS/convert.py +++ b/data/QUGS/convert.py @@ -27,12 +27,26 @@ class DamageFilesIndices(TypedDict): def generate_damage_files_index(**kwargs) -> DamageFilesIndices: + """ + Generate a dictionary of damage scenarios with file indices. + :param kwargs: Keyword arguments to specify parameters. + - prefix: Prefix for the file names (default: "zzzAD"). + - extension: File extension (default: ".TXT"). + - num_damage: Number of damage scenarios. + - file_index_start: Starting index for file names. + - col: Number of files per damage scenario. + - base_path: Base path for the files. + - undamage_file: Name of the undamaged file with extension. + :return: A dictionary where keys are damage scenario indices and values are lists of file paths. + """ + prefix: str = kwargs.get("prefix", "zzzAD") extension: str = kwargs.get("extension", ".TXT") num_damage: int = kwargs.get("num_damage") file_index_start: int = kwargs.get("file_index_start") col: int = kwargs.get("col") base_path: str = kwargs.get("base_path") + undamage_file: str = kwargs.get("undamage_file") damage_scenarios = {} a = file_index_start @@ -41,10 +55,25 @@ def generate_damage_files_index(**kwargs) -> DamageFilesIndices: damage_scenarios[i] = range(a, b) a += col b += col - # return damage_scenarios x = {} + if undamage_file: + try: + x[0] = [] + if base_path: + x[0].append( + os.path.normpath(os.path.join(base_path, f"{undamage_file}")) + ) + else: + x[0].append(f"{prefix}{undamage_file}") + except Exception as e: + print(Fore.RED + f"Error processing undamaged file: {e}") + sys.exit(1) + else: + print(Fore.RED + "No undamaged file specified, terminating.") + sys.exit(1) + for damage, files in damage_scenarios.items(): x[damage] = [] # Initialize each key with an empty list for i, file_index in enumerate(files, start=1): @@ -59,10 +88,11 @@ def generate_damage_files_index(**kwargs) -> DamageFilesIndices: # continue else: x[damage].append(f"{prefix}{file_index}{extension}") + return x # file_path = os.path.join(base_path, f"zzz{prefix}D{file_index}.TXT") - # df = pd.read_csv( file_path, sep="\t", skiprows=10) # Read with explicit column names + # df = pd.read_csv(file_path, sep="\t", skiprows=10) # Read with explicit column names class DataProcessor: @@ -91,7 +121,11 @@ class DataProcessor: """ data = [] # Find the maximum group index to determine the list size - max_group_idx = max(self.file_index.keys()) if self.file_index else 0 + max_group_idx = len(self.file_index) if self.file_index else 0 + + # Handle case when file_index is empty + if max_group_idx == 0: + raise ValueError("No file index provided; file_index is empty.") # Initialize empty lists for _ in range(max_group_idx): @@ -100,7 +134,7 @@ class DataProcessor: # Fill the list with data for group_idx, file_list in self.file_index.items(): # Adjust index to be 0-based - list_idx = group_idx - 1 + list_idx = group_idx data[list_idx] = [self._load_dataframe(file) for file in file_list] return data