1import 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, sweep_baseline, make_report
  7from bench.protocol import DEFAULT_SEEDS
  8
  9EPOCHS, BATCH = 12, 128
 10GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}]
 11
 12
 13def seed_all(seed):
 14    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 15    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 16
 17
 18def controlled_forward(net, x, stabilized=False):
 19    seq = x.reshape(x.shape[0], 8, 3)
 20    gru, head = net.rnn, net.head
 21    h = torch.zeros(1, x.shape[0], gru.hidden_size, device=x.device, dtype=x.dtype)
 22    for k in range(8):
 23        raw, _ = gru(seq[:, k:k+1, :], h)
 24        proposed = raw[:, 0, :]
 25        h = (h[:, 0, :] + (0.5 if stabilized else 1.0) *
 26             (proposed - h[:, 0, :])).unsqueeze(0)
 27    return head(h[0])
 28
 29
 30def train_idea(ds, lr, seed):
 31    seed_all(seed)
 32    net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 33    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 34    try:
 35        net.to(device)
 36        opt = torch.optim.Adam(net.parameters(), lr=lr)
 37        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 38        for _ in range(EPOCHS):
 39            net.train(); perm = torch.randperm(len(x), device=device)
 40            for i in range(0, len(x), BATCH):
 41                ix = perm[i:i+BATCH]
 42                loss = nn.functional.mse_loss(controlled_forward(net, x[ix], True), y[ix])
 43                opt.zero_grad(); loss.backward(); opt.step()
 44        net.eval()
 45        with torch.no_grad():
 46            m = float(((controlled_forward(net, ds['xte'].to(device), True) - ds['yte'].to(device)) ** 2).mean())
 47        return net, m
 48    except RuntimeError:
 49        if device.type != 'cuda': raise
 50        torch.cuda.empty_cache()
 51        net = make_model('rnn_small', ds['input_shape'], ds['out_dim']).cpu()
 52        opt = torch.optim.Adam(net.parameters(), lr=lr)
 53        x, y = ds['xtr'], ds['ytr']
 54        for _ in range(EPOCHS):
 55            perm = torch.randperm(len(x))
 56            for i in range(0, len(x), BATCH):
 57                ix = perm[i:i+BATCH]
 58                loss = nn.functional.mse_loss(controlled_forward(net, x[ix], True), y[ix])
 59                opt.zero_grad(); loss.backward(); opt.step()
 60        with torch.no_grad():
 61            m = float(((controlled_forward(net, ds['xte'], True)-ds['yte'])**2).mean())
 62        return net, m
 63
 64
 65def baseline_factory(cfg):
 66    def run(seed):
 67        ds = get_dataset('dynamics', int(seed), 400, 200)
 68        seed_all(int(seed)); net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 69        _, metric, _ = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_: None)
 70        return metric
 71    return run
 72
 73
 74def idea_factory(cfg):
 75    return lambda seed: train_idea(get_dataset('dynamics', int(seed), 400, 200), cfg['lr'], int(seed))[1]
 76
 77
 78def train_one_baseline(ds, lr, seed):
 79    seed_all(seed); net=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 80    net,m,_=train_model(net,ds,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_:None)
 81    return net,m
 82
 83
 84def trained_exponent(net, x, stabilized):
 85    """Exact JVP estimate on a trained GRU; called on CPU to avoid cuDNN JVP limits."""
 86    net = net.cpu().eval(); x = x.cpu(); gru = net.rnn
 87    seq = x[:1].reshape(1, 8, 3)
 88    q = torch.randn(gru.hidden_size); q /= q.norm(); total = 0.0
 89    h = torch.zeros(gru.hidden_size)
 90    for k in range(8):
 91        u = seq[:, k:k+1, :]
 92        def fmap(v):
 93            raw, _ = gru(u, v.reshape(1,1,-1)); r = raw[0,0]
 94            return v + 0.5*(r-v) if stabilized else r
 95        _, v = torch.autograd.functional.jvp(fmap, h, q)
 96        n = float(v.norm().detach()); total += math.log(max(n,1e-12)); q=v.detach()/max(n,1e-12)
 97        with torch.no_grad():
 98            raw,_=gru(u,h.reshape(1,1,-1)); r=raw[0,0]
 99            h=(h+0.5*(r-h) if stabilized else r).detach()
100    return total/8.0
101
102
103def signature(seed=0):
104    ds=get_dataset('dynamics',seed,400,200)
105    bnet,bm=train_one_baseline(ds,3e-3,seed+1000)
106    inet,im=train_idea(ds,3e-3,seed+1000)
107    eb=trained_exponent(bnet,ds['xte'],False); ei=trained_exponent(inet,ds['xte'],True)
108    return {'trained_baseline_mse':bm,'trained_idea_mse':im,
109            'baseline_discrete_exponent':eb,'controlled_discrete_exponent':ei,
110            'predicted_sign_preserved':bool(eb<0),'observed_sign_preserved':bool(ei<0),
111            'confirmed':bool((eb<0)==(ei<0))}
112
113
114def main():
115    base=sweep_baseline(baseline_factory, GRID, seeds=(0,1,2,3))
116    vals=[idea_factory(base['best_cfg'])(s) for s in DEFAULT_SEEDS]
117    idea={'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)}
118    sig=signature(0)
119    idea_sweep=[]
120    for cfg in GRID:
121        vals4=[idea_factory(cfg)(s) for s in (0,1,2,3)]
122        idea_sweep.append({'cfg':cfg,'mean':float(np.mean(vals4))})
123    rep=make_report('dynamics','rnn_small',base,idea,extra={'mechanism_signature':sig,'idea_sweep':idea_sweep,'track_justification':'Controlled pendulum rollout is the built-in stability/control/Lyapunov match.','protocol':'8 paired seeds; baseline and idea share lr union; standard test MSE.'})
124    out={'bench_report':rep}; print(json.dumps(out,indent=2)); open('bench_report.json','w').write(json.dumps(out,indent=2))
125
126if __name__=='__main__': main()