| """Load safetensors baseline; optionally correct using the released mechanics.""" |
| from pathlib import Path |
| import json |
| import numpy as np |
| import torch |
| from safetensors.torch import load_file |
| from train import features,network |
| from rod import Rod,Controls |
|
|
| def predict(controls,branch='down',model_dir=None,correct=False): |
| folder=Path(model_dir) if model_dir else Path(__file__).resolve().parents[1]/'models' |
| name='branch_conditioned' if branch is not None else 'ordinary' |
| config=json.loads((folder/f'{name}.json').read_text()) |
| net=network(config['architecture'][0]);net.load_state_dict(load_file(str(folder/f'{name}.safetensors')));net.eval() |
| x=(np.array(features(controls,branch))-config['mean'])/config['std'] |
| with torch.no_grad():z=net(torch.tensor(x,dtype=torch.float32)).numpy() |
| rod=Rod(48);c=Controls(**controls) |
| if correct:return rod.solve(c,z) |
| th=rod.A@z;stretch=1+(c.fx+2*c.h)*np.cos(th)+c.fy*np.sin(th) |
| xy=np.vstack([np.zeros(2),np.cumsum(stretch[:,None]*np.column_stack([np.cos(th),np.sin(th)])/48,axis=0)]) |
| return dict(z=z,centerline=xy,tip=xy[-1],note='Raw neural estimate; stability is not certified.') |
|
|
| if __name__=='__main__': |
| c=dict(chi=.0001,fx=-.016428,fy=0.,moment=0.,h=.01,delta=.1,k1=0.,k2=0.,k3=0.) |
| for corrected in [False,True]: |
| r=predict(c,'down',correct=corrected) |
| print('corrected' if corrected else 'neural','tip',r['tip'],'residual',r.get('residual'),'stability',r.get('stability')) |
|
|
