File size: 4,531 Bytes
d37c079 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | #!/usr/bin/env python
"""Cross-subject MI with Euclidean Alignment (EA) — the biggest calibration-free win.
Euclidean Alignment (He & Wu, 2020): for each subject, whiten every trial by that
subject's own reference matrix R^{-1/2} (R = mean spatial covariance over the
subject's trials). R needs no labels, so it can be computed on the *test* subject's
trials too — i.e. this stays calibration-free. It pulls each subject's covariance
cloud to a common center, cutting inter-subject shift before any classifier.
We run a manual leave-one-subject-out loop (so alignment is per-subject) and compare
`none` vs `ea` for the strong Riemannian pipelines.
python src/baselines/riemann_alignment.py --dataset 2a --subjects 9
"""
from __future__ import annotations
import argparse
import os
import warnings
from importlib import import_module
from pathlib import Path
import numpy as np
_DATA = Path(__file__).resolve().parents[2] / "mne_data"
_DATA.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("MNE_DATA", str(_DATA))
os.environ.setdefault("MOABB_RESULTS", str(_DATA))
from sklearn.pipeline import make_pipeline
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from pyriemann.estimation import Covariances
from pyriemann.tangentspace import TangentSpace
from pyriemann.classification import MDM
from moabb.paradigms import LeftRightImagery
warnings.filterwarnings("ignore")
DATASETS = {"2a": "BNCI2014_001", "2b": "BNCI2014_004"}
def load_mi(key: str, n_subjects: int):
ds = getattr(import_module("moabb.datasets"), DATASETS[key])()
if n_subjects:
ds.subject_list = ds.subject_list[:n_subjects]
X, y, meta = LeftRightImagery(fmin=8, fmax=32).get_data(dataset=ds)
return np.asarray(X), np.asarray(y), meta["subject"].to_numpy()
def inv_sqrt(R: np.ndarray) -> np.ndarray:
"""Real symmetric matrix inverse square root via eigendecomposition."""
w, V = np.linalg.eigh(R)
w = np.clip(w, 1e-12, None)
return (V * (1.0 / np.sqrt(w))) @ V.T
def euclidean_align(X: np.ndarray) -> np.ndarray:
"""EA on one subject's trials. X: (n, ch, t) -> aligned (n, ch, t)."""
covs = np.matmul(X, X.transpose(0, 2, 1)) # (n, ch, ch)
W = inv_sqrt(covs.mean(0)) # R^{-1/2}, label-free
return np.matmul(W[None], X)
def make_clf(kind: str):
if kind == "TS+LR":
return make_pipeline(Covariances("oas"), TangentSpace(),
LogisticRegression(max_iter=1000))
if kind == "MDM":
return make_pipeline(Covariances("oas"), MDM())
raise ValueError(kind)
def loso(X, y, subj, clf_kind: str, align: str):
aucs = {}
for te in np.unique(subj):
tr_mask = subj != te
# align each subject independently (train subjects + the held-out subject)
Xtr_parts, ytr_parts = [], []
for s in np.unique(subj[tr_mask]):
m = subj == s
Xs = euclidean_align(X[m]) if align == "ea" else X[m]
Xtr_parts.append(Xs)
ytr_parts.append(y[m])
Xtr = np.concatenate(Xtr_parts)
ytr = np.concatenate(ytr_parts)
Xte = euclidean_align(X[subj == te]) if align == "ea" else X[subj == te]
yte = y[subj == te]
clf = make_clf(clf_kind)
clf.fit(Xtr, ytr)
proba = clf.predict_proba(Xte)[:, 1]
y_bin = (yte == clf.classes_[1]).astype(int)
aucs[te] = roc_auc_score(y_bin, proba)
return aucs
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--dataset", default="2a", choices=list(DATASETS))
ap.add_argument("--subjects", type=int, default=9)
args = ap.parse_args()
X, y, subj = load_mi(args.dataset, args.subjects)
print(f"loaded X={X.shape}, {len(np.unique(subj))} subjects, classes={sorted(set(y))}\n")
print(f"{'pipeline':<10}{'align':<8}{'mean AUC':>10}{'std':>8}")
print("-" * 36)
results = {}
for clf_kind in ["TS+LR", "MDM"]:
for align in ["none", "ea"]:
aucs = loso(X, y, subj, clf_kind, align)
vals = np.array(list(aucs.values()))
results[(clf_kind, align)] = vals
print(f"{clf_kind:<10}{align:<8}{vals.mean():>10.3f}{vals.std():>8.3f}")
print("\nEA improvement (mean AUC):")
for clf_kind in ["TS+LR", "MDM"]:
d = results[(clf_kind, "ea")].mean() - results[(clf_kind, "none")].mean()
print(f" {clf_kind}: {d:+.3f}")
if __name__ == "__main__":
main()
|