diff --git a/data/QUGS/convert.py b/data/QUGS/convert.py index d0788f8..c86e34c 100644 --- a/data/QUGS/convert.py +++ b/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,7 +116,10 @@ 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') - col_indices = tuple[1] # [1, 26] + 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 df = pd.read_csv(file_path, delim_whitespace=True, skiprows=10, header=0, memory_map=True) @@ -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