data/QUGS/convert.py

@@ -103,9 +103,10 @@ def generate_df_tuples(total_dfs=30, group_size=5, prefix="zzzAD", ext="TXT", fi
 
 
 class DataProcessor:
-    def __init__(self, file_index, cache_path: str = None, base_path: str = None):
+    def __init__(self, file_index, cache_path: str = None, base_path: str = None, include_time: bool = False):
         self.file_index = file_index
         self.base_path = base_path
+        self.include_time = include_time
         if cache_path:
             self.data = load(cache_path)
         else:
@@ -115,6 +116,9 @@ class DataProcessor:
         for idxs, group in enumerate(self.file_index):
             for idx, tuple in enumerate(group):
                 file_path = os.path.join(self.base_path, tuple[0]) # ('zzzAD1.TXT')
+                if self.include_time:
+                    col_indices = [0] + tuple[1]  # [1, 26] + [0] -> [0, 1, 26]
+                else:
                     col_indices = tuple[1] # [1, 26]
                 try:
                     # Read the CSV file
@@ -272,30 +276,18 @@ class DataProcessor:
         :param output_dir: Directory to save the CSV files.
         :param file_prefix: Prefix for the output filenames.
         """
-        for group_idx, group in enumerate(self.data, start=0):
+        for group_idx, group in enumerate(self.file_index, start=0):
             group_folder = os.path.join(output_dir, f"{file_prefix}_{group_idx}")
             os.makedirs(group_folder, exist_ok=True)
+
             for test_idx, df in enumerate(group, start=1):
-                # Ensure columns are named uniquely if duplicated
-                df = df.copy()
-                df.columns = ["Time", "Real_0", "Real_1"]  # Rename
-
-                # Export first Real column
-                out1 = os.path.join(
-                    group_folder, f"{file_prefix}_{group_idx}_TEST{test_idx}_01.csv"
-                )
-                df[["Time", "Real_0"]].rename(columns={"Real_0": "Real"}).to_csv(
-                    out1, index=False
-                )
-
-                # Export last Real column
-                out2 = os.path.join(
-                    group_folder, f"{file_prefix}_{group_idx}_TEST{test_idx}_02.csv"
-                )
-                df[["Time", "Real_1"]].rename(columns={"Real_1": "Real"}).to_csv(
-                    out2, index=False
-                )
+                out1 = os.path.join(group_folder, f"{file_prefix}_{group_idx}_TEST{test_idx}_01.csv")
+                cols_to_export = [0, 1] if self.include_time else [1]
+                df.iloc[:, cols_to_export].to_csv(out1, index=False)
 
+                out2 = os.path.join(group_folder, f"{file_prefix}_{group_idx}_TEST{test_idx}_02.csv")
+                cols_to_export = [0, 2] if self.include_time else [2]
+                df.iloc[:, cols_to_export].to_csv(out2, index=False)
 
 # def create_damage_files(base_path, output_base, prefix):
 #     # Initialize colorama