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