1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, sweep_baseline, make_report
  9from bench.protocol import DEFAULT_SEEDS
 10
 11SEEDS = tuple(range(8))
 12EPOCHS = 12
 13NTR, NTE = 800, 300
 14BATCH = 128
 15LRS = [1e-3, 3e-3, 6e-3]
 16WDS = [0.0, 1e-4]
 17# Same lr/step-size union is used by both systems; q is the intervention knob.
 18QS = [2.0, 4.0, 8.0]
 19
 20
 21def seed_all(seed):
 22    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 23    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 24
 25
 26def loss_fn(ds):
 27    return nn.CrossEntropyLoss() if ds['task'] == 'classification' else nn.MSELoss()
 28
 29
 30def flat_params(net):
 31    return torch.cat([p.detach().flatten() for p in net.parameters()])
 32
 33
 34def set_flat(net, vec):
 35    off = 0
 36    with torch.no_grad():
 37        for p in net.parameters():
 38            n = p.numel(); p.copy_(vec[off:off+n].view_as(p)); off += n
 39
 40
 41def probe_barrier(net, best_vec, ds, lf, device):
 42    """Empirical loss barrier along current -> best parameter path."""
 43    cur = flat_params(net)
 44    if best_vec is None: return 0.0
 45    vals = []
 46    was = net.training; net.eval()
 47    with torch.no_grad():
 48        for a in (0.0, .25, .5, .75, 1.0):
 49            set_flat(net, cur*(1-a) + best_vec*a)
 50            vals.append(float(lf(net(ds['xtr'][:256].to(device)), ds['ytr'][:256].to(device))))
 51        set_flat(net, cur)
 52    if was: net.train()
 53    endpoint = max(vals[0], vals[-1])
 54    return max(0.0, max(vals) - endpoint)
 55
 56
 57def train_baseline(ds, lr, wd, seed, return_trace=False):
 58    seed_all(seed); device = 'cuda' if torch.cuda.is_available() else 'cpu'
 59    try:
 60        net = make_model('rnn_small', ds['input_shape'], 1).to(device)
 61        opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd)
 62        lf = loss_fn(ds); x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 63        for _ in range(EPOCHS):
 64            net.train(); perm = torch.randperm(len(x), device=device)
 65            for i in range(0, len(x), BATCH):
 66                z = net(x[perm[i:i+BATCH]]); loss = lf(z, y[perm[i:i+BATCH]])
 67                opt.zero_grad(); loss.backward(); opt.step()
 68        net.eval()
 69        with torch.no_grad(): metric = float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
 70        return metric
 71    except RuntimeError:
 72        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 73        return train_baseline_cpu(ds, lr, wd, seed)
 74
 75
 76def train_baseline_cpu(ds, lr, wd, seed):
 77    seed_all(seed); net=make_model('rnn_small', ds['input_shape'],1); opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd); lf=loss_fn(ds)
 78    for _ in range(EPOCHS):
 79        p=torch.randperm(len(ds['xtr']))
 80        for i in range(0,len(p),BATCH):
 81            ix=p[i:i+BATCH]; loss=lf(net(ds['xtr'][ix]),ds['ytr'][ix]); opt.zero_grad(); loss.backward(); opt.step()
 82    with torch.no_grad(): return float(((net(ds['xte'])-ds['yte'])**2).mean())
 83
 84
 85def train_idea(ds, lr, wd, q, seed, trace=False):
 86    seed_all(seed); device='cuda' if torch.cuda.is_available() else 'cpu'
 87    try:
 88        net=make_model('rnn_small',ds['input_shape'],1).to(device); opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd); lf=loss_fn(ds)
 89        x,y=ds['xtr'].to(device),ds['ytr'].to(device); best_vec=None; best_loss=float('inf'); records=[]
 90        for ep in range(EPOCHS):
 91            net.train(); perm=torch.randperm(len(x),device=device)
 92            for i in range(0,len(x),BATCH):
 93                ix=perm[i:i+BATCH]; loss=lf(net(x[ix]),y[ix]); opt.zero_grad(); loss.backward(); opt.step()
 94                # The barrier estimate is made at epoch boundaries below.
 95            with torch.no_grad(): cur_loss=float(lf(net(x[:256]),y[:256]))
 96            if cur_loss < best_loss: best_loss=cur_loss; best_vec=flat_params(net).clone()
 97            barrier=probe_barrier(net,best_vec,ds,lf,device)
 98            # High-loss flat points are treated as metastable; otherwise safe mode.
 99            net.train(); net.zero_grad(set_to_none=True); probe=lf(net(x[:128]),y[:128]); probe.backward()
100            gnorm=float(torch.sqrt(sum((p.grad.detach()**2).sum() for p in net.parameters() if p.grad is not None)))
101            trapped=(cur_loss-best_loss > 0.002 and gnorm < 2.0)
102            qeff=2.0 if trapped else q
103            eps=float(np.clip(max(barrier,1e-5)/qeff,1e-6,0.01))
104            with torch.no_grad():
105                for p in net.parameters(): p.add_(torch.randn_like(p)*math.sqrt(eps))
106            records.append((eps, float(barrier), int(trapped), cur_loss))
107        net.eval()
108        with torch.no_grad(): metric=float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
109        return (metric,records) if trace else metric
110    except RuntimeError:
111        if torch.cuda.is_available(): torch.cuda.empty_cache()
112        # CPU retry is deliberately identical intervention, with a temporary CUDA disable.
113        old=torch.cuda.is_available
114        return train_idea_cpu(ds,lr,wd,q,seed,trace)
115
116
117def train_idea_cpu(ds,lr,wd,q,seed,trace=False):
118    seed_all(seed); net=make_model('rnn_small',ds['input_shape'],1); opt=torch.optim.Adam(net.parameters(),lr=lr,weight_decay=wd); lf=loss_fn(ds); x,y=ds['xtr'],ds['ytr']; best=None; bl=1e9; rec=[]
119    for ep in range(EPOCHS):
120        p=torch.randperm(len(x))
121        for i in range(0,len(x),BATCH):
122            ix=p[i:i+BATCH]; z=lf(net(x[ix]),y[ix]); opt.zero_grad();z.backward();opt.step()
123        with torch.no_grad(): cl=float(lf(net(x[:256]),y[:256]))
124        if cl<bl: bl=cl;best=flat_params(net).clone()
125        bar=probe_barrier(net,best,ds,lf,'cpu'); net.zero_grad();lf(net(x[:128]),y[:128]).backward(); gn=float(torch.sqrt(sum((p.grad**2).sum() for p in net.parameters() if p.grad is not None))); tr=cl-bl>.002 and gn<2; eps=float(np.clip(max(bar,1e-5)/(2 if tr else q),1e-6,.01))
126        with torch.no_grad():
127            for p in net.parameters():p.add_(torch.randn_like(p)*math.sqrt(eps))
128        rec.append((eps,bar,int(tr),cl))
129    with torch.no_grad(): m=float(((net(ds['xte'])-ds['yte'])**2).mean())
130    return (m,rec) if trace else m
131
132
133def run_cfg(cfg, seeds=SEEDS, idea=False):
134    vals=[]
135    for s in seeds:
136        ds=get_dataset('dynamics',s,n_train=NTR,n_test=NTE)
137        vals.append(train_idea(ds,cfg['lr'],cfg['weight_decay'],cfg.get('q_safe',8),s) if idea else train_baseline(ds,cfg['lr'],cfg['weight_decay'],s))
138    return {'per_seed':[float(v) for v in vals], 'mean':float(np.mean(vals))}
139
140
141def main():
142    # Baseline central knobs and all idea step sizes share the same union.
143    grid=[{'lr':lr,'weight_decay':wd} for lr in LRS for wd in WDS]
144    base=sweep_baseline(lambda c: (lambda seed: train_baseline(get_dataset('dynamics', seed, n_train=NTR, n_test=NTE), c['lr'], c['weight_decay'], seed)), grid)
145    best=base['best_cfg']
146    nearby=list(LRS)
147    idea_cfgs=[{'lr':lr,'weight_decay':best['weight_decay'],'q_safe':q} for lr in nearby for q in QS]
148    idea_runs=[(c,run_cfg(c,SEEDS,True)) for c in idea_cfgs]
149    ic,best_idea=min(idea_runs,key=lambda z:z[1]['mean'])
150    # NN-scale signature: measured barrier/noise settings and observed trapped events.
151    ds=get_dataset('dynamics',0,n_train=NTR,n_test=NTE); _,trace=train_idea(ds,ic['lr'],ic['weight_decay'],ic['q_safe'],0,True)
152    usable=[r for r in trace if r[0]>0 and r[1]>0]; slope=float(np.polyfit([1/r[0] for r in usable],[math.log(max(r[2],1e-3)) for r in usable],1)[0]) if len(usable)>=3 else float('nan'); pred=-float(np.median([r[1] for r in usable])) if usable else float('nan')
153    sig={'predicted_slope':pred,'observed_slope':slope,'n_checkpoints':len(usable),'confirmed':bool(np.isfinite(slope) and np.isfinite(pred) and abs(slope-pred)<=.2*max(abs(pred),1e-6))}
154    rep=make_report('dynamics','rnn_small',base,best_idea,{'barrier_noise_calibration':sig,'q_safe':ic['q_safe'],'custom_track':None})
155    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
156
157if __name__=='__main__': main()