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, sweep_baseline, evaluate, make_report
  8
  9SEEDS = tuple(range(8))
 10# Equal search-space union: every idea setting is also evaluated for baseline.
 11GRID = [
 12    {'lr': 0.0015, 'weight_decay': 0.0},
 13    {'lr': 0.0030, 'weight_decay': 0.0},
 14    {'lr': 0.0060, 'weight_decay': 0.0},
 15]
 16EPOCHS = 15
 17BATCH = 128
 18
 19def seed_all(seed):
 20    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 21    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 22
 23def device():
 24    return 'cuda' if torch.cuda.is_available() else 'cpu'
 25
 26def jacobian2(net, h):
 27    """Jacobian of the first two recurrent coordinates, with other h fixed."""
 28    outs = []
 29    for j in range(2):
 30        g = torch.autograd.grad(net.step_hidden(h)[0, j], h,
 31                                create_graph=True, retain_graph=True)[0]
 32        outs.append(g[0, :2])
 33    return torch.stack(outs)
 34
 35def bt_values(net, h0=None, fd=0.05):
 36    """Centered finite-difference BT monitor on the trained recurrent map."""
 37    if h0 is None:
 38        h0 = torch.zeros(1, net.hidden, device=next(net.parameters()).device)
 39    h0 = h0.detach().requires_grad_(True)
 40    J = jacobian2(net, h0)
 41    # SVD gives the coordinate-free smallest right-singular direction.
 42    _, _, vh = torch.linalg.svd(J.detach())
 43    q = vh[-1]
 44    q = q / (torch.linalg.vector_norm(q) + 1e-12)
 45    hp = (h0.detach() + fd * torch.cat((q, torch.zeros(net.hidden-2, device=h0.device))).view(1,-1)).requires_grad_(True)
 46    hm = (h0.detach() - fd * torch.cat((q, torch.zeros(net.hidden-2, device=h0.device))).view(1,-1)).requires_grad_(True)
 47    # For reporting, calculate Jacobians with graph disabled through explicit helper.
 48    def plain(x, differentiable=False):
 49        x = x.detach().requires_grad_(True)
 50        rows=[]
 51        for j in range(2):
 52            g=torch.autograd.grad(net.step_hidden(x)[0,j],x,retain_graph=True,create_graph=differentiable)[0]
 53            rows.append(g[0,:2])
 54        return torch.stack(rows).detach()
 55    jp, jm = plain(hp, differentiable=h0.requires_grad), plain(hm, differentiable=h0.requires_grad)
 56    a = -0.5 * (torch.det(jp)-torch.det(jm))/(2*fd)
 57    b = (torch.trace(jp)-torch.trace(jm))/(2*fd)
 58    return a, b, J.detach(), q.detach()
 59
 60class BTGRU(nn.Module):
 61    """Same rnn_small architecture, exposing its recurrent map for regularization."""
 62    def __init__(self, out_dim=1, hidden=64):
 63        super().__init__(); self.hidden=hidden
 64        self.rnn=nn.GRU(3,hidden,batch_first=True); self.head=nn.Linear(hidden,out_dim)
 65        self._no_cudnn=False
 66    def step_hidden(self,h):
 67        # GRU accepts a length-one zero-input sequence and explicit hidden state.
 68        x=torch.zeros(1,1,3,device=h.device,dtype=h.dtype)
 69        try: _, z=self.rnn(x,h.unsqueeze(0))
 70        except RuntimeError:
 71            old=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
 72            try: _,z=self.rnn(x,h.unsqueeze(0))
 73            finally: torch.backends.cudnn.enabled=old
 74        return z[-1]
 75    def forward(self,x):
 76        seq=x.view(x.shape[0],-1,3)
 77        try: _,h=self.rnn(seq)
 78        except RuntimeError:
 79            old=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
 80            try: _,h=self.rnn(seq)
 81            finally: torch.backends.cudnn.enabled=old
 82        return self.head(h[-1])
 83
 84def train_one(seed, cfg, idea):
 85    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
 86    net=BTGRU(out_dim=1,hidden=64).to(device())
 87    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
 88    xtr,ytr=ds['xtr'].to(device()),ds['ytr'].to(device())
 89    for ep in range(EPOCHS):
 90        net.train(); perm=torch.randperm(len(xtr),device=xtr.device)
 91        for i in range(0,len(xtr),BATCH):
 92            idx=perm[i:i+BATCH]; pred=net(xtr[idx]); loss=((pred-ytr[idx])**2).mean()
 93            if idea:
 94                a,b,_,_=bt_values(net)
 95                bt=0.01*(torch.relu(torch.tensor(0.02,device=xtr.device)-torch.abs(a))**2 + torch.relu(torch.tensor(0.02,device=xtr.device)-torch.abs(b))**2)
 96                loss=loss+bt
 97            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
 98    net.eval()
 99    with torch.no_grad(): metric=float(((net(ds['xte'].to(device()))-ds['yte'].to(device()))**2).mean())
100    with torch.no_grad(): pass
101    a,b,J,q=bt_values(net)
102    return metric, {'a':float(a),'b':float(b),'ab':float(abs(a*b)),'rank_smin':float(torch.linalg.svdvals(J)[-1])}
103
104def train_metric(seed,cfg,idea): return train_one(seed,cfg,idea)[0]
105def run():
106    # Baseline selection uses the same three configs, then full paired evaluation.
107    base=sweep_baseline(lambda cfg: lambda s: train_metric(s,cfg,False), GRID)
108    best=base['best_cfg']
109    idea_grid=[best, {'lr':0.0015,'weight_decay':0.0}, {'lr':0.006,'weight_decay':0.0}]
110    # Report idea at the best of the same-size idea sweep; baseline has all union settings.
111    candidates=[]
112    for cfg in idea_grid:
113        r=evaluate(lambda s,cfg=cfg: train_metric(s,cfg,True), SEEDS)
114        candidates.append((r,cfg))
115    idea_res,best_idea=min(candidates,key=lambda z:z[0]['mean'])
116    sig=[]
117    for s in SEEDS:
118        _,m0=train_one(s,best,False); _,m1=train_one(s,best_idea,True)
119        sig.append({'seed':s,'baseline':m0,'idea':m1})
120    # Quantitative prediction tested on trained models: regularization should reduce near-degenerate |ab|.
121    t=0.02**2
122    bfrac=float(np.mean([x['baseline']['ab']<t for x in sig])); ifrac=float(np.mean([x['idea']['ab']<t for x in sig]))
123    mechanism={'quantity':'trained recurrent-map BT coefficients a,b from centered FD along SVD kernel direction',
124      'predicted':'BT penalty reduces fraction of checkpoints/models with |a*b| below threshold',
125      'threshold':t,'baseline_near_degenerate_fraction':bfrac,'idea_near_degenerate_fraction':ifrac,
126      'per_seed':sig,'confirmed':bool(ifrac < bfrac)}
127    rep=make_report('dynamics','rnn_small',base,idea_res,mechanism)
128    rep['idea_sweep']=[{'cfg':c,'mean':r['mean'],'per_seed':r['per_seed']} for r,c in candidates]
129    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
130    print(json.dumps(rep,indent=2))
131if __name__=='__main__': run()