"""Separately named pooled published-window protocol reproduction; never primary imagery-only selection.""" import argparse,json,random,time,hashlib from pathlib import Path import numpy as np,torch from torch import nn from scipy.signal import resample from foundation_models import Foundation from bci_data import cached,montage from classical import metrics from gpu_utils import require_cuda,expired,atomic_save,cpu_state,install_signals def transform(x,idx,kind): # Official preprocessing_speech.py: last 768 samples, Fourier resample to 600. # No extra SETU bandpass or rereference in this reproduction arm. return (resample(x[:,idx,-768:],600,axis=-1)/(1000 if kind=='reve' else 100)).astype(np.float32) class ReferenceModel(nn.Module): def __init__(self,kind,names): super().__init__();self.kind=kind;self.encoder=Foundation(kind,names,n_times=600) if kind=='reve': self.query=nn.Parameter(torch.randn(1,1,512));self.norm_weight=nn.Parameter(torch.ones(512));self.head=nn.Sequential(nn.Dropout(.5),nn.Linear(512,5)) else: layers=[nn.Linear(self.encoder.dimension,600),nn.ELU()] if kind=='cbramod':layers.append(nn.Dropout(.1)) layers.extend([nn.Linear(600,200),nn.ELU()]) if kind=='cbramod':layers.append(nn.Dropout(.1)) layers.append(nn.Linear(200,5));self.head=nn.Sequential(*layers) def forward(self,x): f=self.encoder(x) if self.kind=='reve': f=f.reshape(len(x),-1,512);q=self.query.expand(len(x),-1,-1);f=((q@f.transpose(-1,-2)/512**.5).softmax(-1)@f).squeeze(1);f=f*torch.rsqrt(f.float().pow(2).mean(-1,keepdim=True)+1e-6)*self.norm_weight return self.head(f) def train(kind): from experiment_guard import assert_tuning_open assert_tuning_open();device=require_cuda();install_signals();torch.set_num_threads(4);random.seed(2026);np.random.seed(2026);torch.manual_seed(2026) name=f'reference3s_{kind}_full_pooled_64_seed-2026';out=Path(f'experiments/{name}.json') if out.exists():return xx=[];yy=[];vv=[];zz=[] for subject in range(1,16): x,y,names=cached(subject,'train');v,z,_=cached(subject,'validation');idx=montage(names,64);xx.append(transform(x,idx,kind));yy.append(y);vv.append(transform(v,idx,kind));zz.append(z) x=np.concatenate(xx);y=np.concatenate(yy);v=np.concatenate(vv);z=np.concatenate(zz);model=ReferenceModel(kind,names).to(device);opt=torch.optim.AdamW(model.parameters(),lr=.0001,weight_decay=.01) loader=torch.utils.data.DataLoader(torch.utils.data.TensorDataset(torch.tensor(x),torch.tensor(y)),batch_size=32,shuffle=True);vx=torch.tensor(v,device=device);criterion=nn.CrossEntropyLoss() best=-1;best_epoch=-1;history=[];start=0;last=Path(f'checkpoints/{name}_last.pt');cp=Path(f'checkpoints/{name}.pt') if last.exists(): state=torch.load(last,map_location='cpu',weights_only=False);model.load_state_dict(state['model']);opt.load_state_dict(state['optimizer']);best=state['best'];best_epoch=state['best_epoch'];history=state['history'];start=state['epoch']+1;torch.set_rng_state(state['rng']);torch.cuda.set_rng_state_all(state['cuda_rng']) for epoch in range(start,50): if expired():return model.train();losses=[] for a,b in loader: a=a.to(device);b=b.to(device);opt.zero_grad(set_to_none=True) with torch.autocast('cuda',dtype=torch.bfloat16):loss=criterion(model(a),b) if not torch.isfinite(loss):raise FloatingPointError('Nonfinite reference loss') loss.backward();nn.utils.clip_grad_norm_(model.parameters(),1);opt.step();losses.append(loss.item()) model.eval() with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16):pred=torch.cat([model(q) for q in vx.split(32)]).argmax(1).cpu().numpy() m=metrics(z,pred);history.append({'epoch':epoch,'loss':float(np.mean(losses)),**m}) if m['balanced_accuracy']>best: best=m['balanced_accuracy'];best_epoch=epoch;atomic_save({'model':cpu_state(model),'kind':kind,'indices':idx,'names':names,'epoch':epoch,'validation_metrics':m,'protocol':'published_last768_fourier600_pooled'},cp) atomic_save({'model':cpu_state(model),'optimizer':opt.state_dict(),'epoch':epoch,'best':best,'best_epoch':best_epoch,'history':history,'rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all()},last) print(name,epoch,m['balanced_accuracy'],flush=True) if epoch-best_epoch>=15:break if expired():return out.write_text(json.dumps({'name':name,'model':kind,'channels':64,'subjects':list(range(1,16)),'seed':2026,'split':'validation','protocol':'published_last768_fourier600_pooled','metrics':torch.load(cp,map_location='cpu',weights_only=False)['validation_metrics'],'history':history,'recipe_differences':'Bounded single-seed AdamW 50 epochs, patience15, batch32; no mixup, warmup or scheduler; no full paper hyperparameter/seed sweep. Does not claim exact numerical reproduction.'},indent=2)) def evaluate(): manifest=json.load(open('reports/frozen_selection.json'));out=[];torch.set_num_threads(2) for entry in manifest.get('reference_evaluations',[]): cp=Path(entry['checkpoint']);assert hashlib.sha256(cp.read_bytes()).hexdigest()==entry['sha256'];state=torch.load(cp,map_location='cpu',weights_only=False);model=ReferenceModel(state['kind'],state['names']).eval();model.load_state_dict(state['model']);ys=[];ps=[];by_subject=[] for subject in range(1,16): x,y,_=cached(subject,'test',evaluation=True);a=torch.tensor(transform(x,state['indices'],state['kind'])) with torch.inference_mode():prob=torch.cat([model(q).softmax(-1) for q in a.split(16)]).numpy() pred=prob.argmax(1);ys.extend(y);ps.extend(pred);by_subject.append({'subject':subject,**metrics(y,pred)});np.savez(f"reports/predictions/{entry['record']['name']}_sub-{subject:02d}.npz",y=y,pred=pred,prob=prob) scores=[r['accuracy'] for r in by_subject];ci=np.quantile(np.random.default_rng(2026).choice(scores,(10000,15),replace=True).mean(1),[.025,.975]).tolist();out.append({'model':state['kind'],'protocol':state['protocol'],'channels':64,'test_n':750,'metrics':metrics(ys,ps),'subject_results':by_subject,'subject_bootstrap_accuracy_ci95':ci,'selected_on_validation':entry['record']['name']==manifest.get('reference_primary')}) Path('reports/reference3s_test.json').write_text(json.dumps(out,indent=2)) if __name__=='__main__': p=argparse.ArgumentParser();p.add_argument('--model',choices=['reve','csbrain','cbramod']);p.add_argument('--channels',type=int,choices=[64],default=64);p.add_argument('--evaluate',action='store_true');a=p.parse_args();evaluate() if a.evaluate else train(a.model)