"""Baseline-conditioned soft-routing temporal classifier for event-referenced EEG.""" from __future__ import annotations import copy import json import math from dataclasses import asdict, dataclass from pathlib import Path from typing import Dict, List, Mapping, Sequence, Tuple import numpy as np import pandas as pd import torch import torch.nn as nn import torch.nn.functional as F from scipy.stats import wilcoxon from sklearn.metrics import accuracy_score, balanced_accuracy_score, f1_score from torch.utils.data import DataLoader, TensorDataset from .data import LABELED_SUBJECTS, load_fold_assignments, load_subject from .training import augment_waveforms, resolve_device, seed_everything, stable_seed @dataclass(frozen=True) class BaselineMoEConfig: device: str = "cuda" max_epochs: int = 150 patience: int = 20 batch_size_all: int = 32 batch_size_fewshot: int = 8 learning_rate: float = 3e-4 weight_decay: float = 1e-3 label_smoothing: float = 0.05 warmup_epochs: int = 5 gradient_clip: float = 1.0 amplitude_low: float = 0.95 amplitude_high: float = 1.05 max_shift_samples: int = 10 all_trial_seeds: Tuple[int, ...] = (20260720, 20260721, 20260722) def normalize_event_trial(x: torch.Tensor) -> torch.Tensor: """Per-channel normalization over the full baseline+task event window.""" mean = x.mean(dim=-1, keepdim=True) scale = x.std(dim=-1, keepdim=True, unbiased=False).clamp_min(1e-6) return (x - mean) / scale class GatedTemporalBlock(nn.Module): def __init__(self, width: int, dilation: int, state_dim: int, dropout: float) -> None: super().__init__() padding = dilation * 4 self.norm = nn.GroupNorm(8, width) self.depthwise = nn.Conv1d( width, width * 2, kernel_size=9, padding=padding, dilation=dilation, groups=width, bias=False ) self.pointwise = nn.Conv1d(width, width, kernel_size=1, bias=False) self.condition = nn.Linear(state_dim, width * 2) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor, state: torch.Tensor) -> torch.Tensor: values = self.depthwise(self.norm(x)) gate, candidate = values.chunk(2, dim=1) values = torch.sigmoid(gate) * torch.tanh(candidate) gamma, beta = self.condition(state).chunk(2, dim=-1) values = values * (1.0 + gamma.unsqueeze(-1)) + beta.unsqueeze(-1) return x + self.dropout(self.pointwise(values)) class TemporalAdapter(nn.Module): """A small expert: all trials retain the shared representation and receive a soft residual.""" def __init__(self, width: int, dilation: int, dropout: float) -> None: super().__init__() self.norm = nn.GroupNorm(8, width) self.depthwise = nn.Conv1d( width, width, kernel_size=7, padding=dilation * 3, dilation=dilation, groups=width, bias=False ) self.pointwise = nn.Conv1d(width, width, kernel_size=1, bias=False) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.dropout(self.pointwise(F.gelu(self.depthwise(self.norm(x))))) class BaselineRouter(nn.Module): def __init__(self, channels: int = 8, state_dim: int = 64, experts: int = 2) -> None: super().__init__() self.features = nn.Sequential( nn.Conv1d(channels, 32, kernel_size=25, padding=12, groups=8, bias=False), nn.GroupNorm(8, 32), nn.GELU(), nn.Conv1d(32, 48, kernel_size=15, padding=7, bias=False), nn.GroupNorm(8, 48), nn.GELU(), nn.AdaptiveAvgPool1d(1), ) self.state = nn.Sequential(nn.Flatten(), nn.Linear(48, state_dim), nn.GELU()) self.gates = nn.Linear(state_dim, experts) def forward(self, baseline: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: state = self.state(self.features(baseline)) return state, torch.softmax(self.gates(state), dim=-1) class BaselineConditionedMoETCN(nn.Module): """Raw-waveform TCN with baseline-only soft routing into two temporal adapters.""" baseline_samples: int = 200 def __init__(self, channels: int = 8, width: int = 64, state_dim: int = 64, experts: int = 2) -> None: super().__init__() if channels != 8 or experts != 2: raise ValueError("This experiment fixes the eight-channel montage and two soft adapters") self.router = BaselineRouter(channels, state_dim, experts) self.stem = nn.Sequential( nn.Conv1d(channels, 32, kernel_size=25, padding=12, groups=channels, bias=False), nn.GroupNorm(8, 32), nn.GELU(), nn.Conv1d(32, width, kernel_size=1, bias=False), nn.GroupNorm(8, width), nn.GELU(), ) self.blocks = nn.ModuleList( [ GatedTemporalBlock(width, dilation=1, state_dim=state_dim, dropout=0.15), GatedTemporalBlock(width, dilation=2, state_dim=state_dim, dropout=0.15), GatedTemporalBlock(width, dilation=4, state_dim=state_dim, dropout=0.15), ] ) self.adapters = nn.ModuleList( [TemporalAdapter(width, dilation=1, dropout=0.20), TemporalAdapter(width, dilation=4, dropout=0.20)] ) self.attention = nn.Sequential(nn.Conv1d(width, width // 2, 1), nn.Tanh(), nn.Conv1d(width // 2, 1, 1)) self.classifier = nn.Sequential( nn.LayerNorm(width * 3), nn.Dropout(0.35), nn.Linear(width * 3, 64), nn.GELU(), nn.Dropout(0.20), nn.Linear(64, 2), ) def forward(self, x: torch.Tensor, return_routing: bool = False): if x.ndim != 3 or tuple(x.shape[1:]) != (8, 600): raise ValueError(f"Expected (batch, 8, 600), got {tuple(x.shape)}") state, gates = self.router(x[:, :, : self.baseline_samples]) values = self.stem(x[:, :, self.baseline_samples :]) for block in self.blocks: values = block(values, state) adapted = torch.stack([adapter(values) for adapter in self.adapters], dim=1) values = values + (adapted * gates[:, :, None, None]).sum(dim=1) scores = torch.softmax(self.attention(values), dim=-1) attentive = (values * scores).sum(dim=-1) mean = values.mean(dim=-1) log_std = torch.log(values.var(dim=-1, unbiased=False).add(1e-6)) logits = self.classifier(torch.cat((attentive, mean, log_std), dim=-1)) return (logits, gates) if return_routing else logits def parameter_count(model: nn.Module) -> int: return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad) def _loader(x: np.ndarray, y: np.ndarray, indices: np.ndarray, batch_size: int, shuffle: bool, seed: int) -> DataLoader: dataset = TensorDataset(torch.from_numpy(np.asarray(x[indices], dtype=np.float32)), torch.from_numpy(np.asarray(y[indices], dtype=np.int64))) generator = torch.Generator().manual_seed(seed) return DataLoader(dataset, batch_size=min(batch_size, len(dataset)), shuffle=shuffle, num_workers=0, generator=generator) def _scheduler(epoch: int, max_epochs: int, warmup: int) -> float: if epoch < warmup: return float(epoch + 1) / max(warmup, 1) progress = (epoch - warmup) / max(max_epochs - warmup, 1) return 0.5 * (1.0 + math.cos(math.pi * min(progress, 1.0))) def _evaluate(model: BaselineConditionedMoETCN, loader: DataLoader, device: torch.device, criterion: nn.Module) -> Dict[str, object]: model.eval() total_loss = 0.0 labels: List[np.ndarray] = [] predictions: List[np.ndarray] = [] gates: List[np.ndarray] = [] with torch.no_grad(): for x, y in loader: x = normalize_event_trial(x.to(device)) y = y.to(device) logits, routing = model(x, return_routing=True) total_loss += float(criterion(logits, y).item()) * len(y) labels.append(y.cpu().numpy()) predictions.append(logits.argmax(dim=-1).cpu().numpy()) gates.append(routing.cpu().numpy()) actual = np.concatenate(labels) predicted = np.concatenate(predictions) route = np.concatenate(gates) return { "loss": total_loss / len(actual), "accuracy": float(accuracy_score(actual, predicted)), "balanced_accuracy": float(balanced_accuracy_score(actual, predicted)), "macro_f1": float(f1_score(actual, predicted, average="macro", zero_division=0)), "route_expert1_mean": float(route[:, 0].mean()), "route_entropy": float((-(route * np.log(route + 1e-8)).sum(axis=1)).mean()), } def _train_fold(x: np.ndarray, y: np.ndarray, split: Mapping[str, np.ndarray], seed: int, config: BaselineMoEConfig, checkpoint: Path | None) -> Dict[str, object]: seed_everything(seed) device = resolve_device(config.device) model = BaselineConditionedMoETCN().to(device) total = sum(len(split[key]) for key in split) batch_size = config.batch_size_fewshot if total <= 40 else config.batch_size_all loaders = {key: _loader(x, y, values, batch_size, key == "train", stable_seed(seed, key)) for key, values in split.items()} criterion = nn.CrossEntropyLoss(label_smoothing=config.label_smoothing) optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda epoch: _scheduler(epoch, config.max_epochs, config.warmup_epochs)) best_state = None best_score, best_loss, best_epoch, stale = -float("inf"), float("inf"), 0, 0 history: List[Dict[str, float]] = [] for epoch in range(1, config.max_epochs + 1): model.train() for batch_x, batch_y in loaders["train"]: batch_x = normalize_event_trial(batch_x.to(device)) batch_x = augment_waveforms(batch_x, config) batch_y = batch_y.to(device) optimizer.zero_grad(set_to_none=True) loss = criterion(model(batch_x), batch_y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip) optimizer.step() validation = _evaluate(model, loaders["val"], device, criterion) scheduler.step() history.append({"epoch": epoch, "val_loss": validation["loss"], "val_balanced_accuracy": validation["balanced_accuracy"]}) score, loss_value = validation["balanced_accuracy"], validation["loss"] if score > best_score + 1e-8 or (abs(score - best_score) <= 1e-8 and loss_value < best_loss - 1e-8): best_state = copy.deepcopy(model.state_dict()) best_score, best_loss, best_epoch, stale = score, loss_value, epoch, 0 else: stale += 1 if stale >= config.patience: break if best_state is None: raise RuntimeError("No valid baseline-conditioned checkpoint") model.load_state_dict(best_state) validation = _evaluate(model, loaders["val"], device, criterion) test = _evaluate(model, loaders["test"], device, criterion) result = { "model": "baseline_soft_moe_tcn", "representation": "raw_time_baseline_conditioned", "fold_seed": seed, "train_size": len(split["train"]), "val_size": len(split["val"]), "test_size": len(split["test"]), "best_epoch": best_epoch, "epochs_ran": len(history), "parameter_count": parameter_count(model), **{f"val_{key}": value for key, value in validation.items()}, **test, } if checkpoint is not None: checkpoint.parent.mkdir(parents=True, exist_ok=True) torch.save({"model_state": best_state, "config": asdict(config), "result": result, "history": history}, checkpoint) return result def _summary(folds: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]: subject = folds.groupby(["protocol", "model", "subject_id"], as_index=False).agg( mean_accuracy=("accuracy", "mean"), std_accuracy=("accuracy", "std"), mean_balanced_accuracy=("balanced_accuracy", "mean"), mean_macro_f1=("macro_f1", "mean"), mean_route_expert1=("route_expert1_mean", "mean"), mean_route_entropy=("route_entropy", "mean"), ) group = subject.groupby(["protocol", "model"], as_index=False).agg( subjects=("subject_id", "nunique"), mean_accuracy=("mean_accuracy", "mean"), std_subject_accuracy=("mean_accuracy", "std"), mean_balanced_accuracy=("mean_balanced_accuracy", "mean"), std_subject_balanced_accuracy=("mean_balanced_accuracy", "std"), mean_macro_f1=("mean_macro_f1", "mean"), mean_route_expert1=("mean_route_expert1", "mean"), mean_route_entropy=("mean_route_entropy", "mean"), ) return subject, group def _paired(new_subject: pd.DataFrame, shallow_subject: pd.DataFrame) -> pd.DataFrame: rows=[] rng=np.random.default_rng(20260720) for protocol in ("all", "20x20"): a=new_subject[new_subject.protocol==protocol].set_index("subject_id")["mean_accuracy"] b=shallow_subject[shallow_subject.protocol==protocol].set_index("subject_id")["mean_accuracy"] common=a.index.intersection(b.index) diff=(a.loc[common]-b.loc[common]).to_numpy() boot=np.array([rng.choice(diff, size=len(diff), replace=True).mean() for _ in range(10000)]) try: p=float(wilcoxon(diff, alternative="two-sided", method="auto").pvalue) except ValueError: p=float("nan") rows.append({"protocol":protocol,"comparison":"Baseline soft-MoE TCN - ShallowConvNet task","subjects":len(diff),"mean_difference_pp":100*float(diff.mean()),"bootstrap_ci_low_pp":100*float(np.quantile(boot,.025)),"bootstrap_ci_high_pp":100*float(np.quantile(boot,.975)),"wilcoxon_p":p,"improved_subjects":int((diff>0).sum())}) return pd.DataFrame(rows) def run_outer(cache: Path, run_dir: Path, config: BaselineMoEConfig, subjects: Sequence[str] = LABELED_SUBJECTS) -> Dict[str, Path]: run_dir.mkdir(parents=True, exist_ok=True) assignments=load_fold_assignments(cache/"fold_assignments.csv") rows=[] for protocol in ("all", "20x20"): split_sets=["all"] if protocol=="all" else [f"20_seed_{seed}" for seed in range(20260713,20260718)] train_seeds=config.all_trial_seeds if protocol=="all" else (20260720,) for split_set in split_sets: for train_seed in train_seeds: for subject in subjects: x,y=load_subject(cache,subject) for fold in range(1,6): split=assignments[(split_set,subject,fold)] fold_seed=stable_seed("baseline-soft-moe",protocol,split_set,subject,fold,train_seed) print(f"[outer] {protocol} {split_set} {subject} fold={fold}",flush=True) result=_train_fold(x,y,split,fold_seed,config,run_dir/"checkpoints"/subject/f"{protocol}_{split_set}_fold{fold}.pt") rows.append({"protocol":protocol,"split_set":split_set,"train_seed":train_seed,"subject_id":subject,"fold":fold,**result}) folds=pd.DataFrame(rows) expected=10*5*(len(config.all_trial_seeds)+5) if len(folds)!=expected: raise RuntimeError(f"Expected {expected} fold results, received {len(folds)}") subject,group=_summary(folds) shallow=pd.read_csv(Path(__file__).resolve().parents[1]/"runs"/"fold_results.csv") shallow=shallow[shallow.model=="shallow_task"].copy() shallow_subject=shallow.groupby(["protocol","subject_id"],as_index=False).agg(mean_accuracy=("accuracy","mean")) paired=_paired(subject,shallow_subject) folds.to_csv(run_dir/"fold_results.csv",index=False,encoding="utf-8-sig") subject.to_csv(run_dir/"subject_summary.csv",index=False,encoding="utf-8-sig") group.to_csv(run_dir/"group_summary.csv",index=False,encoding="utf-8-sig") paired.to_csv(run_dir/"paired_comparisons.csv",index=False,encoding="utf-8-sig") (run_dir/"run_config.json").write_text(json.dumps(asdict(config),indent=2),encoding="utf-8") return {"folds":run_dir/"fold_results.csv","subject":run_dir/"subject_summary.csv","group":run_dir/"group_summary.csv","paired":run_dir/"paired_comparisons.csv"}