1import sys, json, time
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset, make_model, sweep_baseline, make_report
 7
 8TRACK='directional_polytope'; MODEL='mlp_tiny'
 9SEEDS=tuple(range(8)); BOUNDS=(-1.0,1.0)
10
11def legendre_box(z, eps=0.1, lo=-1., hi=1.):
12    # argmin z*w + eps[-log(w-lo)-log(hi-w)] independently.
13    # Solves z + eps[-1/(w-lo)+1/(hi-w)] = 0 in closed form.
14    a = z
15    L = hi-lo
16    disc = torch.sqrt((a*L)**2 + 4.0*eps*eps)
17    # numerically stable positive root for s=w-lo: a*s^2 + ...
18    # equivalent centered solution: q=w-(lo+hi)/2, q = (eps*a? )
19    # stationarity with center c and radius r: z + eps*(-1/(r+q)+1/(r-q))=0
20    # => z*(r^2-q^2)+2 eps*q=0; q=(sqrt(eps^2+z^2 r^2)-eps)/z
21    r=L/2; c=(lo+hi)/2
22    root=torch.sqrt(eps*eps + (a*r)**2)
23    q=torch.where(a.abs()>1e-8, -a*r*r/(root+eps), torch.zeros_like(a))
24    return (c+q).clamp(min=lo+1e-7,max=hi-1e-7)
25
26def train(seed, lr, kind, penalty=1.0, eps=0.1, epochs=20, return_net=False):
27    torch.manual_seed(seed); np.random.seed(seed)
28    d=get_dataset(TRACK, seed, n_train=400, n_test=200)
29    net=make_model(MODEL,d['input_shape'],d['out_dim'])
30    opt=torch.optim.Adam(net.parameters(),lr=lr)
31    xtr,ytr,xte,yte=d['xtr'],d['ytr'],d['xte'],d['yte']
32    if d.get('out_dim', 1) > 1 and ytr.shape[0] != xtr.shape[0]:
33        ytr=ytr.reshape(xtr.shape[0], d['out_dim']); yte=yte.reshape(xte.shape[0], d['out_dim'])
34    for ep in range(epochs):
35        p=torch.randperm(len(xtr)); net.train()
36        for i in range(0,len(xtr),128):
37            ix=p[i:i+128]; raw=net(xtr[ix])
38            if kind=='idea': out=legendre_box(raw,eps)
39            else: out=raw
40            loss=((out-ytr[ix])**2).mean()
41            if kind=='baseline':
42                loss=loss + penalty*(torch.relu(out-1).square()+torch.relu(-1-out).square()).mean()
43            opt.zero_grad(); loss.backward(); opt.step()
44    net.eval()
45    with torch.no_grad():
46        raw=net(xte); out=legendre_box(raw,eps) if kind=='idea' else raw
47        metric=((out-yte)**2).mean().item()
48    if return_net: return metric,net,d
49    return metric
50
51def base_factory(cfg):
52    return lambda seed: train(seed,cfg['lr'],'baseline',cfg['penalty'])
53def idea_factory(cfg):
54    return lambda seed: train(seed,cfg['lr'],'idea',eps=cfg['eps'])
55
56def signature():
57    metric,net,d=train(0,0.01,'idea',eps=.1,return_net=True)
58    z=net(d['xte'][:1]).detach().requires_grad_(True)
59    w=legendre_box(z,.1)
60    # diagonal observed J via autograd and analytic barrier inverse Hessian
61    obs=[]
62    for j in range(w.shape[1]):
63        obs.append(torch.autograd.grad(w[0,j],z,retain_graph=True)[0][0,j].item())
64    zz=z.detach()[0]; ww=legendre_box(zz,.1); h=.1*((1/(ww+1)**2)+(1/(1-ww)**2))
65    pred=(-1/h).numpy(); obs=np.array(obs)
66    rel=float(np.max(np.abs(obs-pred)/(np.abs(pred)+1e-12)))
67    return {'quantity':'diagonal Jacobian of trained feasibility head', 'predicted_mean_abs':float(np.mean(np.abs(pred))), 'observed_mean_abs':float(np.mean(np.abs(obs))), 'max_relative_error':rel, 'confirmed':bool(rel<1e-5)}
68
69def main():
70    # Union parity: all lr and baseline method knobs used by either side are swept here.
71    grid=[{'lr':lr,'penalty':p} for lr in (.003,.01,.03) for p in (0.0,1.0,10.0)]
72    t=time.time(); base=sweep_baseline(base_factory,grid,seeds=(0,1,2,3))
73    # explicit idea 3-config sweep at the same lr union; report best by 4-seed mean
74    igrid=[{'lr':lr,'eps':e} for lr,e in ((base['best_cfg']['lr'],.05),(base['best_cfg']['lr'],.1),(base['best_cfg']['lr'],.2))]
75    itried=[]
76    for cfg in igrid:
77        vals=[idea_factory(cfg)(s) for s in (0,1,2,3)]
78        itried.append({'cfg':cfg,'mean':float(np.mean(vals))})
79    ibest=min(itried,key=lambda x:x['mean'])['cfg']
80    idea={'best_cfg':ibest,'sweep':itried,'full':__import__('bench').evaluate(idea_factory(ibest),SEEDS)}
81    rep=make_report(TRACK,MODEL,base,idea['full'],signature())
82    rep['idea_sweep']=idea['sweep']; rep['idea_best_cfg']=idea['best_cfg']
83    rep['runtime_seconds']=time.time()-t; rep['protocol_note']='Registered directional_polytope; same MLP, end-to-end trained systems; box log barrier feasibility.'
84    print(json.dumps(rep,indent=2))
85    open('bench_report.json','w').write(json.dumps(rep,indent=2))
86if __name__=='__main__': main()