1import os, sys, json, 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 evaluate
  8
  9# Characteristic-invariant BT monitor on a measured 2D local recurrent chart.
 10# The architecture remains exactly bench rnn_small; h[0:2] are perturbed and
 11# the other recurrent coordinates are held at zero for a cheap local monitor.
 12def monitor(model):
 13    rnn = model.rnn
 14    W = rnn.weight_hh_l0
 15    b = rnn.bias_ih_l0 + rnn.bias_hh_l0
 16    z = torch.zeros(2, device=W.device, dtype=W.dtype, requires_grad=True)
 17    h = torch.cat([z, torch.zeros(W.shape[1]-2, device=W.device, dtype=W.dtype)])
 18    # PyTorch GRU gate order is reset, update, new.
 19    gates = torch.mv(W, h) + b
 20    r, u, npre = gates.chunk(3, 0)
 21    rr, uu = torch.sigmoid(r), torch.sigmoid(u)
 22    nnv = torch.tanh(npre + rr * 0)  # zero external input; recurrent reset is in Wn below
 23    # Recompute candidate with reset applied to recurrent hidden contribution.
 24    whr, whu, whn = W.chunk(3, 0)
 25    bir, biu, bin_ = b.chunk(3, 0)
 26    rr = torch.sigmoid(torch.mv(whr, h) + bir)
 27    uu = torch.sigmoid(torch.mv(whu, h) + biu)
 28    nnv = torch.tanh(torch.mv(whn, rr*h) + bin_)
 29    hp = (1-uu)*nnv + uu*h
 30    out = hp[:2]
 31    rows=[]
 32    for i in range(2):
 33        rows.append(torch.autograd.grad(out[i], z, create_graph=True, retain_graph=True)[0])
 34    A = torch.stack(rows)
 35    # Invariants of the continuous vector field f=h'-h: J=A-I.
 36    J = A - torch.eye(2, device=A.device, dtype=A.dtype)
 37    e1 = torch.trace(J); e2 = torch.linalg.det(J)
 38    delta, tau = e2, e1
 39    q = torch.linalg.svd(J.detach()).Vh[-1]
 40    gd = torch.autograd.grad(delta, z, create_graph=True, retain_graph=True)[0]
 41    gt = torch.autograd.grad(tau, z, create_graph=True, retain_graph=True)[0]
 42    a = -0.5 * torch.dot(gd, q)
 43    bb = torch.dot(gt, q)
 44    return delta, tau, a, bb, J.detach()
 45
 46def seed_all(s):
 47    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 48
 49def train_idea(seed, cfg, return_model=False):
 50    seed_all(seed)
 51    ds=get_dataset('dynamics', seed, n_train=400, n_test=400)
 52    model=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 53    dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 54    try:
 55        model.to(dev); x,y=ds['xtr'].to(dev),ds['ytr'].to(dev)
 56        opt=torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 57        lossf=nn.MSELoss()
 58        for ep in range(cfg['epochs']):
 59            model.train(); perm=torch.randperm(len(x),device=dev)
 60            for ix in perm.split(128):
 61                pred=model(x[ix]); loss=lossf(pred,y[ix])
 62                d,t,a,b,_=monitor(model)
 63                loss=loss + cfg['rho']*(d*d+t*t)
 64                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
 65        model.eval()
 66        with torch.no_grad(): metric=float(lossf(model(ds['xte'].to(dev)),ds['yte'].to(dev)).cpu())
 67    except Exception:
 68        # Required robust fallback: restart fully on CPU after any CUDA failure.
 69        seed_all(seed); dev=torch.device('cpu'); model=make_model('rnn_small',ds['input_shape'],ds['out_dim']).to(dev)
 70        x,y=ds['xtr'],ds['ytr']; opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
 71        for ep in range(cfg['epochs']):
 72            for ix in torch.randperm(len(x)).split(128):
 73                loss=lossf(model(x[ix]),y[ix]); d,t,a,b,_=monitor(model); loss=loss+cfg['rho']*(d*d+t*t)
 74                opt.zero_grad(); loss.backward(); opt.step()
 75        with torch.no_grad(): metric=float(lossf(model(ds['xte']),ds['yte']))
 76    if return_model: return metric, model, ds
 77    return metric
 78
 79def train_base(seed,cfg, return_model=False):
 80    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=400)
 81    model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 82    raw=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,weight_decay=cfg['weight_decay'],log=lambda *_:None)
 83    metric = float(raw[1] if isinstance(raw,(tuple,list)) else raw.get('test', raw.get('metric')) if isinstance(raw,dict) else raw)
 84    return (metric,model,ds) if return_model else metric
 85
 86def signature(cfg):
 87    vals=[]
 88    for s in range(8):
 89        m,model,ds=train_idea(s,cfg,True)
 90        d,t,a,b,J=monitor(model)
 91        eig=torch.linalg.eigvals(J).cpu().numpy()
 92        # observed recurrent persistence is spectral radius of the trained local map.
 93        vals.append((float(abs(d)),float(abs(t)),float(np.max(np.abs(eig+1))),float(np.max(eig.real))))
 94    v=np.asarray(vals)
 95    return {'prediction':'BT monitor drives delta,tau toward zero and the trained local dynamics toward a near-zero continuous eigenvalue','trained_mean_abs_delta':float(v[:,0].mean()),'trained_mean_abs_tau':float(v[:,1].mean()),'observed_mean_discrete_radius':float(v[:,2].mean()),'observed_mean_continuous_max_real':float(v[:,3].mean()),'confirmed':bool(v[:,0].mean()<0.08 and v[:,1].mean()<0.08 and v[:,2].mean()>0.85)}
 96
 97def main():
 98    base_grid=[{'lr':x,'weight_decay':w,'epochs':8} for x in (0.001,0.003,0.01) for w in (0.0,1e-4)]
 99    # Same union of decisive Adam knobs is used for baseline and idea.
100    base=sweep_baseline(lambda c: (lambda s: train_base(s,c)),base_grid)
101    idea_grid=[dict(base['best_cfg'],rho=r) for r in (0.001,0.01,0.05)]
102    # idea uses best baseline lr/wd and two a-priori regularizer strengths
103    idea_runs=[(c,evaluate(lambda s,c=c: train_idea(s,c))) for c in idea_grid]
104    best_i,best_r=min(idea_runs,key=lambda z:z[1]['mean'])
105    rep=make_report('dynamics','rnn_small',base,best_r,{'idea_config':best_i,'idea_sweep':[{'cfg':c,'mean':r['mean']} for c,r in idea_runs],**signature(best_i)})
106    rep['custom_track']=None
107    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
108    print(json.dumps(rep,indent=2))
109if __name__=='__main__': main()