1import sys, json, 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, evaluate, sweep_baseline, make_report
 8
 9SEEDS = tuple(range(8)); EPOCHS = 30; BATCH = 128
10GRID = [{'lr': 1e-3, 'rho_cap': None}, {'lr': 3e-3, 'rho_cap': None}, {'lr': 5e-3, 'rho_cap': None}]
11IDEA_GRID = [{'lr': 1e-3, 'rho_cap': .90}, {'lr': 3e-3, 'rho_cap': .90}, {'lr': 5e-3, 'rho_cap': .90}]
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
17def operator_gain(W):
18    return float(np.linalg.svd(W, compute_uv=False)[0])
19
20def project_gru(net, cap):
21    with torch.no_grad():
22        w = net.rnn.weight_hh_l0
23        blocks = w.chunk(3, 0)
24        gains = [float(torch.linalg.matrix_norm(b, 2).cpu()) for b in blocks]
25        g = max(gains)
26        if g > cap: w.mul_(cap / (g + 1e-12))
27        return g
28
29def train(seed, lr, cap=None, return_net=False):
30    seed_all(seed); d = get_dataset('dynamics', seed, n_train=400, n_test=200)
31    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
32    device = 'cuda' if torch.cuda.is_available() else 'cpu'
33    try:
34        net.to(device); x, y = d['xtr'].to(device), d['ytr'].to(device)
35        opt = torch.optim.Adam(net.parameters(), lr=lr); lossf = nn.MSELoss(); hist=[]
36        for _ in range(EPOCHS):
37            net.train(); perm=torch.randperm(len(x), device=device); total=0.
38            for i in range(0,len(x),BATCH):
39                ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]), y[ix])
40                opt.zero_grad(); loss.backward(); opt.step()
41                if cap is not None: project_gru(net, cap)
42                total += float(loss.detach())*len(ix)
43            hist.append(total/len(x))
44        net.eval()
45        with torch.no_grad(): metric=float(lossf(net(d['xte'].to(device)), d['yte'].to(device)).cpu())
46        return (metric, net, d) if return_net else metric
47    except RuntimeError:
48        seed_all(seed); net=make_model('rnn_small', d['input_shape'], d['out_dim']); net.to('cpu')
49        x,y=d['xtr'],d['ytr']; opt=torch.optim.Adam(net.parameters(),lr=lr); lossf=nn.MSELoss()
50        for _ in range(EPOCHS):
51            perm=torch.randperm(len(x))
52            for i in range(0,len(x),BATCH):
53                ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward(); opt.step()
54                if cap is not None: project_gru(net,cap)
55        with torch.no_grad(): metric=float(lossf(net(d['xte']),d['yte']))
56        return (metric,net,d) if return_net else metric
57
58def baseline_factory(cfg): return lambda seed: train(seed, cfg['lr'], None)
59def idea_factory(cfg): return lambda seed: train(seed, cfg['lr'], cfg['rho_cap'])
60
61def signature():
62    rows=[]
63    for seed in (0,1,2,3):
64        _, net, d = train(seed, 3e-3, .90, True); dev=next(net.parameters()).device
65        x=d['xte'][:32].to(dev); delta=torch.randn_like(x)*1e-3
66        with torch.no_grad(): z1=net(x); z2=net(x+delta)
67        observed=float(((z2-z1).norm(dim=1)/(delta.view(len(delta),-1).norm(dim=1)+1e-12)).mean().cpu())
68        blocks=net.rnn.weight_hh_l0.detach().cpu().numpy().reshape(3,64,64)
69        pred=max(operator_gain(b) for b in blocks)
70        rows.append({'seed':seed,'predicted_gain_proxy':pred,'observed_output_gain':observed})
71    maxpred=max(r['predicted_gain_proxy'] for r in rows)
72    return {'prediction':'projected recurrent operator-gain proxy stays <= 0.90', 'predicted_max':maxpred,
73            'observed':rows, 'confirmed': bool(maxpred <= .90+1e-4)}
74
75def main():
76    base=sweep_baseline(baseline_factory, GRID, seeds=(0,1,2,3))
77    nearby={str(c['lr']): evaluate(idea_factory(c), seeds=SEEDS) for c in IDEA_GRID}
78    best=min(nearby, key=lambda k: nearby[k]['mean']); idea=nearby[best]
79    rep=make_report('dynamics','rnn_small',base,idea,{'signature':signature(), 'idea_sweep':nearby,
80      'selected_idea_lr':float(best), 'structure':'GRU recurrent dynamics; spectral-gain projection is the only intervention'})
81    rep['baseline_sweep_union_note']='Baseline evaluated all idea learning rates; baseline method has no additional decisive knob.'
82    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
83if __name__=='__main__': main()