1import sys, json, math, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report, all_track_names
 8
 9TRACK = 'two_view_gauge_localization_v2'
10
11def seed_all(s):
12    random.seed(s); np.random.seed(s); torch.manual_seed(s)
13
14def rot(a):
15    c, s = np.cos(a), np.sin(a)
16    return np.array([[c,-s],[s,c]], dtype=np.float32)
17
18def two_view_init_numpy(x):
19    q1,q2,ev1,ev2,et=[x[:,2*i:2*i+2] for i in range(5)]
20    dq=q2-q1; dl=ev2-ev1
21    psi=np.arctan2(dq[:,1],dq[:,0])-np.arctan2(dl[:,1],dl[:,0])
22    out=[]
23    for i,a in enumerate(psi):
24        R=rot(float(a)); xx=q1[i]-R@ev1[i]; out.append(xx+R@et[i])
25    return np.asarray(out,np.float32)[:, :1]
26
27class SharedMLP(nn.Module):
28    def __init__(self, bias=None):
29        super().__init__()
30        self.backbone=nn.Sequential(nn.Linear(10,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh())
31        self.head=nn.Linear(64,1)
32        if bias is not None:
33            with torch.no_grad():
34                self.head.weight.mul_(0.05); self.head.bias.copy_(torch.as_tensor(bias))
35    def forward(self,x): return self.head(self.backbone(x))
36
37def run_one(kind, seed, lr, epochs):
38    seed_all(seed)
39    ds=get_dataset(TRACK, seed=seed, n_train=400, n_test=400)
40    bias=None
41    if kind=='idea': bias=two_view_init_numpy(ds['xtr'].numpy()).mean(0)
42    net, metric, hist=train_model(SharedMLP(bias), ds, epochs=epochs, lr=lr, batch=128, log=lambda *_:None)
43    return float(metric), net, ds
44
45def make_train_fn(kind, cfg):
46    return lambda seed: run_one(kind, int(seed), cfg['lr'], cfg['epochs'])[0]
47
48def main():
49    assert TRACK in all_track_names(), all_track_names()
50    seeds=tuple(range(8)); epochs=18
51    grid=[{'lr':v,'epochs':epochs} for v in (0.001,0.003,0.01)]
52    base=sweep_baseline(lambda cfg: make_train_fn('baseline',cfg), grid, seeds=seeds[:4])
53    idea_cfg_results=[]
54    for cfg in grid:
55        r=evaluate(make_train_fn('idea',cfg), seeds=seeds)
56        idea_cfg_results.append({'cfg':cfg,'mean':r['mean'],'std':r['std'],'per_seed':r['per_seed']})
57    ib=min(idea_cfg_results,key=lambda z:z['mean'])
58    idea={'config':ib['cfg'],'mean':ib['mean'],'std':ib['std'],'per_seed':ib['per_seed'],'n':8}
59    rep=make_report(TRACK,'mlp_tiny',base,idea)
60    # Trained-model signature: compare standard test MSE of trained models against
61    # observed targets, plus the two-view estimate's prediction on the same test data.
62    sig=[]
63    for s in seeds:
64        bm, bn, ds=run_one('baseline',s,base['best_cfg']['lr'],epochs)
65        im, inn, _=run_one('idea',s,ib['cfg']['lr'],epochs)
66        analytic=two_view_init_numpy(ds['xte'].numpy())
67        with torch.no_grad():
68            bd=next(bn.parameters()).device; idv=next(inn.parameters()).device
69            with torch.no_grad():
70                bp=bn(ds['xte'].to(bd)).detach().cpu().numpy(); ip=inn(ds['xte'].to(idv)).detach().cpu().numpy()
71        y=ds['yte'].numpy()
72        sig.append((float(np.mean((analytic-y)**2)),float(np.mean((bp-y)**2)),float(np.mean((ip-y)**2))))
73    vals=np.asarray(sig)
74    signature={'prediction':'two-view analytic gauge initialization reduces target prediction error','analytic_test_mse_mean':float(vals[:,0].mean()),'trained_baseline_test_mse_mean':float(vals[:,1].mean()),'trained_idea_test_mse_mean':float(vals[:,2].mean()),'confirmed':bool(vals[:,2].mean() < vals[:,1].mean())}
75    rep['mechanism_signature']={'values':signature,'confirmed':signature['confirmed']}
76    rep['custom_track']={'name':TRACK,'file':'two_view_registered_track_v2.py','domain':'geometry/localization'}
77    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
78    print(json.dumps(rep,indent=2))
79
80if __name__=='__main__': main()