| 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) |
|
|
