refactor(src): restructure compute_stft function to be pure function and include return parameters and improve clarity

code/src/process_stft.py

@@ -1,7 +1,8 @@
 import os
 import pandas as pd
 import numpy as np
-from scipy.signal import stft, hann
+from scipy.signal import stft
+from scipy.signal.windows import hann
 import glob
 import multiprocessing  # Added import for multiprocessing
 
@@ -19,27 +20,32 @@ for dir_path in output_dirs.values():
     os.makedirs(dir_path, exist_ok=True)
 
 # Define STFT parameters
-window_size = 1024
-hop_size = 512
-window = hann(window_size)
-Fs = 1024
 
 # Number of damage cases (adjust as needed)
 num_damage_cases = 0  # Change to 30 if you have 30 damage cases
 
 # Function to perform STFT and return magnitude
-def compute_stft(vibration_data, Fs=Fs, window_size=window_size, hop_size=hop_size):
-    frequencies, times, Zxx = stft(
-        vibration_data, 
-        fs=Fs, 
-        window=window, 
-        nperseg=window_size, 
-        noverlap=window_size - hop_size
-    )
-    stft_magnitude = np.abs(Zxx)
-    return stft_magnitude.T  # Transpose to have frequencies as columns
+def compute_stft(vibration_data, return_param=False):
+    window_size = 1024
+    hop_size = 512
+    window = hann(window_size)
+    Fs = 1024
 
-def process_damage_case(damage_num, Fs=Fs, window_size=window_size, hop_size=hop_size, output_dirs=output_dirs):
+    frequencies, times, Zxx = stft(
+                                    vibration_data, 
+                                    fs=Fs, 
+                                    window=window, 
+                                    nperseg=window_size, 
+                                    noverlap=window_size - hop_size
+                                )
+    stft_magnitude = np.abs(Zxx)
+
+    if return_param:
+        return stft_magnitude.T, [window_size, hop_size, Fs]  # Transpose to have frequencies as columns
+    else:
+        return stft_magnitude.T
+
+def process_damage_case(damage_num, Fs=Fs,):
     damage_folder = os.path.join(damage_base_path, f'DAMAGE_{damage_num}')
     if damage_num == 0:
         # Number of test runs per damage case
@@ -83,8 +89,8 @@ def process_damage_case(damage_num, Fs=Fs, window_size=window_size, hop_size=hop
             vibration_data = df.iloc[:, 1].values
             
             # Perform STFT
-            stft_magnitude = compute_stft(vibration_data, Fs=Fs, window_size=window_size, hop_size=hop_size)
-            
+            stft_magnitude, (window_size, hop_size, Fs) = compute_stft(vibration_data, return_param=True)
+
             # Convert STFT result to DataFrame
             df_stft = pd.DataFrame(
                 stft_magnitude,