""" 5-Fold Stratified Cross-Validation Evaluation for Cleaned TCH EEG Dataset. Evaluates EEGNet, ShallowConvNet, and FBCSP+Classifiers. Strictly fits scaling, spatial filters, and models on training folds only (No Data Leakage). """ from __future__ import annotations from pathlib import Path import numpy as np import scipy.signal as signal from scipy.linalg import eigh import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import TensorDataset, DataLoader from sklearn.model_selection import StratifiedKFold from sklearn.metrics import accuracy_score from sklearn.svm import SVC from sklearn.discriminant_analysis import LinearDiscriminantAnalysis from sklearn.linear_model import RidgeClassifier # Set random seed for reproducibility np.random.seed(42) torch.manual_seed(42) # --- 1. EEGNet Architecture --- class EEGNet(nn.Module): def __init__(self, n_classes=2, channels=8, samples=400, F1=8, D=2, F2=16, kernel_length=32, dropout_rate=0.25): super().__init__() self.conv1 = nn.Conv2d(1, F1, (1, kernel_length), padding=(0, kernel_length // 2), bias=False) self.bn1 = nn.BatchNorm2d(F1) self.depthwise = nn.Conv2d(F1, F1 * D, (channels, 1), groups=F1, bias=False) self.bn2 = nn.BatchNorm2d(F1 * D) self.act1 = nn.ELU() self.avgpool1 = nn.AvgPool2d((1, 4)) self.drop1 = nn.Dropout(dropout_rate) self.separable = nn.Sequential( nn.Conv2d(F1 * D, F1 * D, (1, 16), padding=(0, 8), groups=F1 * D, bias=False), nn.Conv2d(F1 * D, F2, (1, 1), bias=False), nn.BatchNorm2d(F2), nn.ELU(), nn.AvgPool2d((1, 8)), nn.Dropout(dropout_rate) ) out_len = samples // 32 self.fc = nn.Linear(F2 * out_len, n_classes) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.depthwise(x) x = self.bn2(x) x = self.act1(x) x = self.avgpool1(x) x = self.drop1(x) x = self.separable(x) x = self.fc(x.flatten(start_dim=1)) return x # --- 2. ShallowConvNet Architecture (Schirrmeister et al., 2017) --- class ShallowConvNet(nn.Module): def __init__(self, n_classes=2, channels=8, samples=400, n_filters=40, kernel_length=25, dropout_rate=0.25): super().__init__() self.conv_temporal = nn.Conv2d(1, n_filters, (1, kernel_length), bias=False) self.conv_spatial = nn.Conv2d(n_filters, n_filters, (channels, 1), bias=False) self.bn = nn.BatchNorm2d(n_filters) self.pooling = nn.AvgPool2d(kernel_size=(1, 75), stride=(1, 15)) self.drop = nn.Dropout(dropout_rate) with torch.no_grad(): dummy = torch.zeros(1, 1, channels, samples) x = self.conv_temporal(dummy) x = self.conv_spatial(x) x = self.bn(x) x = x ** 2 x = self.pooling(x) x = torch.log(torch.clamp(x, min=1e-6)) x = self.drop(x) flat_dim = x.flatten(start_dim=1).shape[1] self.fc = nn.Linear(flat_dim, n_classes) def forward(self, x): x = self.conv_temporal(x) x = self.conv_spatial(x) x = self.bn(x) x = x ** 2 x = self.pooling(x) x = torch.log(torch.clamp(x, min=1e-6)) x = self.drop(x) return self.fc(x.flatten(start_dim=1)) # --- 3. Data Loader --- def load_subject_cleaned_data(subj_prefix: str, cache_dir: Path, task_samples: tuple[int, int] = (200, 600)): cleaned_dir = cache_dir / "cleaned_sessions" x_files = sorted(cleaned_dir.glob(f"{subj_prefix}*_x.npy")) X_list, y_list = [], [] for x_f in x_files: sess_name = x_f.stem.replace("_x", "") y_f = cleaned_dir / f"{sess_name}_y.npy" X_list.append(np.load(x_f)) y_list.append(np.load(y_f)) X_all = np.concatenate(X_list, axis=0) # (trials, 8, 600) y_all = np.concatenate(y_list, axis=0) # Extract task window [0, 2] s -> samples 200:600 (400 samples) X_task = X_all[:, :, task_samples[0]:task_samples[1]] return X_task, y_all # --- 4. Deep Learning Model Evaluator (EEGNet & ShallowConvNet) --- def eval_deep_model(model_cls, X: np.ndarray, y: np.ndarray, epochs: int = 100, batch_size: int = 16, lr: float = 1e-3, apply_car: bool = False): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) fold_accs = [] for fold, (tr_idx, te_idx) in enumerate(skf.split(X, y)): X_tr, y_tr = X[tr_idx].copy(), y[tr_idx] X_te, y_te = X[te_idx].copy(), y[te_idx] if apply_car: X_tr = X_tr - np.mean(X_tr, axis=1, keepdims=True) X_te = X_te - np.mean(X_te, axis=1, keepdims=True) # Fit channel standardization strictly on train fold tr_mean = np.mean(X_tr, axis=(0, 2), keepdims=True) tr_std = np.std(X_tr, axis=(0, 2), keepdims=True) + 1e-6 X_tr_norm = (X_tr - tr_mean) / tr_std X_te_norm = (X_te - tr_mean) / tr_std ds_tr = TensorDataset(torch.tensor(X_tr_norm[:, None, :, :], dtype=torch.float32), torch.tensor(y_tr, dtype=torch.long)) ds_te = TensorDataset(torch.tensor(X_te_norm[:, None, :, :], dtype=torch.float32), torch.tensor(y_te, dtype=torch.long)) loader_tr = DataLoader(ds_tr, batch_size=batch_size, shuffle=True) loader_te = DataLoader(ds_te, batch_size=len(ds_te), shuffle=False) model = model_cls(n_classes=2, channels=X.shape[1], samples=X.shape[2]).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-2) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) best_acc = 0.0 for epoch in range(epochs): model.train() for bx, by in loader_tr: bx, by = bx.to(device), by.to(device) optimizer.zero_grad() out = model(bx) loss = criterion(out, by) loss.backward() optimizer.step() scheduler.step() model.eval() with torch.no_grad(): for bx, by in loader_te: bx, by = bx.to(device), by.to(device) preds = torch.argmax(model(bx), dim=1) acc = accuracy_score(by.cpu().numpy(), preds.cpu().numpy()) if acc > best_acc: best_acc = acc fold_accs.append(best_acc) return fold_accs # --- 5. FBCSP Evaluator --- def eval_fbcsp(X: np.ndarray, y: np.ndarray, bands=[(4, 8), (8, 12), (12, 16), (16, 24), (24, 32)], fs: int = 200): skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) accs_svm, accs_lda, accs_ridge = [], [], [] for fold, (tr_idx, te_idx) in enumerate(skf.split(X, y)): X_tr, y_tr = X[tr_idx], y[tr_idx] X_te, y_te = X[te_idx], y[te_idx] feats_tr, feats_te = [], [] for b_low, b_high in bands: sos = signal.butter(4, (b_low, b_high), btype="bandpass", fs=fs, output="sos") X_tr_b = np.array([signal.sosfiltfilt(sos, t, axis=-1) for t in X_tr]) X_te_b = np.array([signal.sosfiltfilt(sos, t, axis=-1) for t in X_te]) c0 = np.mean([np.cov(t) for t in X_tr_b[y_tr == 0]], axis=0) c1 = np.mean([np.cov(t) for t in X_tr_b[y_tr == 1]], axis=0) reg = 1e-4 * np.eye(c0.shape[0]) w, v = eigh(c0 + reg, c0 + c1 + 2 * reg) ix = np.argsort(w)[::-1] v = v[:, ix] W = v[:, [0, 1, -2, -1]].T def extract_logvar(X_in): return np.array([np.log(np.var(W @ t, axis=-1) + 1e-10) for t in X_in]) feats_tr.append(extract_logvar(X_tr_b)) feats_te.append(extract_logvar(X_te_b)) F_tr = np.hstack(feats_tr) F_te = np.hstack(feats_te) f_mean = np.mean(F_tr, axis=0) f_std = np.std(F_tr, axis=0) + 1e-6 F_tr_scaled = (F_tr - f_mean) / f_std F_te_scaled = (F_te - f_mean) / f_std clf_svm = SVC(C=1.0, kernel="rbf") clf_svm.fit(F_tr_scaled, y_tr) accs_svm.append(accuracy_score(y_te, clf_svm.predict(F_te_scaled))) clf_lda = LinearDiscriminantAnalysis() clf_lda.fit(F_tr_scaled, y_tr) accs_lda.append(accuracy_score(y_te, clf_lda.predict(F_te_scaled))) clf_ridge = RidgeClassifier(alpha=1.0) clf_ridge.fit(F_tr_scaled, y_tr) accs_ridge.append(accuracy_score(y_te, clf_ridge.predict(F_te_scaled))) return accs_svm, accs_lda, accs_ridge # --- 6. Main Execution --- def main(): package = Path(__file__).resolve().parents[1] cache_dir = package / "output_cache" print("==========================================================") print(" TCH EEG Cleaned Dataset 5-Fold CV Evaluation ") print("==========================================================") for subj_prefix, subj_name in [("s1", "Subject 1 (s1)"), ("s2", "Subject 2 (s2)")]: X, y = load_subject_cleaned_data(subj_prefix, cache_dir) print(f"\n>>> {subj_name}: Total Trials = {len(y)} (Left={(y==0).sum()}, Right={(y==1).sum()})") # 1. EEGNet Evaluation eegnet_folds = eval_deep_model(EEGNet, X, y, epochs=100) eegnet_mean = np.mean(eegnet_folds) print(f" [EEGNet Deep Learning] 5-Fold Acc: {eegnet_mean*100:.2f}% | Folds: {[round(a*100, 1) for a in eegnet_folds]}") # 2. ShallowConvNet Evaluation shallow_folds = eval_deep_model(ShallowConvNet, X, y, epochs=100, apply_car=True) shallow_mean = np.mean(shallow_folds) print(f" [ShallowConvNet Deep Learning] 5-Fold Acc: {shallow_mean*100:.2f}% | Folds: {[round(a*100, 1) for a in shallow_folds]}") # 3. FBCSP + SVM Evaluation svm_folds, lda_folds, _ = eval_fbcsp(X, y) print(f" [FBCSP + SVM (RBF)] 5-Fold Acc: {np.mean(svm_folds)*100:.2f}% | Folds: {[round(a*100, 1) for a in svm_folds]}") print(f" [FBCSP + LDA] 5-Fold Acc: {np.mean(lda_folds)*100:.2f}% | Folds: {[round(a*100, 1) for a in lda_folds]}") if __name__ == "__main__": main()