"""Leakage-safe physiology routing and confidence-aware Shallow experts.""" from __future__ import annotations import copy import json import math import warnings 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 from scipy.signal import welch from scipy.stats import wilcoxon from sklearn.cluster import KMeans from sklearn.decomposition import PCA from sklearn.metrics import ( accuracy_score, adjusted_rand_score, balanced_accuracy_score, confusion_matrix, f1_score, roc_auc_score, silhouette_score, ) from sklearn.preprocessing import StandardScaler from .data import LABELED_SUBJECTS, SUBSAMPLE_SEEDS, load_fold_assignments, load_subject from .kmeans_shallow import baseline_features as current_baseline_features from .kmeans_shallow import fallback_reason from .models import ShallowConvNet, count_parameters from .training import ( TrainConfig, _scheduler_lambda, augment_waveforms, evaluate_loader, make_loader, prepare_model_input, resolve_device, seed_everything, stable_seed, train_fold, ) MODEL_GLOBAL = "global_shallow" MODEL_HARD = "optimized_hard_k2" MODEL_SOFT = "optimized_soft_k2" FINAL_MODELS = (MODEL_GLOBAL, MODEL_HARD, MODEL_SOFT) FEATURE_SETS = ("current", "corrected", "motor_focused", "motor_plus_lowfreq") SCREEN_SPLITS = tuple(f"20_seed_{seed}" for seed in SUBSAMPLE_SEEDS[:2]) @dataclass(frozen=True) class OptimizedRoutingConfig: device: str = "cuda" seed: int = 20260720 all_trial_seeds: Tuple[int, ...] = (20260720, 20260721, 20260722) kmeans_clusters: int = 2 kmeans_n_init: int = 50 pca_components: int = 8 global_epochs: int = 150 global_patience: int = 20 expert_epochs: int = 50 expert_patience: int = 10 weight_decay: float = 1e-3 label_smoothing: float = 0.05 warmup_epochs: int = 5 batch_size_all: int = 32 batch_size_fewshot: int = 8 gradient_clip: float = 1.0 amplitude_low: float = 0.95 amplitude_high: float = 1.05 max_shift_samples: int = 10 bootstrap_router_repeats: int = 20 tau_candidates: Tuple[float, ...] = (0.25, 0.5, 1.0, 2.0) gamma_candidates: Tuple[float, ...] = (0.25, 0.5, 0.75, 1.0) def training_config(self) -> TrainConfig: return TrainConfig( device=self.device, seed=self.seed, all_trial_seeds=self.all_trial_seeds, max_epochs=self.global_epochs, patience=self.global_patience, batch_size_all=self.batch_size_all, batch_size_fewshot=self.batch_size_fewshot, learning_rate=3e-4, weight_decay=self.weight_decay, label_smoothing=self.label_smoothing, warmup_epochs=self.warmup_epochs, amplitude_low=self.amplitude_low, amplitude_high=self.amplitude_high, max_shift_samples=self.max_shift_samples, gradient_clip=self.gradient_clip, ) @dataclass(frozen=True) class AdaptationSpec: name: str temporal_lr: float spatial_lr: float batchnorm_lr: float head_lr: float l2sp_lambda: float = 0.0 ADAPTATION_SPECS = ( AdaptationSpec("uniform_full", 1e-4, 1e-4, 1e-4, 1e-4), AdaptationSpec("head_only", 0.0, 0.0, 0.0, 5e-4), AdaptationSpec("spatial_head", 0.0, 1e-4, 1e-4, 5e-4), AdaptationSpec("full_l2_1e4", 2e-5, 1e-4, 1e-4, 5e-4, 1e-4), AdaptationSpec("full_l2_1e3", 2e-5, 1e-4, 1e-4, 5e-4, 1e-3), AdaptationSpec("full_l2_1e2", 2e-5, 1e-4, 1e-4, 5e-4, 1e-2), ) SPEC_BY_NAME = {spec.name: spec for spec in ADAPTATION_SPECS} SPEC_COMPLEXITY = { "head_only": 0, "spatial_head": 1, "full_l2_1e4": 2, "full_l2_1e3": 2, "full_l2_1e2": 2, "uniform_full": 3, } FEATURE_COMPLEXITY = {name: index for index, name in enumerate(FEATURE_SETS)} def _baseline_arrays( x: np.ndarray, sample_rate: int = 200, baseline_samples: int = 200 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: baseline = np.asarray(x[:, :, :baseline_samples], dtype=np.float64) centered = baseline - baseline.mean(axis=-1, keepdims=True) frequencies, power = welch(centered, fs=sample_rate, nperseg=128, axis=-1) return baseline, frequencies, power def _band_mean(power: np.ndarray, frequencies: np.ndarray, low: float, high: float) -> np.ndarray: mask = (frequencies >= low) & (frequencies <= high) return power[:, :, mask].mean(axis=-1) + 1e-12 def _fisher_correlations(baseline: np.ndarray) -> np.ndarray: correlations = np.asarray([np.corrcoef(trial) for trial in baseline]) correlations = np.nan_to_num(correlations, nan=0.0, posinf=0.0, neginf=0.0) correlations = np.arctanh(np.clip(correlations, -0.995, 0.995)) triangle = np.triu_indices(baseline.shape[1], k=1) return correlations[:, triangle[0], triangle[1]] def routing_features(x: np.ndarray, feature_set: str) -> np.ndarray: """Return label-free event-preceding state features for one feature definition.""" if feature_set not in FEATURE_SETS: raise ValueError(f"Unknown feature set: {feature_set}") if feature_set == "current": return current_baseline_features(x) baseline, frequencies, power = _baseline_arrays(x) centered = baseline - baseline.mean(axis=-1, keepdims=True) log_sd = np.log(centered.std(axis=-1) + 1e-12) correlations = _fisher_correlations(centered) total = _band_mean(power, frequencies, 1.0, 40.0) if feature_set == "corrected": relative = [ np.log(_band_mean(power, frequencies, low, high) / total) for low, high in ((1, 4), (4, 8), (8, 13), (13, 20), (20, 30), (30, 40)) ] values = np.concatenate([log_sd, *relative, correlations], axis=1) else: bands = ((8, 13), (13, 20), (20, 30)) relative = [np.log(_band_mean(power, frequencies, low, high) / total) for low, high in bands] absolute = [np.log10(_band_mean(power, frequencies, low, high)) for low, high in bands] c3, c4 = 3, 4 central_absolute = np.concatenate( [values[:, [c3, c4]] for values in absolute], axis=1 ) asymmetry = np.stack([values[:, c3] - values[:, c4] for values in absolute], axis=1) parts = [log_sd, *relative, central_absolute, asymmetry, correlations] if feature_set == "motor_plus_lowfreq": central = (2, 3, 4, 5) low_power = np.log10(_band_mean(power, frequencies, 1.0, 4.0)[:, central]) time = np.linspace(-0.5, 0.5, baseline.shape[-1], dtype=np.float64) denominator = np.sum(time * time) slopes = np.einsum("t,nct->nc", time, centered[:, central, :]) / denominator slopes = slopes / (centered[:, central, :].std(axis=-1) + 1e-12) parts.extend([low_power, slopes]) values = np.concatenate(parts, axis=1) if not np.isfinite(values).all(): raise RuntimeError(f"Non-finite {feature_set} routing features") return values.astype(np.float32) class PhysiologyKMeansRouter: """Train-only scaler/PCA/KMeans router with normalized soft distances.""" def __init__(self, config: OptimizedRoutingConfig, seed: int) -> None: self.config = config self.seed = seed self.scaler = StandardScaler() self.pca: PCA | None = None self.kmeans = KMeans( n_clusters=config.kmeans_clusters, n_init=config.kmeans_n_init, random_state=seed, ) self.fitted_samples = 0 self.distance_scale = float("nan") self.train_embedding: np.ndarray | None = None self.train_labels: np.ndarray | None = None def fit(self, features: np.ndarray) -> "PhysiologyKMeansRouter": standardized = self.scaler.fit_transform(features) components = min(self.config.pca_components, standardized.shape[1], len(features) - 2) if components < 1: raise ValueError("Insufficient samples for PCA") self.pca = PCA(n_components=components, whiten=True, random_state=self.seed) embedded = self.pca.fit_transform(standardized) labels = self.kmeans.fit_predict(embedded) distances = self.kmeans.transform(embedded) self.distance_scale = max(float(np.median(np.min(distances, axis=1))), 1e-6) self.fitted_samples = len(features) self.train_embedding = embedded self.train_labels = labels.astype(np.int64) return self @property def pca_components(self) -> int: if self.pca is None: raise RuntimeError("Router has not been fitted") return int(self.pca.n_components_) def transform(self, features: np.ndarray) -> np.ndarray: if self.pca is None: raise RuntimeError("Router has not been fitted") return self.pca.transform(self.scaler.transform(features)) def normalized_distances(self, features: np.ndarray) -> np.ndarray: return self.kmeans.transform(self.transform(features)) / self.distance_scale def predict(self, features: np.ndarray) -> np.ndarray: return self.kmeans.predict(self.transform(features)).astype(np.int64) def soft_weights(self, features: np.ndarray, tau: float) -> np.ndarray: distances = self.normalized_distances(features) logits = -distances / float(tau) logits -= logits.max(axis=1, keepdims=True) weights = np.exp(logits) return weights / weights.sum(axis=1, keepdims=True) def silhouette(self) -> float: if self.train_embedding is None or self.train_labels is None: raise RuntimeError("Router has not been fitted") counts = np.bincount(self.train_labels, minlength=self.config.kmeans_clusters) if (counts < 2).any(): return float("nan") return float(silhouette_score(self.train_embedding, self.train_labels)) def bootstrap_ari(self) -> float: if self.train_embedding is None or self.train_labels is None: raise RuntimeError("Router has not been fitted") rng = np.random.default_rng(stable_seed("router-bootstrap", self.seed)) scores = [] for repeat in range(self.config.bootstrap_router_repeats): indices = rng.integers(0, len(self.train_embedding), len(self.train_embedding)) candidate = KMeans( n_clusters=self.config.kmeans_clusters, n_init=10, random_state=stable_seed(self.seed, repeat), ).fit(self.train_embedding[indices]) scores.append(adjusted_rand_score(self.train_labels, candidate.predict(self.train_embedding))) return float(np.mean(scores)) def _criterion(labels: np.ndarray, config: OptimizedRoutingConfig, device: torch.device) -> nn.Module: counts = np.bincount(labels, minlength=2).astype(np.float64) if (counts == 0).any(): raise ValueError("Cannot train expert with a missing class") weights = counts.sum() / (2.0 * counts) return nn.CrossEntropyLoss( weight=torch.tensor(weights, dtype=torch.float32, device=device), label_smoothing=config.label_smoothing, ) def _optimizer_groups(model: ShallowConvNet, spec: AdaptationSpec) -> List[Dict[str, object]]: for parameter in model.parameters(): parameter.requires_grad = False groups: List[Dict[str, object]] = [] definitions = ( (model.features[0], spec.temporal_lr, "temporal"), (model.features[1], spec.spatial_lr, "spatial"), (model.features[2], spec.batchnorm_lr, "batchnorm"), (model.classifier, spec.head_lr, "head"), ) for module, learning_rate, name in definitions: if learning_rate <= 0: continue parameters = list(module.parameters()) for parameter in parameters: parameter.requires_grad = True groups.append({"params": parameters, "lr": learning_rate, "name": name}) if not groups: raise RuntimeError(f"Adaptation {spec.name} has no trainable parameters") return groups def _l2sp_penalty( model: ShallowConvNet, anchors: Mapping[str, torch.Tensor] ) -> torch.Tensor: penalty = torch.zeros((), device=next(model.parameters()).device) for name in ("features.0.weight", "features.1.weight"): penalty = penalty + torch.sum(torch.square(dict(model.named_parameters())[name] - anchors[name])) return penalty def fit_adapted_expert( x: np.ndarray, y: np.ndarray, train_indices: np.ndarray, validation_indices: np.ndarray, global_state: Mapping[str, torch.Tensor], spec: AdaptationSpec, config: OptimizedRoutingConfig, seed: int, ) -> Tuple[Dict[str, torch.Tensor], Dict[str, object]]: seed_everything(seed) device = resolve_device(config.device) model = ShallowConvNet(channels=8, samples=400).to(device) model.load_state_dict(global_state) groups = _optimizer_groups(model, spec) anchors = { name: dict(model.named_parameters())[name].detach().clone() for name in ("features.0.weight", "features.1.weight") } optimizer = torch.optim.AdamW(groups, weight_decay=config.weight_decay) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lambda epoch: _scheduler_lambda(epoch, config.expert_epochs, config.warmup_epochs), ) batch_size = ( config.batch_size_fewshot if len(train_indices) + len(validation_indices) <= 40 else config.batch_size_all ) train_loader = make_loader( x, y, train_indices, batch_size, True, stable_seed(seed, "train"), device.type == "cuda" ) validation_loader = make_loader( x, y, validation_indices, batch_size, False, stable_seed(seed, "val"), device.type == "cuda" ) criterion = _criterion(y[train_indices], config, device) train_config = config.training_config() best_state = None best_score = -float("inf") best_loss = float("inf") best_epoch = 0 stale = 0 initial_l2sp = float(_l2sp_penalty(model, anchors).detach().cpu()) validation_has_both = len(np.unique(y[validation_indices])) == 2 for epoch in range(1, config.expert_epochs + 1): model.train() if spec.name != "uniform_full": model.features[2].eval() for raw, target in train_loader: raw = raw.to(device, non_blocking=True) target = target.to(device, non_blocking=True) values = prepare_model_input(raw, "shallow_task", train_config, training=False) values = augment_waveforms(values, train_config) optimizer.zero_grad(set_to_none=True) loss = criterion(model(values), target) if spec.l2sp_lambda > 0: loss = loss + spec.l2sp_lambda * _l2sp_penalty(model, anchors) loss.backward() torch.nn.utils.clip_grad_norm_( [parameter for parameter in model.parameters() if parameter.requires_grad], config.gradient_clip, ) optimizer.step() with warnings.catch_warnings(): warnings.simplefilter("ignore") validation = evaluate_loader( model, validation_loader, "shallow_task", train_config, device, criterion ) scheduler.step() score = ( float(validation["balanced_accuracy"]) if validation_has_both else -float(validation["loss"]) ) loss_value = float(validation["loss"]) improved = score > best_score + 1e-8 or ( abs(score - best_score) <= 1e-8 and loss_value < best_loss - 1e-8 ) if improved: best_state = {key: value.detach().cpu().clone() for key, value in model.state_dict().items()} best_score, best_loss, best_epoch, stale = score, loss_value, epoch, 0 else: stale += 1 if stale >= config.expert_patience: break if best_state is None: raise RuntimeError("Expert training produced no checkpoint") return best_state, { "best_epoch": best_epoch, "epochs_ran": epoch, "validation_loss": best_loss, "validation_has_both_classes": validation_has_both, "trainable_parameters": sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad), "total_parameters": sum(parameter.numel() for parameter in model.parameters()), "initial_l2sp": initial_l2sp, "final_l2sp": float(_l2sp_penalty(model, anchors).detach().cpu()), } def _predict_state( state: Mapping[str, torch.Tensor], x: np.ndarray, y: np.ndarray, indices: np.ndarray, config: OptimizedRoutingConfig, ) -> Tuple[np.ndarray, np.ndarray]: device = resolve_device(config.device) model = ShallowConvNet(channels=8, samples=400).to(device) model.load_state_dict(state) loader = make_loader( x, y, indices, config.batch_size_all, False, stable_seed("predict", *indices.tolist()), device.type == "cuda" ) result = evaluate_loader( model, loader, "shallow_task", config.training_config(), device, nn.CrossEntropyLoss() ) return result["labels"], result["probabilities"] def classification_metrics(labels: np.ndarray, probabilities: np.ndarray) -> Dict[str, object]: predictions = (probabilities >= 0.5).astype(np.int64) matrix = confusion_matrix(labels, predictions, labels=(0, 1)) try: auc = float(roc_auc_score(labels, probabilities)) except ValueError: auc = float("nan") return { "accuracy": float(accuracy_score(labels, predictions)), "balanced_accuracy": float(balanced_accuracy_score(labels, predictions)), "macro_f1": float(f1_score(labels, predictions, average="macro", zero_division=0)), "roc_auc": auc, "tn": int(matrix[0, 0]), "fp": int(matrix[0, 1]), "fn": int(matrix[1, 0]), "tp": int(matrix[1, 1]), } def confidence_soft_probabilities( global_probability: np.ndarray, expert_probabilities: np.ndarray, normalized_distances: np.ndarray, tau: float, gamma: float, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: logits = -np.asarray(normalized_distances, dtype=np.float64) / float(tau) logits -= logits.max(axis=1, keepdims=True) weights = np.exp(logits) weights /= weights.sum(axis=1, keepdims=True) certainty = 2.0 * weights.max(axis=1) - 1.0 expert_mix = np.sum(weights * expert_probabilities, axis=1) blend = float(gamma) * certainty final = (1.0 - blend) * global_probability + blend * expert_mix return final.astype(np.float32), weights.astype(np.float32), certainty.astype(np.float32) def _save(path: Path, frame: pd.DataFrame) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") frame.to_csv(temporary, index=False, encoding="utf-8-sig") temporary.replace(path) def _load_selection(run_dir: Path) -> Dict[str, object]: path = run_dir / "selection.json" return json.loads(path.read_text(encoding="utf-8")) if path.exists() else {} def _write_selection(run_dir: Path, selection: Mapping[str, object]) -> None: run_dir.mkdir(parents=True, exist_ok=True) (run_dir / "selection.json").write_text( json.dumps(dict(selection), ensure_ascii=False, indent=2), encoding="utf-8" ) def _global_checkpoint( run_dir: Path, split_set: str, train_seed: int, subject: str, fold: int ) -> Path: return ( run_dir / "screen_checkpoints" / subject / f"{split_set}_seed{train_seed}_fold{fold}_global.pt" ) def _load_or_train_global( x: np.ndarray, y: np.ndarray, split: Mapping[str, np.ndarray], split_set: str, train_seed: int, subject: str, fold: int, run_dir: Path, config: OptimizedRoutingConfig, evaluate_test: bool, ) -> Tuple[Mapping[str, torch.Tensor], Dict[str, object]]: checkpoint_path = _global_checkpoint(run_dir, split_set, train_seed, subject, fold) expected_test = bool(evaluate_test) if checkpoint_path.exists(): checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) result = dict(checkpoint["result"]) has_test = "accuracy" in result if has_test == expected_test or (has_test and not expected_test): return checkpoint["model_state"], result seed = stable_seed("outer", "shallow_task", "waveform", split_set, train_seed, subject, fold) result = train_fold( x, y, split, "shallow_task", "waveform", seed, config.training_config(), config.global_epochs, config.global_patience, evaluate_test=evaluate_test, checkpoint_path=checkpoint_path, ) checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) return checkpoint["model_state"], result def _cluster_indices( split: Mapping[str, np.ndarray], routed: Mapping[str, np.ndarray], cluster: int ) -> Tuple[np.ndarray, np.ndarray, str]: train = split["train"][routed["train"] == cluster] validation = split["val"][routed["val"] == cluster] return train, validation, "" def _fit_experts( x: np.ndarray, y: np.ndarray, split: Mapping[str, np.ndarray], routed: Mapping[str, np.ndarray], global_state: Mapping[str, torch.Tensor], spec: AdaptationSpec, config: OptimizedRoutingConfig, seed_parts: Sequence[object], ) -> Tuple[Dict[int, Mapping[str, torch.Tensor]], Dict[int, str], List[Dict[str, object]]]: states: Dict[int, Mapping[str, torch.Tensor]] = {} fallbacks: Dict[int, str] = {} metadata: List[Dict[str, object]] = [] for cluster in range(config.kmeans_clusters): train_indices, validation_indices, _ = _cluster_indices(split, routed, cluster) reason = fallback_reason(y[train_indices], y[validation_indices]) if reason: fallbacks[cluster] = reason states[cluster] = global_state metadata.append( {"cluster": cluster, "fallback_reason": reason, "best_epoch": 0, "epochs_ran": 0, "validation_loss": float("nan"), "validation_has_both_classes": False, "trainable_parameters": 0, "total_parameters": count_parameters(ShallowConvNet(channels=8, samples=400)), "initial_l2sp": 0.0, "final_l2sp": 0.0} ) continue state, meta = fit_adapted_expert( x, y, train_indices, validation_indices, global_state, spec, config, stable_seed("expert", "hard_k2_global_init", *seed_parts, cluster), ) states[cluster] = state metadata.append({"cluster": cluster, "fallback_reason": "", **meta}) return states, fallbacks, metadata def _predict_bundle( states: Mapping[int, Mapping[str, torch.Tensor]], global_state: Mapping[str, torch.Tensor], x: np.ndarray, y: np.ndarray, indices: np.ndarray, clusters: np.ndarray, distances: np.ndarray, config: OptimizedRoutingConfig, tau: float, gamma: float, ) -> Tuple[Dict[str, Dict[str, object]], Dict[str, np.ndarray]]: labels, global_probability = _predict_state(global_state, x, y, indices, config) expert_probability = np.empty((len(indices), config.kmeans_clusters), dtype=np.float32) for cluster in range(config.kmeans_clusters): _, expert_probability[:, cluster] = _predict_state(states[cluster], x, y, indices, config) hard_probability = expert_probability[np.arange(len(indices)), clusters] soft_probability, weights, certainty = confidence_soft_probabilities( global_probability, expert_probability, distances, tau, gamma ) metrics = { MODEL_GLOBAL: classification_metrics(labels, global_probability), MODEL_HARD: classification_metrics(labels, hard_probability), MODEL_SOFT: classification_metrics(labels, soft_probability), } arrays = { "labels": labels, "global_probability": global_probability, "expert_probability": expert_probability, "hard_probability": hard_probability, "soft_probability": soft_probability, "weights": weights, "certainty": certainty, } return metrics, arrays def _selection_summary(frame: pd.DataFrame, candidate: str) -> pd.DataFrame: subject = frame.groupby([candidate, "subject_id"], as_index=False).agg( mean_validation_bacc=("balanced_accuracy", "mean"), mean_validation_accuracy=("accuracy", "mean"), mean_fallback_clusters=("fallback_clusters", "mean"), ) return subject.groupby(candidate, as_index=False).agg( subjects=("subject_id", "nunique"), mean_validation_bacc=("mean_validation_bacc", "mean"), mean_validation_accuracy=("mean_validation_accuracy", "mean"), mean_fallback_clusters=("mean_fallback_clusters", "mean"), ) def _pick_with_tolerance( summary: pd.DataFrame, candidate: str, complexity: Mapping[str, int] ) -> str: maximum = float(summary.mean_validation_bacc.max()) eligible = summary[summary.mean_validation_bacc >= maximum - 0.005].copy() eligible["complexity"] = eligible[candidate].map(complexity) eligible = eligible.sort_values(["complexity", "mean_validation_bacc"], ascending=[True, False]) return str(eligible.iloc[0][candidate]) def run_adaptation_screen( cache: Path, run_dir: Path, config: OptimizedRoutingConfig | None = None, subjects: Sequence[str] = LABELED_SUBJECTS, max_new_candidates: int | None = None, ) -> Path: config = config or OptimizedRoutingConfig() output = run_dir / "adaptation_screen.csv" frame = pd.read_csv(output) if output.exists() else pd.DataFrame() complete = set() if not frame.empty: complete = set(zip(frame.split_set, frame.subject_id, frame.fold.astype(int), frame.adaptation)) rows = frame.to_dict("records") if not frame.empty else [] assignments = load_fold_assignments(cache / "fold_assignments.csv") new_candidates = 0 for split_set in SCREEN_SPLITS: for subject in subjects: x, y = load_subject(cache, subject) features = routing_features(x, "current") for fold in range(1, 6): split = assignments[(split_set, subject, fold)] router_seed = config.seed router = PhysiologyKMeansRouter(config, router_seed).fit(features[split["train"]]) routed = {role: router.predict(features[indices]) for role, indices in split.items()} global_state, _ = _load_or_train_global( x, y, split, split_set, config.seed, subject, fold, run_dir, config, False ) for spec in ADAPTATION_SPECS: key = (split_set, subject, fold, spec.name) if key in complete: continue print(f"[adapt-screen] {split_set} {subject} fold={fold} {spec.name}", flush=True) states, fallbacks, metadata = _fit_experts( x, y, split, routed, global_state, spec, config, (split_set, config.seed, subject, fold), ) distances = router.normalized_distances(features[split["val"]]) metrics, _ = _predict_bundle( states, global_state, x, y, split["val"], routed["val"], distances, config, tau=1.0, gamma=1.0, ) rows.append({ "split_set": split_set, "subject_id": subject, "fold": fold, "adaptation": spec.name, "feature_set": "current", "fallback_clusters": len(fallbacks), "trainable_parameters": max(meta["trainable_parameters"] for meta in metadata), **metrics[MODEL_HARD], }) _save(output, pd.DataFrame(rows)) complete.add(key) new_candidates += 1 if max_new_candidates is not None and new_candidates >= max_new_candidates: return output frame = pd.DataFrame(rows) summary = _selection_summary(frame, "adaptation") selected = _pick_with_tolerance(summary, "adaptation", SPEC_COMPLEXITY) _save(run_dir / "adaptation_screen_summary.csv", summary) selection = _load_selection(run_dir) selection.update({"adaptation": selected, "adaptation_selection_metric": "subject-mean validation balanced accuracy"}) _write_selection(run_dir, selection) return output def run_feature_screen( cache: Path, run_dir: Path, config: OptimizedRoutingConfig | None = None, subjects: Sequence[str] = LABELED_SUBJECTS, max_new_candidates: int | None = None, ) -> Path: config = config or OptimizedRoutingConfig() selection = _load_selection(run_dir) if "adaptation" not in selection: raise RuntimeError("Run adaptation screen first") spec = SPEC_BY_NAME[str(selection["adaptation"])] output = run_dir / "feature_screen.csv" prediction_path = run_dir / "feature_validation_predictions.csv" frame = pd.read_csv(output) if output.exists() else pd.DataFrame() prediction_frame = pd.read_csv(prediction_path) if prediction_path.exists() else pd.DataFrame() complete = set() if not frame.empty: complete = set(zip(frame.split_set, frame.subject_id, frame.fold.astype(int), frame.feature_set)) rows = frame.to_dict("records") if not frame.empty else [] prediction_rows = prediction_frame.to_dict("records") if not prediction_frame.empty else [] assignments = load_fold_assignments(cache / "fold_assignments.csv") new_candidates = 0 for split_set in SCREEN_SPLITS: for subject in subjects: x, y = load_subject(cache, subject) feature_cache = {name: routing_features(x, name) for name in FEATURE_SETS} for fold in range(1, 6): split = assignments[(split_set, subject, fold)] global_state, _ = _load_or_train_global( x, y, split, split_set, config.seed, subject, fold, run_dir, config, False ) for feature_set in FEATURE_SETS: key = (split_set, subject, fold, feature_set) if key in complete: continue print(f"[feature-screen] {split_set} {subject} fold={fold} {feature_set}", flush=True) features = feature_cache[feature_set] router = PhysiologyKMeansRouter(config, config.seed).fit(features[split["train"]]) routed = {role: router.predict(features[indices]) for role, indices in split.items()} states, fallbacks, metadata = _fit_experts( x, y, split, routed, global_state, spec, config, (split_set, config.seed, subject, fold), ) distances = router.normalized_distances(features[split["val"]]) metrics, arrays = _predict_bundle( states, global_state, x, y, split["val"], routed["val"], distances, config, tau=1.0, gamma=1.0, ) rows.append({ "split_set": split_set, "subject_id": subject, "fold": fold, "adaptation": spec.name, "feature_set": feature_set, "feature_count": features.shape[1], "fallback_clusters": len(fallbacks), "silhouette": router.silhouette(), "bootstrap_ari": router.bootstrap_ari(), **metrics[MODEL_HARD], }) for position, trial_index in enumerate(split["val"]): prediction_rows.append({ "split_set": split_set, "subject_id": subject, "fold": fold, "feature_set": feature_set, "trial_index": int(trial_index), "label": int(arrays["labels"][position]), "cluster": int(routed["val"][position]), "distance_0": float(distances[position, 0]), "distance_1": float(distances[position, 1]), "global_probability": float(arrays["global_probability"][position]), "expert_0_probability": float(arrays["expert_probability"][position, 0]), "expert_1_probability": float(arrays["expert_probability"][position, 1]), "fallback_0": fallbacks.get(0, ""), "fallback_1": fallbacks.get(1, ""), }) _save(prediction_path, pd.DataFrame(prediction_rows)) _save(output, pd.DataFrame(rows)) complete.add(key) new_candidates += 1 if max_new_candidates is not None and new_candidates >= max_new_candidates: return output frame = pd.DataFrame(rows) summary = _selection_summary(frame, "feature_set") selected = _pick_with_tolerance(summary, "feature_set", FEATURE_COMPLEXITY) _save(run_dir / "feature_screen_summary.csv", summary) selection.update({"feature_set": selected, "feature_selection_metric": "subject-mean validation balanced accuracy"}) _write_selection(run_dir, selection) return output def run_routing_screen(run_dir: Path, config: OptimizedRoutingConfig | None = None) -> Path: config = config or OptimizedRoutingConfig() selection = _load_selection(run_dir) if "feature_set" not in selection: raise RuntimeError("Run feature screen first") predictions = pd.read_csv(run_dir / "feature_validation_predictions.csv") predictions = predictions[predictions.feature_set == selection["feature_set"]].copy() rows = [] bundle_columns = ["split_set", "subject_id", "fold"] for keys, bundle in predictions.groupby(bundle_columns, sort=False): labels = bundle.label.to_numpy(np.int64) global_probability = bundle.global_probability.to_numpy(np.float64) expert_probability = bundle[["expert_0_probability", "expert_1_probability"]].to_numpy(np.float64) distances = bundle[["distance_0", "distance_1"]].to_numpy(np.float64) for tau in config.tau_candidates: for gamma in config.gamma_candidates: probability, weights, certainty = confidence_soft_probabilities( global_probability, expert_probability, distances, tau, gamma ) rows.append({ "split_set": keys[0], "subject_id": keys[1], "fold": int(keys[2]), "tau": tau, "gamma": gamma, "mean_certainty": float(certainty.mean()), **classification_metrics(labels, probability), }) frame = pd.DataFrame(rows) output = run_dir / "routing_screen.csv" _save(output, frame) subject = frame.groupby(["tau", "gamma", "subject_id"], as_index=False).agg( mean_validation_bacc=("balanced_accuracy", "mean"), mean_validation_accuracy=("accuracy", "mean"), ) summary = subject.groupby(["tau", "gamma"], as_index=False).agg( subjects=("subject_id", "nunique"), mean_validation_bacc=("mean_validation_bacc", "mean"), mean_validation_accuracy=("mean_validation_accuracy", "mean"), ).sort_values(["mean_validation_bacc", "gamma"], ascending=[False, True]) _save(run_dir / "routing_screen_summary.csv", summary) best = summary.iloc[0] selection.update({ "tau": float(best.tau), "gamma": float(best.gamma), "routing_selection_metric": "subject-mean validation balanced accuracy", }) _write_selection(run_dir, selection) return output def _bundle_diversity(labels: np.ndarray, arrays: Mapping[str, np.ndarray], clusters: np.ndarray) -> Dict[str, float]: expert_predictions = arrays["expert_probability"] >= 0.5 global_predictions = arrays["global_probability"] >= 0.5 correct = expert_predictions == labels[:, None] exclusive = correct[:, 0] ^ correct[:, 1] routed_correct = correct[np.arange(len(labels)), clusters] return { "expert_disagreement": float(np.mean(expert_predictions[:, 0] != expert_predictions[:, 1])), "expert_oracle_accuracy": float(np.mean(correct.any(axis=1))), "router_hit_rate_exclusive": float(np.mean(routed_correct[exclusive])) if exclusive.any() else float("nan"), "global_expert0_disagreement": float(np.mean(global_predictions != expert_predictions[:, 0])), "global_expert1_disagreement": float(np.mean(global_predictions != expert_predictions[:, 1])), } def run_final_evaluation( cache: Path, run_dir: Path, config: OptimizedRoutingConfig | None = None, subjects: Sequence[str] = LABELED_SUBJECTS, max_new_folds: int | None = None, ) -> Path: config = config or OptimizedRoutingConfig() selection = _load_selection(run_dir) required = {"adaptation", "feature_set", "tau", "gamma"} if not required.issubset(selection): raise RuntimeError(f"Incomplete selection: missing {sorted(required - set(selection))}") spec = SPEC_BY_NAME[str(selection["adaptation"])] feature_set = str(selection["feature_set"]) tau, gamma = float(selection["tau"]), float(selection["gamma"]) final_dir = run_dir / "final" result_path = final_dir / "fold_results.csv" diagnostic_path = final_dir / "cluster_diagnostics.csv" prediction_path = final_dir / "trial_predictions.csv" expert_path = final_dir / "expert_training.csv" results = pd.read_csv(result_path) if result_path.exists() else pd.DataFrame() diagnostics = pd.read_csv(diagnostic_path) if diagnostic_path.exists() else pd.DataFrame() predictions = pd.read_csv(prediction_path) if prediction_path.exists() else pd.DataFrame() experts = pd.read_csv(expert_path) if expert_path.exists() else pd.DataFrame() complete = set() if not results.empty: counts = results.groupby(["split_set", "train_seed", "subject_id", "fold"])["model"].nunique() complete = {tuple(key) for key, count in counts.items() if count == len(FINAL_MODELS)} result_rows = results.to_dict("records") if not results.empty else [] diagnostic_rows = diagnostics.to_dict("records") if not diagnostics.empty else [] prediction_rows = predictions.to_dict("records") if not predictions.empty else [] expert_rows = experts.to_dict("records") if not experts.empty else [] assignments = load_fold_assignments(cache / "fold_assignments.csv") split_specs = [("all", config.all_trial_seeds)] + [ (f"20_seed_{seed}", (config.seed,)) for seed in SUBSAMPLE_SEEDS ] new_folds = 0 for split_set, train_seeds in split_specs: protocol = "all" if split_set == "all" else "20x20" for train_seed in train_seeds: for subject in subjects: x, y = load_subject(cache, subject) features = routing_features(x, feature_set) for fold in range(1, 6): key = (split_set, train_seed, subject, fold) if key in complete: continue print(f"[optimized-final] {protocol} {split_set} {subject} fold={fold}", flush=True) split = assignments[(split_set, subject, fold)] router_seed = config.seed router = PhysiologyKMeansRouter(config, router_seed).fit(features[split["train"]]) routed = {role: router.predict(features[indices]) for role, indices in split.items()} distances = {role: router.normalized_distances(features[indices]) for role, indices in split.items()} global_state, global_result = _load_or_train_global( x, y, split, split_set, train_seed, subject, fold, final_dir, config, True ) states, fallbacks, metadata = _fit_experts( x, y, split, routed, global_state, spec, config, (split_set, train_seed, subject, fold), ) metrics, arrays = _predict_bundle( states, global_state, x, y, split["test"], routed["test"], distances["test"], config, tau, gamma, ) common = { "protocol": protocol, "split_set": split_set, "train_seed": train_seed, "subject_id": subject, "fold": fold, "adaptation": spec.name, "feature_set": feature_set, "tau": tau, "gamma": gamma, "train_size": len(split["train"]), "val_size": len(split["val"]), "test_size": len(split["test"]), "fallback_clusters": len(fallbacks), } for model in FINAL_MODELS: row = {**common, "model": model, **metrics[model]} if model == MODEL_GLOBAL: row.update({ "fallback_clusters": 0, "best_epoch": global_result["best_epoch"], "epochs_ran": global_result["epochs_ran"], "parameter_count": global_result["parameter_count"], }) result_rows.append(row) diversity = _bundle_diversity(arrays["labels"], arrays, routed["test"]) for cluster in range(config.kmeans_clusters): train_indices = split["train"][routed["train"] == cluster] val_indices = split["val"][routed["val"] == cluster] test_indices = split["test"][routed["test"] == cluster] diagnostic_rows.append({ **common, "cluster": cluster, "feature_count": features.shape[1], "pca_components": router.pca_components, "distance_scale": router.distance_scale, "silhouette": router.silhouette(), "bootstrap_ari": router.bootstrap_ari(), "train_left": int((y[train_indices] == 0).sum()), "train_right": int((y[train_indices] == 1).sum()), "val_left": int((y[val_indices] == 0).sum()), "val_right": int((y[val_indices] == 1).sum()), "test_left": int((y[test_indices] == 0).sum()), "test_right": int((y[test_indices] == 1).sum()), "fallback_reason": fallbacks.get(cluster, ""), **diversity, }) for meta in metadata: expert_rows.append({**common, **meta}) for position, trial_index in enumerate(split["test"]): prediction_rows.append({ **common, "trial_index": int(trial_index), "label": int(arrays["labels"][position]), "cluster": int(routed["test"][position]), "distance_0": float(distances["test"][position, 0]), "distance_1": float(distances["test"][position, 1]), "weight_0": float(arrays["weights"][position, 0]), "weight_1": float(arrays["weights"][position, 1]), "certainty": float(arrays["certainty"][position]), "global_probability": float(arrays["global_probability"][position]), "expert_0_probability": float(arrays["expert_probability"][position, 0]), "expert_1_probability": float(arrays["expert_probability"][position, 1]), "hard_probability": float(arrays["hard_probability"][position]), "soft_probability": float(arrays["soft_probability"][position]), }) _save(diagnostic_path, pd.DataFrame(diagnostic_rows)) _save(prediction_path, pd.DataFrame(prediction_rows)) _save(expert_path, pd.DataFrame(expert_rows)) _save(result_path, pd.DataFrame(result_rows)) complete.add(key) new_folds += 1 if max_new_folds is not None and new_folds >= max_new_folds: summarize_final(final_dir, config.seed) return result_path summarize_final(final_dir, config.seed) (final_dir / "run_config.json").write_text( json.dumps({"config": asdict(config), "selection": selection}, ensure_ascii=False, indent=2), encoding="utf-8", ) return result_path def _bootstrap_ci(values: np.ndarray, seed: int, draws: int = 10000) -> Tuple[float, float]: rng = np.random.default_rng(seed) samples = rng.choice(values, size=(draws, len(values)), replace=True).mean(axis=1) return tuple(np.quantile(samples, (0.025, 0.975)).tolist()) def _holm_adjust(p_values: Sequence[float]) -> List[float]: values = np.asarray(p_values, dtype=float) order = np.argsort(values) adjusted = np.empty_like(values) running = 0.0 count = len(values) for rank, index in enumerate(order): running = max(running, (count - rank) * values[index]) adjusted[index] = min(running, 1.0) return adjusted.tolist() def summarize_final(final_dir: Path, seed: int) -> None: folds = pd.read_csv(final_dir / "fold_results.csv") 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"), std_balanced_accuracy=("balanced_accuracy", "std"), mean_macro_f1=("macro_f1", "mean"), mean_roc_auc=("roc_auc", "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"), ) comparisons = [] for protocol in ("all", "20x20"): pivot = subject[subject.protocol == protocol].pivot( index="subject_id", columns="model", values="mean_accuracy" ) for model in (MODEL_HARD, MODEL_SOFT): if not {MODEL_GLOBAL, model}.issubset(pivot.columns): continue differences = (pivot[model] - pivot[MODEL_GLOBAL]).dropna().to_numpy() if len(differences) == 0: continue low, high = _bootstrap_ci(differences, stable_seed("optimized-bootstrap", seed, protocol, model)) if np.allclose(differences, 0.0): p_value = 1.0 else: try: p_value = float(wilcoxon(differences, method="approx").pvalue) except ValueError: p_value = 1.0 comparisons.append({ "protocol": protocol, "comparison": f"{model} - {MODEL_GLOBAL}", "subjects": len(differences), "mean_difference_pp": 100 * float(differences.mean()), "bootstrap_ci_low_pp": 100 * low, "bootstrap_ci_high_pp": 100 * high, "wilcoxon_p": p_value, "improved_subjects": int((differences > 0).sum()), }) adjusted = _holm_adjust([row["wilcoxon_p"] for row in comparisons]) if comparisons else [] for row, value in zip(comparisons, adjusted): row["holm_p"] = value _save(final_dir / "subject_summary.csv", subject) _save(final_dir / "group_summary.csv", group) _save(final_dir / "paired_comparisons.csv", pd.DataFrame(comparisons)) def smoke_test(cache: Path, device: str = "cpu") -> Dict[str, object]: config = OptimizedRoutingConfig( device=device, global_epochs=2, global_patience=2, expert_epochs=2, expert_patience=2, bootstrap_router_repeats=2, ) x, y = load_subject(cache, "sub1") shapes = {name: list(routing_features(x[:12], name).shape) for name in FEATURE_SETS} assignments = load_fold_assignments(cache / "fold_assignments.csv") split = assignments[(SCREEN_SPLITS[0], "sub1", 1)] features = routing_features(x, "motor_focused") router = PhysiologyKMeansRouter(config, config.seed).fit(features[split["train"]]) weights = router.soft_weights(features[split["val"]], 1.0) model = ShallowConvNet(channels=8, samples=400) global_state = copy.deepcopy(model.state_dict()) train_clusters = router.predict(features[split["train"]]) val_clusters = router.predict(features[split["val"]]) trained = None for cluster in range(2): train_indices = split["train"][train_clusters == cluster] val_indices = split["val"][val_clusters == cluster] if not fallback_reason(y[train_indices], y[val_indices]): trained, meta = fit_adapted_expert( x, y, train_indices, val_indices, global_state, SPEC_BY_NAME["head_only"], config, stable_seed("smoke", cluster), ) break return { "device": device, "feature_shapes": shapes, "router_fit_samples": router.fitted_samples, "pca_components": router.pca_components, "soft_weight_sum_max_error": float(np.max(np.abs(weights.sum(axis=1) - 1))), "silhouette": router.silhouette(), "bootstrap_ari": router.bootstrap_ari(), "expert_trained": trained is not None, "shallow_parameters": count_parameters(model), }