1import os, sys, json, math, random
  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, train_model, evaluate, sweep_baseline, make_report
  7
  8SEEDS=tuple(range(8))
  9# Union is shared: baseline and idea both evaluated at all lr/epochs settings.
 10GRID=[{'lr':1e-3,'epochs':10},{'lr':3e-3,'epochs':10},{'lr':6e-3,'epochs':10}]
 11BATCH=128
 12
 13def sanity():
 14    # A monotone map, independently of learned models, checks the claimed arithmetic.
 15    def F(z): return np.tanh(np.array([[.7,.2],[.1,.8]]) @ z)
 16    rng=np.random.default_rng(1907); rates=[]; widths=[]
 17    for h in [.2,.1,.05,.025]:
 18        ok=[]; ww=[]
 19        for _ in range(3000):
 20            lo=rng.uniform(-.5,.3,2); hi=lo+rng.uniform(.01,.15,2)
 21            yl=F(lo); yh=F(hi)
 22            lower=h*np.floor((yl-.01)/h)
 23            upper=h*np.floor((yh+.01+h)/h)
 24            z=lo+rng.random(2)*(hi-lo); q=h*np.floor(F(z)/h)
 25            ok.append(np.all((q>=lower-1e-9)&(q<=upper+1e-9)))
 26            ww.append(np.mean(upper-lower))
 27        rates.append(float(np.mean(ok))); widths.append(float(np.mean(ww)))
 28    return {'containment':rates,'h': [.2,.1,.05,.025], 'mean_width':widths,
 29            'all_contained':bool(min(rates)>=.999),
 30            'width_decreases':bool(all(widths[i+1] <= widths[i]+1e-9 for i in range(3)))}
 31
 32def seed_all(s):
 33    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 34    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 35
 36def train_base(seed,cfg,return_model=False):
 37    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=120)
 38    net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 39    net,metric,hist=train_model(net,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=BATCH,weight_decay=0.0,log=lambda x:None)
 40    if return_model: return float(metric),net,ds
 41    return float(metric)
 42
 43def interval_loss(net,x,y,h=.05,delta=.015):
 44    # Local state-action cell: opposite corners are formed by coordinatewise +/- h/2.
 45    # This is a training surrogate for T- / T+ containment; target residual is inflated.
 46    lo=x-h/2; hi=x+h/2
 47    yl=net(lo); yh=net(hi)
 48    lower=torch.minimum(yl,yh)-delta
 49    upper=torch.maximum(yl,yh)+delta+h
 50    contain=torch.relu(lower-y)+torch.relu(y-upper)
 51    width=torch.relu(upper-lower)
 52    # retain task accuracy while discouraging unnecessarily broad certificates
 53    return (contain**2).mean() + .02*width.mean()
 54
 55def train_idea(seed,cfg,return_model=False):
 56    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=120)
 57    net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 58    # Own loop is required because the idea changes the loss; all other settings match train_model.
 59    device='cuda' if torch.cuda.is_available() else 'cpu'
 60    try:
 61        net=net.to(device); xtr,ytr=ds['xtr'].to(device),ds['ytr'].to(device)
 62        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=0.0)
 63        for ep in range(cfg['epochs']):
 64            net.train(); perm=torch.randperm(len(xtr),device=device)
 65            for i in range(0,len(xtr),BATCH):
 66                ix=perm[i:i+BATCH]; x=xtr[ix]; y=ytr[ix]
 67                pred=net(x); mse=((pred-y)**2).mean()
 68                loss=mse + .15*interval_loss(net,x,y)
 69                opt.zero_grad(); loss.backward(); opt.step()
 70        net.eval()
 71        with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
 72        if return_model:return metric,net,ds
 73        return metric
 74    except RuntimeError:
 75        # CPU fallback from a fresh identical seed/model.
 76        seed_all(seed); net=make_model('rnn_small',ds['input_shape'],ds['out_dim']).cpu()
 77        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
 78        for ep in range(cfg['epochs']):
 79            perm=torch.randperm(len(ds['xtr']))
 80            for i in range(0,len(perm),BATCH):
 81                ix=perm[i:i+BATCH]; pred=net(ds['xtr'][ix]); y=ds['ytr'][ix]
 82                loss=((pred-y)**2).mean()+.15*interval_loss(net,ds['xtr'][ix],y)
 83                opt.zero_grad();loss.backward();opt.step()
 84        with torch.no_grad(): metric=float(((net(ds['xte'])-ds['yte'])**2).mean())
 85        if return_model:return metric,net,ds
 86        return metric
 87
 88def signature(cfg):
 89    vals=[]; widths=[]; violations=[]
 90    for s in SEEDS:
 91        metric,net,ds=train_idea(s,cfg,True)
 92        x,y=ds['xte'],ds['yte']; dev=next(net.parameters()).device
 93        with torch.no_grad():
 94            z=x.to(dev); yl=net(z-.025); yh=net(z+.025)
 95            lower=torch.minimum(yl,yh)-.015; upper=torch.maximum(yl,yh)+.015+.05
 96            yy=y.to(dev); vals.append(float(((yy>=lower)&(yy<=upper)).all(1).float().mean()))
 97            widths.append(float((upper-lower).mean()))
 98            violations.append(float(torch.relu(lower-yy).mean()+torch.relu(yy-upper).mean()))
 99    return {'predicted':'smaller h should reduce interval width while residual inflation preserves containment',
100            'observed_containment_mean':float(np.mean(vals)), 'observed_width_mean':float(np.mean(widths)),
101            'observed_violation_mean':float(np.mean(violations)),
102            'confirmed':bool(np.mean(vals)>=.95 and np.mean(violations)<1e-5)}
103
104def main():
105    san=sanity()
106    # Baseline sweep uses all three configs on four seeds; each union lr is also idea-tested.
107    base=sweep_baseline(lambda c: (lambda s: train_base(s,c)),GRID,seeds=(0,1,2,3))
108    idea_cfgs=[base['best_cfg']]+[c for c in GRID if c!=base['best_cfg']]
109    idea_runs=[]
110    for c in idea_cfgs:
111        r=evaluate(lambda s,c=c: train_idea(s,c),SEEDS)
112        idea_runs.append({'cfg':c,'result':r})
113    best=min(idea_runs,key=lambda q:q['result']['mean'])
114    report=make_report('dynamics','rnn_small',base,best['result'],{'mechanism_signature':signature(best['cfg']),
115        'sanity_check':san,'idea_config_sweep':idea_runs,
116        'protocol_note':'Baseline and idea share rnn_small, dataset, lr/epoch union, batch and seeds; idea differs only by interval training loss.'})
117    report['custom_track']=None
118    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
119    print(json.dumps(report,indent=2))
120if __name__=='__main__': main()