File size: 3,000 Bytes
4fcc185 | 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 | import os
os.environ.setdefault('OPENBLAS_NUM_THREADS','1');os.environ.setdefault('OMP_NUM_THREADS','1')
import argparse,json,time,joblib
from pathlib import Path
import numpy as np
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score,balanced_accuracy_score,f1_score,confusion_matrix
from pyriemann.estimation import Covariances
from pyriemann.tangentspace import TangentSpace
from pyriemann.classification import MDM
from scipy.signal import welch
from bci_data import cached,preprocess,montage
from threadpoolctl import threadpool_limits
threadpool_limits(1)
def metrics(y,p):
return dict(accuracy=accuracy_score(y,p),balanced_accuracy=balanced_accuracy_score(y,p),macro_f1=f1_score(y,p,average='macro',zero_division=0),confusion_matrix=confusion_matrix(y,p,labels=range(5)).tolist())
def features(x,kind):
if kind in ['ts','mdm']:return Covariances(estimator='oas').fit_transform(x)
f,p=welch(x,fs=200,nperseg=200,axis=-1)
return np.concatenate([np.log(p[:,:,(f>=a)&(f<b)].mean(-1)+1e-12) for a,b in [(1,4),(4,8),(8,13),(13,20),(20,30),(30,40)]],axis=1)
def estimator(kind,c):
if kind=='mdm':return MDM(metric='logeuclid')
parts=[TangentSpace(metric='riemann')] if kind=='ts' else []
return make_pipeline(*parts,StandardScaler(),LogisticRegression(C=c,max_iter=1500,class_weight='balanced'))
def run(subject,n,kind,selection):
name=f'strict_{kind}_{selection}_{n}_sub-{subject:02d}';out=Path('experiments')/(name+'.json')
if out.exists():return
t=time.time();x,y,names=cached(subject,'train');v,vy,_=cached(subject,'validation');idx=montage(names,n,selection)
a=features(preprocess(x,idx),kind);b=features(preprocess(v,idx),kind)
best=None;scores=[]
for c in ([1] if kind=='mdm' else [.001,.01,.1,1,10]):
model=estimator(kind,c);model.fit(a,y);pred=model.predict(b);m=metrics(vy,pred);scores.append({'C':c,**m})
if best is None or m['balanced_accuracy']>best[0]:best=(m['balanced_accuracy'],model,c,m)
joblib.dump({'model':best[1],'indices':idx,'kind':kind,'names':names,'subject':subject,'preprocessing':'car_1-40_imagined0to2s_200hz'},Path('checkpoints')/(name+'.joblib'))
out.write_text(json.dumps(dict(name=name,subject=subject,channels=n,montage=selection,model=kind,split='validation',selected_C=best[2],metrics=best[3],search=scores,seconds=time.time()-t),indent=2))
print(name,round(best[0],4),round(time.time()-t,1),flush=True)
if __name__=='__main__':
from experiment_guard import assert_tuning_open
assert_tuning_open()
p=argparse.ArgumentParser();p.add_argument('--channels',nargs='+',type=int,default=[64]);p.add_argument('--subjects',nargs='+',type=int,default=list(range(1,16)));p.add_argument('--models',nargs='+',default=['bandpower','mdm','ts']);p.add_argument('--montage',default='distributed');args=p.parse_args()
for n in args.channels:
for kind in args.models:
for s in args.subjects:run(s,n,kind,args.montage)
|