| """Matched ordinary/branch-conditioned MLP baselines. Grouped splits only.""" |
| import os |
| os.environ.setdefault('OMP_NUM_THREADS','1') |
| from pathlib import Path |
| import json,time |
| import numpy as np |
| import torch |
| from torch import nn |
| from safetensors.torch import save_file |
| ROOT=Path(__file__).resolve().parents[1] |
| torch.set_num_threads(2) |
|
|
| def features(c,branch=None): |
| chi=c['chi'];a=c['fx']+2*c['h'] |
| x=[np.log10(chi),a/chi,c['fy']/chi,c['moment']/chi,4*c['h']*c['delta']/chi,c['k1'],c['k2'],c['k3']] |
| if branch is not None:x += [float(branch==b) for b in ['down','straight','up']] |
| return x |
|
|
| def network(d):return nn.Sequential(nn.Linear(d,96),nn.Tanh(),nn.Linear(96,96),nn.Tanh(),nn.Linear(96,48)) |
|
|
| def reconstruct(z,rows): |
| full=np.column_stack([np.zeros(len(z)),z]);th=(full[:,1:]+full[:,:-1])/2 |
| a=np.array([v['controls']['fx']+2*v['controls']['h'] for v in rows])[:,None] |
| fy=np.array([v['controls']['fy'] for v in rows])[:,None] |
| stretch=1+a*np.cos(th)+fy*np.sin(th) |
| xy=np.stack([np.cos(th),np.sin(th)],axis=-1)*stretch[:,:,None]/48 |
| xy=np.concatenate([np.zeros((len(z),1,2)),np.cumsum(xy,axis=1)],axis=1) |
| return xy |
|
|
| def main(): |
| rows=[json.loads(l) for l in (ROOT/'data/equilibria.jsonl').open()] |
| splits={s:np.array([i for i,r in enumerate(rows) if r['split']==s]) for s in ['train','validation','test','parameter_ood']} |
| truth=np.array([r['z'] for r in rows],dtype='float32') |
| out=ROOT/'models';out.mkdir(exist_ok=True);report={} |
| for conditioned in [False,True]: |
| name='branch_conditioned' if conditioned else 'ordinary' |
| torch.manual_seed(14);rng=np.random.default_rng(14) |
| x=np.array([features(r['controls'],r['branch'] if conditioned else None) for r in rows],dtype='float32') |
| mean=x[splits['train']].mean(0);std=x[splits['train']].std(0);std[std<1e-6]=1 |
| X=torch.from_numpy((x-mean)/std);Y=torch.from_numpy(truth) |
| net=network(x.shape[1]);opt=torch.optim.Adam(net.parameters(),lr=.002) |
| best=float('inf');best_epoch=0;start=time.perf_counter() |
| for epoch in range(401): |
| net.train() |
| for idx in np.array_split(rng.permutation(splits['train']),19): |
| opt.zero_grad();loss=(net(X[idx])-Y[idx]).square().mean();loss.backward();opt.step() |
| net.eval() |
| with torch.no_grad():val=float((net(X[splits['validation']])-Y[splits['validation']]).square().mean()) |
| if val<best: |
| best=val;best_epoch=epoch;weights={k:v.clone().contiguous() for k,v in net.state_dict().items()} |
| if epoch%100==0:print(name,epoch,'val angle MSE',val,flush=True) |
| net.load_state_dict(weights);save_file(weights,str(out/f'{name}.safetensors')) |
| (out/f'{name}.json').write_text(json.dumps(dict(mean=mean.tolist(),std=std.tolist(),conditioned=conditioned,architecture=[x.shape[1],96,96,48],activation='tanh',n=48,features='see inference.py / features',best_epoch=best_epoch),indent=2)) |
| with torch.no_grad():pred=net(X).numpy() |
| predicted=reconstruct(pred,rows);target=reconstruct(truth,rows) |
| result={'training_seconds':time.perf_counter()-start,'best_validation_angle_mse':best,'best_epoch':best_epoch,'splits':{}} |
| for s,idx in splits.items(): |
| dist=np.linalg.norm(predicted[idx]-target[idx],axis=2) |
| br=np.where(np.max(np.abs(pred[idx]),axis=1)<.01,'straight',np.where(predicted[idx,-1,1]>0,'up','down')) |
| result['splits'][s]=dict(states=len(idx),centerline_rmse_over_L=float(np.sqrt(np.mean(dist**2))),tip_mean_error_over_L=float(np.mean(dist[:,-1])),tip_p95_error_over_L=float(np.quantile(dist[:,-1],.95)),branch_accuracy=float(np.mean(br==np.array([rows[i]['branch'] for i in idx])))) |
| result['splits'][s]['families']={} |
| for family in ['end_load','intrinsic_curvature','microforce']: |
| select=np.array([rows[i]['family']==family for i in idx]) |
| result['splits'][s]['families'][family]=dict(centerline_rmse_over_L=float(np.sqrt(np.mean(dist[select]**2))),tip_mean_error_over_L=float(np.mean(dist[select,-1]))) |
| report[name]=result |
| (ROOT/'reports/baselines.json').write_text(json.dumps(report,indent=2));print(json.dumps(report,indent=2)) |
|
|
| if __name__=='__main__':main() |
|
|
