data_reader refactor

data/src/consts.py

@@ -1,7 +1,8 @@
 import numpy as np
 
 SAMPLING_RATE = 250
-SAMPLES_PER_5_SEC = 1250
+WINDOW_TIME_FRAME = 5
+SAMPLES_PER_5_SEC = SAMPLING_RATE * WINDOW_TIME_FRAME
 ADAPT_TRIALS = 8
 
 EXPECTED_FREQS = [6.66, 7.50, 8.57, 10.00, 12.00]

data/src/data_reader.py

@@ -1,95 +1,130 @@
 from collections import defaultdict
 import scipy.io
 from data_classes import *
-from consts import *
-
-
-def load_data(filename: str) -> tuple[np.ndarray, np.ndarray, list[int]]:
-    mat_data = scipy.io.loadmat(filename)
-    eeg_data = mat_data[EEG_KEY]
-    din_data = mat_data[DIN_KEY]
-
-    latencies_ms = [int(din_data[1, i].item()) for i in range(din_data.shape[1])]
-    latencies = [int(round(l * SAMPLING_RATE / 1000)) for l in latencies_ms]
-
-    return eeg_data, latencies_ms, latencies
-
-
-def group_din_markers_into_trials(
-    latencies: list[int], gap_threshold_samples: int
-) -> list[trial_group]:
-    trial_groups = []
-    current_start_idx = 0
-    current_group_start = latencies[0]
-    current_group_end = latencies[0]
-
-    for i in range(1, len(latencies)):
-        if latencies[i] - latencies[i - 1] > gap_threshold_samples:
-            trial_groups.append(
-                trial_group(
-                    start_sample=current_group_start,
-                    end_sample=current_group_end,
-                    start_idx=current_start_idx,
-                    end_idx=i - 1,
+from consts import (
+    EEG_KEY,
+    DIN_KEY,
+    SAMPLING_RATE,
+    EXPECTED_FREQS,
+    ADAPT_TRIALS,
+    WINDOW_TIME_FRAME,
+)
+from abc import ABC, abstractmethod
+
+
+class DataReader(ABC):
+    @abstractmethod
+    def _load_data(self, filename: str) -> tuple[np.ndarray, np.ndarray, list[int]]:
+        pass
+
+    @abstractmethod
+    def get_all_trials(self, filename: str) -> list[trial_info]:
+        pass
+
+
+class MatDataReader(DataReader):
+    def _load_data(
+        self, filename: str, sampling_rate: int = SAMPLING_RATE
+    ) -> tuple[np.ndarray, np.ndarray, list[int]]:
+        mat_data = scipy.io.loadmat(filename)
+        eeg_data = mat_data[EEG_KEY]
+        din_data = mat_data[DIN_KEY]
+
+        latencies_ms = [int(din_data[1, i].item()) for i in range(din_data.shape[1])]
+        latencies = [int(round(l * sampling_rate / 1000)) for l in latencies_ms]
+
+        return eeg_data, latencies_ms, latencies
+
+    def _group_din_markers_into_trials(
+        self, latencies: list[int], gap_threshold_samples: int
+    ) -> list[trial_group]:
+        trial_groups = []
+        current_start_idx = 0
+        current_group_start = latencies[0]
+        current_group_end = latencies[0]
+
+        for i in range(1, len(latencies)):
+            if latencies[i] - latencies[i - 1] > gap_threshold_samples:
+                trial_groups.append(
+                    trial_group(
+                        start_sample=current_group_start,
+                        end_sample=current_group_end,
+                        start_idx=current_start_idx,
+                        end_idx=i - 1,
+                    )
                 )
+                current_group_start = latencies[i]
+                current_start_idx = i
+            current_group_end = latencies[i]
+
+        trial_groups.append(
+            trial_group(
+                start_sample=current_group_start,
+                end_sample=current_group_end,
+                start_idx=current_start_idx,
+                end_idx=len(latencies) - 1,
             )
-            current_group_start = latencies[i]
-            current_start_idx = i
-        current_group_end = latencies[i]
-
-    trial_groups.append(
-        trial_group(
-            start_sample=current_group_start,
-            end_sample=current_group_end,
-            start_idx=current_start_idx,
-            end_idx=len(latencies) - 1,
         )
-    )
-
-    return trial_groups
 
+        return trial_groups
+
+    def _extract_trials(
+        self,
+        eeg_data: np.ndarray,
+        latencies_ms: list[int],
+        trial_groups: list[trial_group],
+        adapt_trials: int = ADAPT_TRIALS,
+        trial_time_frame_sec: int = WINDOW_TIME_FRAME,
+        sampling_rate: int = SAMPLING_RATE,
+    ) -> list[trial_info]:
+        all_trials = []
+        samples_per_trial = sampling_rate * trial_time_frame_sec
+
+        for i, tg in enumerate(trial_groups):
+            start_sample = tg.start_sample
+            epoch = eeg_data[:, start_sample : start_sample + samples_per_trial]
+
+            if epoch.shape[1] != samples_per_trial:
+                continue
+
+            dins_ms = latencies_ms[tg.start_idx : tg.end_idx + 1]
+            n_dins = len(dins_ms)
+            true_freq = 0.0
+            closest = 0.0
+
+            if n_dins > 1:
+                intervals = [
+                    dins_ms[j] - dins_ms[j - 1] for j in range(1, len(dins_ms))
+                ]
+                true_freq = (1000.0 / np.mean(intervals)) / 2.0
+                closest = min(EXPECTED_FREQS, key=lambda f: abs(f - true_freq))
+
+            trial_type = TrialType.ADAPT if i < adapt_trials else TrialType.TEST
+
+            trial_record = trial_info(
+                epoch=epoch,
+                trial=i,
+                true_freq=round(true_freq, 2),
+                closest_freq=closest,
+                n_dins=n_dins,
+                type=trial_type,
+            )
 
-def extract_trials(
-    eeg_data: np.ndarray, latencies_ms: list[int], trial_groups: list[trial_group]
-) -> list[trial_info]:
-    all_trials = []
-
-    for i, tg in enumerate(trial_groups):
-        start_sample = tg.start_sample
-        epoch = eeg_data[:, start_sample : start_sample + SAMPLES_PER_5_SEC]
-
-        if epoch.shape[1] != SAMPLES_PER_5_SEC:
-            continue
-
-        dins_ms = latencies_ms[tg.start_idx : tg.end_idx + 1]
-        n_dins = len(dins_ms)
-        true_freq = 0.0
-        closest = 0.0
-
-        if n_dins > 1:
-            intervals = [dins_ms[j] - dins_ms[j - 1] for j in range(1, len(dins_ms))]
-            true_freq = (1000.0 / np.mean(intervals)) / 2.0
-            closest = min(EXPECTED_FREQS, key=lambda f: abs(f - true_freq))
-
-        trial_type = TrialType.ADAPT if i < ADAPT_TRIALS else TrialType.TEST
+            all_trials.append(trial_record)
 
-        trial_record = trial_info(
-            epoch=epoch,
-            trial=i,
-            true_freq=round(true_freq, 2),
-            closest_freq=closest,
-            n_dins=n_dins,
-            type=trial_type,
-        )
+        return all_trials
 
-        all_trials.append(trial_record)
+    def get_all_trials(self, filename: str) -> list[trial_info]:
+        eeg_data, latencies_ms, latencies = self._load_data(filename)
+        trial_groups = self._group_din_markers_into_trials(latencies, 250)
+        all_trials = self._extract_trials(eeg_data, latencies_ms, trial_groups)
 
-    return all_trials
+        return all_trials
 
 
-def main(filename: str) -> list[trial_info]:
-    eeg_data, latencies_ms, latencies = load_data(filename)
-    trial_groups = group_din_markers_into_trials(latencies, 250)
-    all_trials = extract_trials(eeg_data, latencies_ms, trial_groups)
+class CSVDataReader(DataReader):
+    def _load_data(self, filename):
+        return super()._load_data(filename)
 
-    return all_trials
+    def get_all_trials(self, filename):
+        return super().get_all_trials(filename)

data/src/fft_analysis.py

@@ -4,7 +4,7 @@
 
 from consts import (
     SAMPLING_RATE,
-    SAMPLES_PER_5_SEC,
+    WINDOW_TIME_FRAME,
     SNR_NEIGHBOR_BINS,
     SNR_EXCLUDE_BINS,
     SNR_THRESHOLD_DB,
@@ -16,8 +16,8 @@
 from data_classes import trial_info, TrialType
 from utils import group_trials_by_frequency, get_trials_by_type
 
-N = SAMPLES_PER_5_SEC
-freqs_fft = np.fft.rfftfreq(N, d=1.0 / SAMPLING_RATE)
+samples_per_trial = SAMPLING_RATE * WINDOW_TIME_FRAME
+freqs_fft = np.fft.rfftfreq(samples_per_trial, d=1.0 / SAMPLING_RATE)
 
 
 def compute_rfft(signal: np.ndarray) -> np.ndarray:
@@ -27,7 +27,7 @@ def compute_rfft(signal: np.ndarray) -> np.ndarray:
 
 def compute_fft_magnitude(signal: np.ndarray) -> np.ndarray:
     fft_result = compute_rfft(signal)
-    magnitude = (2.0 / N) * np.abs(fft_result)
+    magnitude = (2.0 / samples_per_trial) * np.abs(fft_result)
     return magnitude
 
 

data/src/main.py

@@ -1,5 +1,5 @@
 import argparse
-from data_reader import main as read_data
+from data_reader import DataReader, MatDataReader, CSVDataReader
 from filter import main as filter_data
 from fft_analysis import main as fft_analysis_main
 from filter_plotter import (
@@ -14,10 +14,22 @@
 from consts import OUTPUT_DIR
 import numpy as np
 
+# Data file type to reader map
+DATA_READERS = {
+    "csv": CSVDataReader,
+    "mat": MatDataReader,
+}
 
-def main(filename: str, plot_filter: bool, plot_fft: bool):
+
+def main(
+    filename: str,
+    plot_filter: bool,
+    plot_fft: bool,
+    reader: DataReader,
+    output_dir: str = OUTPUT_DIR,
+):
     print("\nReading data...")
-    trials = read_data(filename)
+    trials = reader.get_all_trials(filename)
 
     # Save raw (unfiltered) epochs before applying any filtering
     raw_epochs = np.array([t.epoch for t in trials])
@@ -50,20 +62,32 @@ def main(filename: str, plot_filter: bool, plot_fft: bool):
                     r["ch_idx"],
                     r["channel"],
                     r["snr_db"],
-                    OUTPUT_DIR,
+                    output_dir,
                 )
 
         for ef in all_results:
-            plot_grouped_fft(ef, all_results[ef], stats, OUTPUT_DIR)
+            plot_grouped_fft(ef, all_results[ef], stats, output_dir)
 
-        plot_snr_heatmap(stats, OUTPUT_DIR)
+        plot_snr_heatmap(stats, output_dir)
 
     print("\nPipeline complete!")
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run the full SSVEP pipeline.")
-    parser.add_argument("--filename", type=str, help="Path to the .mat file")
+    parser.add_argument(
+        "--filename",
+        type=str,
+        required=True,
+        help="Path to the data file",
+    )
+    parser.add_argument(
+        "--data_type",
+        type=str,
+        required=True,
+        choices=["csv", "mat"],
+        help="csv or mat",
+    )
     parser.add_argument(
         "--plot_filter",
         action="store_true",
@@ -74,6 +98,11 @@ def main(filename: str, plot_filter: bool, plot_fft: bool):
         action="store_true",
         help="Generate plots for the FFT results (SNR heatmap, FFT spectrum)",
     )
+    parser.add_argument(
+        "--output", default=OUTPUT_DIR, type=str, help="output directory"
+    )
     args = parser.parse_args()
 
-    main(args.filename, args.plot_filter, args.plot_fft)
+    reader_class = DATA_READERS.get(args.data_type)
+
+    main(args.filename, args.plot_filter, args.plot_fft, reader_class(), args.output)