1import sys, json, itertools, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  7
  8SEEDS = tuple(range(8))
  9GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 5e-3}]
 10EPOCHS = 20
 11BATCH = 128
 12B = 1.0
 13
 14# Finite-candidate KKT shield for the scalar feasible interval [-B,B].
 15# Candidates are the empty active set z and the two one-constraint boundaries.
 16def interval_shield(z, bound=B):
 17    z = torch.as_tensor(z)
 18    return torch.minimum(torch.maximum(z, torch.as_tensor(-bound, dtype=z.dtype, device=z.device)),
 19                         torch.as_tensor(bound, dtype=z.dtype, device=z.device))
 20
 21def kkt_math_check(seed=2906):
 22    # Also check the 2-D disk+box finite candidate construction against SLSQP.
 23    from scipy.optimize import minimize
 24    rng = np.random.default_rng(seed); box=np.array([1.35,1.10]); errs=[]; viol=[]; interior=[]
 25    def g(r,R): return np.array([r[0]-box[0],-r[0]-box[0],r[1]-box[1],-r[1]-box[1],r@r-R*R])
 26    for _ in range(100):
 27        z=rng.uniform(-2,2,2); R=rng.uniform(.65,1.3)
 28        cand=[z.copy()]
 29        for i in range(2):
 30            for s in (-1,1):
 31                q=z.copy(); q[i]=s*box[i]; cand.append(q)
 32        nz=np.linalg.norm(z)
 33        if nz>R: cand.append(z*R/nz)
 34        for s0 in (-1,1):
 35            for s1 in (-1,1): cand.append(np.array([s0*box[0],s1*box[1]]))
 36        for i in range(2):
 37            j=1-i; rem=R*R-box[i]**2
 38            if rem>=0:
 39                for s in (-1,1):
 40                    q=np.zeros(2); q[i]=box[i]; q[j]=s*np.sqrt(rem); cand.append(q)
 41                    q=np.zeros(2); q[i]=-box[i]; q[j]=s*np.sqrt(rem); cand.append(q)
 42        feas=[(0.5*np.sum((q-z)**2),q) for q in cand if np.max(g(q,R))<=1e-9]
 43        if not feas: continue
 44        got=min(feas,key=lambda a:a[0])[1]
 45        ref=minimize(lambda q:.5*np.sum((q-z)**2), np.clip(z,-box,box),
 46          constraints=[{'type':'ineq','fun':lambda q,R=R:R*R-q@q},
 47                       {'type':'ineq','fun':lambda q:box-q},{'type':'ineq','fun':lambda q:box+q}],
 48          method='SLSQP', options={'ftol':1e-12,'maxiter':300})
 49        if ref.success: errs.append(np.linalg.norm(got-ref.x))
 50        viol.append(max(0.,np.max(g(got,R))))
 51        if np.max(g(z,R))<=0: interior.append(np.linalg.norm(got-z))
 52    return {'max_error_vs_slsqp':float(max(errs) if errs else np.nan),
 53            'max_constraint_violation':float(max(viol) if viol else np.nan),
 54            'interior_identity_max_error':float(max(interior) if interior else np.nan),
 55            'n_reference':len(errs)}
 56
 57def seed_all(s):
 58    np.random.seed(s); torch.manual_seed(s)
 59    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 60
 61def train_metric(track, s, lr, shielded):
 62    seed_all(s)
 63    d=get_dataset(track, seed=s, n_train=400, n_test=400)
 64    base=make_model('rnn_small', d['input_shape'], d['out_dim'])
 65    if shielded:
 66        class Shielded(torch.nn.Module):
 67            def __init__(self, net): super().__init__(); self.net=net
 68            def forward(self,x): return interval_shield(self.net(x))
 69        net=Shielded(base)
 70    else: net=base
 71    _, metric, _ = train_model(net,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
 72    return float(metric)
 73
 74def cfg_fn(shielded):
 75    def make(cfg):
 76        return lambda s: train_metric('dynamics', int(s), cfg['lr'], shielded)
 77    return make
 78
 79def observed_signature(s, lr):
 80    seed_all(s); d=get_dataset('dynamics',seed=s,n_train=400,n_test=400)
 81    raw=make_model('rnn_small',d['input_shape'],d['out_dim'])
 82    trained,_,_=train_model(raw,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
 83    raw = trained if trained is not None else raw
 84    raw.eval()
 85    dev=next(raw.parameters()).device
 86    with torch.no_grad(): z=raw(d['xte'].to(dev)); r=interval_shield(z); z=z.detach().cpu().numpy().ravel(); r=r.detach().cpu().numpy().ravel()
 87    interior=np.abs(z)<=B
 88    return {'bound':B,'n_test':len(z),'interior_fraction':float(interior.mean()),
 89            'interior_identity_max_error':float(np.max(np.abs(r[interior]-z[interior])) if interior.any() else 0.),
 90            'max_constraint_violation':float(np.max(np.maximum(np.abs(r)-B,0))),
 91            'mean_shaping_distance':float(np.mean(np.abs(r-z))),
 92            'predicted_interior_error':0.0,'predicted_max_violation':0.0,
 93            'confirmed': bool((np.max(np.abs(r[interior]-z[interior])) if interior.any() else 0.)<1e-7 and np.max(np.maximum(np.abs(r)-B,0))<1e-7)}
 94
 95def main():
 96    t=time.time(); math=kkt_math_check()
 97    # Official sweep (first four seeds), plus parity full evaluations for every lr.
 98    sb=sweep_baseline(cfg_fn(False),GRID,seeds=(0,1,2,3))
 99    baseline_all=[]
100    for c in GRID: baseline_all.append({'cfg':c,'full':evaluate(cfg_fn(False)(c),SEEDS)})
101    best=min(baseline_all,key=lambda q:q['full']['mean'])
102    base_block={'best_cfg':best['cfg'],'sweep':sb['sweep'],'full':best['full'],
103                'parity_full':baseline_all}
104    idea_all=[]
105    for c in GRID: idea_all.append({'cfg':c,'full':evaluate(cfg_fn(True)(c),SEEDS)})
106    ib=min(idea_all,key=lambda q:q['full']['mean'])
107    sig=observed_signature(0,best['cfg']['lr'])
108    sig['math_check']=math
109    sig['trained_baseline_mean']=best['full']['mean']
110    sig['trained_idea_best_mean']=ib['full']['mean']
111    rep=make_report('dynamics','rnn_small',base_block,ib['full'],sig)
112    rep['idea_sweep']=idea_all; rep['runtime_sec']=time.time()-t
113    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
114    print(json.dumps(rep,indent=2))
115if __name__=='__main__': main()