1import json, random, sys
 2import numpy as np
 3import torch
 4sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 5from bench import make_model, train_model, sweep_baseline, evaluate, make_report
 6from bench.data import custom_tracks
 7TRACK='spherical_incompressible_field'
 8RAW = custom_tracks()[TRACK]
 9
10def get_raw(seed,ntr,nte):
11    return RAW.get_dataset(seed=seed,n_train=ntr,n_test=nte)
12
13
14SEEDS=[11,29,47,71,89,107,131,149]
15GRID=[{'lr':1e-3},{'lr':3e-3},{'lr':1e-2}]
16EPOCHS=16
17BATCH=64
18
19def seed_all(s):
20    np.random.seed(s); random.seed(s); torch.manual_seed(s)
21    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
22
23def frame(x):
24    # Deterministic local tangent frame: e1 is the normalized projection of
25    # a fixed axis, e2 completes the oriented tangent frame.
26    q=np.array([1.,0.,0.],dtype=np.float32)
27    dot=x[:,0:1]
28    e1=q[None,:]-dot*x
29    e1=e1/np.maximum(np.linalg.norm(e1,axis=1,keepdims=True),1e-7)
30    e2=np.cross(x,e1)
31    return e1.astype(np.float32),e2.astype(np.float32)
32
33def transformed(seed, idea, ntr=400, nte=200):
34    d=get_raw(seed,ntr,nte)
35    xtr=np.asarray(d['xtr']); xte=np.asarray(d['xte'])
36    if idea:
37        e1,e2=frame(xtr); yt=np.asarray(d['ytr'])
38        ytr=np.concatenate([np.sum(yt*e1,1,keepdims=True),np.sum(yt*e2,1,keepdims=True),np.zeros((len(yt),1),np.float32)],1)
39        e1,e2=frame(xte); yt=np.asarray(d['yte'])
40        yte=np.concatenate([np.sum(yt*e1,1,keepdims=True),np.sum(yt*e2,1,keepdims=True),np.zeros((len(yt),1),np.float32)],1)
41    else:
42        ytr=np.asarray(d['ytr']); yte=np.asarray(d['yte'])
43    return {'xtr':torch.from_numpy(xtr).float(),'ytr':torch.from_numpy(ytr).float(),
44            'xte':torch.from_numpy(xte).float(),'yte':torch.from_numpy(yte).float(),
45            'task':'regression','metric':'mse','input_shape':(3,), 'out_dim':3}
46
47def run_one(seed,idea,lr,collect=False):
48    seed_all(seed); ds=transformed(seed,idea)
49    model=make_model('mlp_tiny',ds['input_shape'],ds['out_dim'])
50    model,metric,_=train_model(model,ds,epochs=EPOCHS,lr=lr,batch=BATCH)
51    if not collect:return float(metric)
52    dev=next(model.parameters()).device; model.eval()
53    with torch.no_grad(): pred=model(ds['xte'].to(dev)).cpu().numpy()
54    pts=np.asarray(get_raw(seed,400,200)['xte'])
55    e1,e2=frame(pts)
56    if idea: pred_phys=pred[:,0:1]*e1+pred[:,1:2]*e2
57    else: pred_phys=pred
58    obs=np.asarray(get_raw(seed,400,200)['yte'])
59    return float(metric),{'predicted':'tangent-frame output should suppress normal component while preserving physical tangent field',
60      'observed_prediction_normal_abs':float(np.mean(np.abs(np.sum(pred_phys*pts,1)))),
61      'observed_target_normal_abs':float(np.mean(np.abs(np.sum(obs*pts,1)))),
62      'observed_physical_rmse':float(np.sqrt(np.mean((pred_phys-obs)**2))),
63      'confirmed':bool(np.mean(np.abs(np.sum(pred_phys*pts,1))) < np.mean(np.abs(np.sum(pred*pts,1))) if idea else False)}
64
65def main():
66    # The registered track is fixed; this check verifies its fields are tangent.
67    d=get_raw(0,400,100); pts=np.asarray(d['xte']); yy=np.asarray(d['yte'])
68    math_check={'target_tangency_max_abs':float(np.max(np.abs(np.sum(pts*yy,1))))}
69    base=sweep_baseline(lambda c:lambda s:run_one(s,False,c['lr']),GRID,seeds=SEEDS)
70    isweep=[]
71    for c in GRID:
72        r=evaluate(lambda s:run_one(s,True,c['lr']),seeds=SEEDS)
73        isweep.append({'cfg':c,'mean':r['mean']})
74    best=min(isweep,key=lambda z:z['mean'])['cfg']
75    idea=evaluate(lambda s:run_one(s,True,best['lr']),seeds=SEEDS)
76    sigb=run_one(SEEDS[0],False,base['best_cfg']['lr'],True)[1]
77    sigi=run_one(SEEDS[0],True,best['lr'],True)[1]
78    sig={'predicted_effect':'tangent-frame pullback should reduce normal leakage and preserve tangent coefficients',
79         'baseline':sigb,'idea':sigi,'confirmed':sigi['observed_prediction_normal_abs'] < sigb['observed_prediction_normal_abs']}
80    base['parity_grid']=GRID; base['idea_grid']=isweep
81    rep=make_report(TRACK,'mlp_tiny',base,idea,{'mechanism_signature':sig,
82      'protocol_notes':{'seeds':SEEDS,'epochs':EPOCHS,'batch':BATCH,'same_architecture':True,
83                        'same_lr_union':True},
84      'registered_track':True})
85    out={'bench_report':rep,'math_check':math_check}
86    open('registered_bench_report.json','w').write(json.dumps(out,indent=2))
87    print(json.dumps(out,indent=2))
88if __name__=='__main__':main()