scripts/reference_gpu.py

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