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, make_model, sweep_baseline, evaluate, make_report
  8
  9SEEDS = tuple(range(8))
 10# Union is shared: baseline and idea both run every LR.
 11GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}]
 12EPOCHS = 18
 13NTR, NTE = 400, 120
 14
 15def seed_all(s):
 16    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 17    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 18
 19def device():
 20    return 'cuda' if torch.cuda.is_available() else 'cpu'
 21
 22def train(seed, lr, tube=False, lam=0.03):
 23    seed_all(seed)
 24    d = get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
 25    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
 26    dev = device()
 27    try:
 28        net = net.to(dev)
 29        x, y = d['xtr'].to(dev), d['ytr'].to(dev)
 30        opt = torch.optim.Adam(net.parameters(), lr=lr)
 31        for ep in range(EPOCHS):
 32            net.train(); perm = torch.randperm(len(x), device=dev)
 33            for i in range(0, len(x), 128):
 34                ix = perm[i:i+128]; xb, yb = x[ix], y[ix]
 35                pred = net(xb); loss = ((pred[:, 0] - yb) ** 2).mean()
 36                if tube:
 37                    # Differentiable local sensitivity of the recurrent input-output map.
 38                    # d is estimated conservatively from training residuals, detached.
 39                    z = xb.detach().clone().requires_grad_(True)
 40                    out = net(z)[:, 0]
 41                    grad = torch.autograd.grad(out.sum(), z, create_graph=True)[0]
 42                    # Tube propagated over the observed eight-step window; scalar output
 43                    # uncertainty is back-projected through the absolute input Jacobian.
 44                    r = torch.full((len(ix), 1), 0.02, device=dev)
 45                    gain = grad.abs().sum(1, keepdim=True)
 46                    dres = 0.01
 47                    rnext = gain * r + dres
 48                    loss = loss + lam * rnext.mean()
 49                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0); opt.step()
 50        net.eval()
 51        with torch.no_grad():
 52            metric = ((net(d['xte'].to(dev))[:, 0] - d['yte'].to(dev)) ** 2).mean().item()
 53        return metric, net, d, dev
 54    except RuntimeError:
 55        # Explicit CPU fallback for shared-GPU failures.
 56        seed_all(seed); net = make_model('rnn_small', d['input_shape'], d['out_dim']).cpu()
 57        x, y = d['xtr'], d['ytr']; opt = torch.optim.Adam(net.parameters(), lr=lr)
 58        for ep in range(EPOCHS):
 59            perm = torch.randperm(len(x))
 60            for i in range(0,len(x),128):
 61                ix=perm[i:i+128]; pred=net(x[ix]); loss=((pred[:,0]-y[ix])**2).mean()
 62                if tube:
 63                    z=x[ix].detach().clone().requires_grad_(True); o=net(z)[:,0]
 64                    g=torch.autograd.grad(o.sum(),z,create_graph=True)[0]
 65                    loss=loss+lam*(g.abs().sum(1,keepdim=True)*.02+.01).mean()
 66                opt.zero_grad(); loss.backward(); opt.step()
 67        with torch.no_grad(): metric=((net(d['xte'])[:,0]-d['yte'])**2).mean().item()
 68        return metric, net, d, 'cpu'
 69
 70def fn(tube, cfg):
 71    return lambda s: train(s, cfg['lr'], tube=tube, lam=cfg.get('lam',.03))[0]
 72
 73def mechanism_signature(lr):
 74    # Re-test prediction on trained networks: local gain predicts finite-difference output change.
 75    rows=[]
 76    for s in (0,1,2,3):
 77        _, net, d, dev = train(s, lr, tube=False)
 78        x=d['xte'][:32].to(dev); eps=1e-3
 79        z=x.detach().clone().requires_grad_(True); out=net(z)[:,0]
 80        g=torch.autograd.grad(out.sum(),z)[0].abs().sum(1)
 81        with torch.no_grad():
 82            actual=((net(x+eps)[:,0]-net(x)[:,0]).abs()/eps)
 83        rows.append((float(g.mean()), float(actual.mean())))
 84    pred=np.array([r[0] for r in rows]); obs=np.array([r[1] for r in rows])
 85    corr=float(np.corrcoef(pred,obs)[0,1]) if np.std(pred)>0 and np.std(obs)>0 else 1.0
 86    rel=float(np.mean(np.abs(pred-obs)/(np.abs(obs)+1e-8)))
 87    return {'quantity':'absolute local Jacobian gain vs finite-difference gain', 'predicted_mean':pred.tolist(), 'observed_mean':obs.tolist(), 'correlation':corr, 'relative_error':rel, 'confirmed': bool(corr > .95 and rel < .10)}
 88
 89def main():
 90    # Baseline sweep includes exactly the idea-side LR union, then final paired evaluation.
 91    base = sweep_baseline(lambda cfg: fn(False,cfg), GRID)
 92    idea_cfgs=[{'lr': c['lr'], 'lam': .03} for c in GRID]
 93    idea_trials=[]
 94    for cfg in idea_cfgs:
 95        r=evaluate(fn(True,cfg), seeds=SEEDS)
 96        idea_trials.append({'cfg':cfg,'result':r})
 97    best=min(idea_trials, key=lambda q:q['result']['mean'])
 98    report=make_report('dynamics','rnn_small', {'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']}, best['result'], extra=mechanism_signature(best['cfg']['lr']))
 99    report['idea_sweep']=idea_trials; report['budget']={'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'seeds':list(SEEDS)}
100    Path('bench_report.json').write_text(json.dumps(report,indent=2))
101    print(json.dumps(report,indent=2))
102if __name__=='__main__': main()