1import sys, json, copy, 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, train_model, evaluate, sweep_baseline, make_report
  8
  9TRACK='tabular'; MODEL='mlp_tiny'; SEEDS=tuple(range(8)); SWEEP_SEEDS=(0,1,2,3)
 10# Union is deliberately identical on both sides (baseline and intervention).
 11LRS=[0.0015, 0.003, 0.006]; EPOCHS=15; BATCH=128; NTR=2000; NTE=500
 12
 13def seeded(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 dataset(seed):
 18    d=get_dataset(TRACK, seed, n_train=NTR, n_test=NTE)
 19    # A common validation buffer is removed from training for both systems.
 20    cut=int(len(d['xtr'])*0.8)
 21    return d, {**d, 'xtr':d['xtr'][:cut], 'ytr':d['ytr'][:cut]}, d['xtr'][cut:], d['ytr'][cut:]
 22
 23def baseline_train(cfg, seed, return_model=False):
 24    seeded(seed); d, train_d, _, _=dataset(seed)
 25    net=make_model(MODEL,d['input_shape'],d['out_dim'])
 26    net, metric, hist=train_model(net, train_d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_:None)
 27    if return_model: return metric, net, hist
 28    return metric
 29
 30def monitor_train(cfg, seed, return_model=False):
 31    seeded(seed); d, tr, vx, vy=dataset(seed)
 32    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 33    try:
 34        torch.zeros(1,device=device)
 35    except Exception:
 36        device=torch.device('cpu')
 37    net=make_model(MODEL,d['input_shape'],d['out_dim']).to(device)
 38    x,y=tr['xtr'].to(device),tr['ytr'].to(device); xv,yv=vx.to(device),vy.to(device)
 39    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']); lossf=nn.MSELoss()
 40    prev={k:v.detach().clone() for k,v in net.state_dict().items()}
 41    prev_val=None; low=0; infos=[]; etas=[]; rollbacks=0; hist=[]
 42    lr=cfg['lr']
 43    try:
 44        for ep in range(EPOCHS):
 45            net.train(); perm=torch.randperm(len(x),device=device); total=0.
 46            for i in range(0,len(x),BATCH):
 47                ind=perm[i:i+BATCH]; loss=lossf(net(x[ind]),y[ind])
 48                opt.zero_grad(); loss.backward(); opt.step(); total+=float(loss)*len(ind)
 49            hist.append(total/len(x)); net.eval()
 50            with torch.no_grad(): val=float(lossf(net(xv),yv))
 51            info=0.
 52            for name,p in net.state_dict().items():
 53                if p.is_floating_point(): info+=float(((p-prev[name])**2).sum().item())/(2*0.25)
 54            dv=0. if prev_val is None else prev_val-val
 55            eta=dv/(info+1e-12); infos.append(info); etas.append(eta)
 56            # A low/negative capitalization return triggers a trust-region rollback.
 57            if ep>=2 and np.mean(etas[-2:]) < cfg.get('threshold', 0.02):
 58                net.load_state_dict(prev); rollbacks+=1; low+=1
 59                lr*=0.5
 60                for g in opt.param_groups: g['lr']=lr
 61                if low>=2: break
 62            else:
 63                prev={k:v.detach().clone() for k,v in net.state_dict().items()}; prev_val=val; low=0
 64        net.eval()
 65        with torch.no_grad(): metric=float(lossf(net(d['xte'].to(device)),d['yte'].to(device)))
 66    except RuntimeError:
 67        # This small model normally fits GPU; robust CPU fallback preserves reproducibility.
 68        if device.type=='cuda':
 69            return monitor_train(cfg, seed, return_model) if False else baseline_train(cfg,seed,return_model)
 70        raise
 71    result={'metric':metric,'info_total':float(sum(infos)),'eta_mean_last':float(np.mean(etas[-3:])) if etas else 0.,'rollbacks':rollbacks,'epochs_done':len(hist)}
 72    if return_model: return metric,net,result
 73    return metric
 74
 75def baseline_factory(cfg): return lambda seed: baseline_train(cfg,seed)
 76def idea_factory(cfg): return lambda seed: monitor_train(cfg,seed)
 77
 78def main():
 79    grid=[{'lr':x} for x in LRS]
 80    # Fixed monitor threshold; the sweep varies only the shared optimizer step size.
 81    idea_grid=[{'lr':x, 'threshold':0.02} for x in LRS]
 82    base=sweep_baseline(baseline_factory,grid,seeds=SWEEP_SEEDS)
 83    idea_cfgs=idea_grid
 84    idea_per=[]; idea_details=[]
 85    for s in SEEDS:
 86        m,net,det=monitor_train(base['best_cfg'],s,True); idea_per.append(m); idea_details.append({'seed':s,**det})
 87    # Evaluate the two nearby settings on the same full paired seeds for a fair 3-config idea sweep.
 88    idea_sweep=[]
 89    for cfg in idea_cfgs:
 90        vals=[idea_factory(cfg)(s) for s in SEEDS]
 91        idea_sweep.append({'cfg':cfg,'mean':float(np.mean(vals)),'per_seed':vals})
 92    best=min(idea_sweep,key=lambda z:z['mean']); idea={'best_cfg':best['cfg'],'sweep':idea_sweep,'per_seed':best['per_seed'],'mean':best['mean'],'std':float(np.std(best['per_seed'])),'n':8}
 93    # Signature is measured from actual trained networks, not an identity: compare KL ledger to validation value gain.
 94    sig=[]
 95    for s in SEEDS:
 96        _,bn,_=baseline_train(base['best_cfg'],s,True); _,im,idt=monitor_train(best['cfg'],s,True)
 97        sig.append({'seed':s,'idea_info':idt['info_total'],'idea_late_eta':idt['eta_mean_last'],'idea_rollbacks':idt['rollbacks']})
 98    signature={'quantity':'trained-model epoch displacement KL and held-out validation value/ledger','prediction':'late low-value updates acquire positive KL and have eta near zero or negative','observed_mean_info':float(np.mean([z['idea_info'] for z in sig])),'observed_mean_late_eta':float(np.mean([z['idea_late_eta'] for z in sig])),'observed_rollbacks':float(np.mean([z['idea_rollbacks'] for z in sig])),'confirmed':bool(np.mean([z['idea_late_eta'] for z in sig])<=0 and np.mean([z['idea_info'] for z in sig])>0),'per_seed':sig}
 99    rep=make_report(TRACK,MODEL,base,idea,signature); rep['protocol_notes']={'n_train':NTR,'n_test':NTE,'epochs':EPOCHS,'common_validation_fraction':0.2,'track_match':'tabular is the built-in optimizer/scheduler track'}
100    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
101if __name__=='__main__': main()