1import sys, json, math, 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, evaluate, sweep_baseline, make_report
  7
  8TRACK='tabular'; MODEL='mlp_tiny'; SEEDS=tuple(range(8))
  9# Union grid is used for both methods: baseline receives every lr tried by idea.
 10LRS=[1e-3, 3e-3, 6e-3]
 11TRAINED_SIGNATURE=[]
 12EPOCHS=18; BATCH=64
 13
 14def seed_all(seed):
 15    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 16    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 17
 18def loss_fn(ds): return nn.MSELoss()
 19
 20def train_baseline(seed, cfg):
 21    seed_all(seed); ds=get_dataset(TRACK, seed, 400, 200)
 22    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
 23    # Baseline is standard Adam; custom loop is used so the model and minibatch
 24    # ordering are exactly matched to the intervention.
 25    return run(net, ds, cfg['lr'], pulse=False, seed=seed)[0]
 26
 27def run(net, ds, lr, pulse=False, seed=0):
 28    device='cuda' if torch.cuda.is_available() else 'cpu'
 29    try:
 30        return _run(net, ds, lr, pulse, seed, device)
 31    except RuntimeError:
 32        return _run(net.cpu(), ds, lr, pulse, seed, 'cpu')
 33
 34def _run(net, ds, lr, pulse, seed, device):
 35    net=net.to(device); x=ds['xtr'].to(device); y=ds['ytr'].to(device)
 36    opt=torch.optim.Adam(net.parameters(), lr=lr)
 37    lf=loss_fn(ds); rng=torch.Generator(device=device); rng.manual_seed(seed+991)
 38    qs=[]; responses=[]; triggers=0; recent=[]
 39    for ep in range(EPOCHS):
 40        net.train(); perm=torch.randperm(len(x), generator=rng, device=device)
 41        # Probe every 4 epochs. q is squared parameter displacement / lr.
 42        pulse_now=False
 43        if pulse and ep > 0 and ep % 4 == 0:
 44            # short temperature pulse = doubled batch noise proxy by halving batch
 45            pulse_now=True
 46        batch=BATCH//2 if pulse_now else BATCH
 47        epq=[]
 48        for i in range(0,len(x),batch):
 49            idx=perm[i:i+batch]; old=[p.detach().clone() for p in net.parameters() if p.requires_grad]
 50            out=net(x[idx].view(len(idx),-1)); loss=lf(out,y[idx])
 51            opt.zero_grad(); loss.backward(); opt.step()
 52            q=sum(((p.detach()-o)**2).sum().item() for p,o in zip((p for p in net.parameters() if p.requires_grad),old))/max(lr,1e-12)
 53            epq.append(q)
 54        qmean=float(np.mean(epq)); qs.append(qmean)
 55        if pulse_now:
 56            base=float(np.mean(recent[-2:])) if recent else qmean
 57            excess=qmean-base
 58            responses.append({'epoch':ep,'q_ss':base,'q_pulse':qmean,'C_train':excess/(1.0/(len(x)))})
 59            # detector only changes training when dissipation is an outlier.
 60            if len(recent)>=2 and qmean > 1.8*float(np.median(recent[-4:])):
 61                for g in opt.param_groups: g['lr']*=0.7
 62                triggers+=1
 63        else: recent.append(qmean)
 64    net.eval()
 65    with torch.no_grad(): metric=float(((net(ds['xte'].to(device)).view(-1,1)-ds['yte'].to(device))**2).mean())
 66    return metric, {'q':qs,'responses':responses,'triggers':triggers}
 67
 68def train_idea(seed,cfg):
 69    seed_all(seed); ds=get_dataset(TRACK, seed, 400, 200)
 70    net=make_model(MODEL, ds['input_shape'], ds['out_dim'])
 71    metric, detail = run(net, ds, cfg['lr'], pulse=True, seed=seed)
 72    TRAINED_SIGNATURE.append(detail)
 73    return metric
 74
 75def main():
 76    # Cheap numerical verification of the transfer formula and stability claim.
 77    eta=np.linspace(.2,1.99,10); lam=1.; T=1.; delta=.05; P=20
 78    Cs=[]
 79    for e in eta:
 80        a=1-e*lam; v=2*e*T/(1-a*a); v2=v; q0=4*T/(1+a); ex=0
 81        for _ in range(P):
 82            v2=a*a*v2+2*e*(T+delta); ex += (1-a)**2*v2/e+2*(T+delta)-q0
 83        Cs.append(ex/delta)
 84    sanity={'stable_below_2':bool(np.all(np.abs(1-eta)<1)), 'C_monotone':bool(np.all(np.diff(Cs)>0)), 'C_values':Cs}
 85    grid=[{'lr':v} for v in LRS]
 86    base=sweep_baseline(lambda c: (lambda s: train_baseline(s,c)),grid)
 87    # Three idea settings are exactly the shared union grid, with best selected.
 88    idea_trials=[]
 89    for c in grid:
 90        r=evaluate(lambda s,c=c: train_idea(s,c),SEEDS)
 91        idea_trials.append({'cfg':c,'mean':r['mean']})
 92    best=min(idea_trials,key=lambda z:z['mean'])['cfg']
 93    idea=evaluate(lambda s: train_idea(s,best),SEEDS)
 94    ratios=[r['q_pulse']/r['q_ss'] for d in TRAINED_SIGNATURE for r in d.get('responses',[]) if r['q_ss']>1e-12]
 95    trained={'n_pulses':len(ratios), 'mean_pulse_to_baseline_q_ratio':float(np.mean(ratios)) if ratios else None, 'mean_triggers':float(np.mean([d.get('triggers',0) for d in TRAINED_SIGNATURE])) if TRAINED_SIGNATURE else 0.0}
 96    signature={'prediction':'temperature/minibatch-noise pulse produces transient excess update dissipation', 'predicted':{'pulse_to_baseline_q_ratio':'>1'}, 'observed_on_trained_models':trained, 'confirmed':bool(ratios and trained['mean_pulse_to_baseline_q_ratio']>1.0), 'toy_sanity':sanity}
 97    rep=make_report(TRACK,MODEL,base,idea,signature)
 98    rep['idea']['sweep']=idea_trials; rep['sanity_check']=sanity
 99    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
100    print(json.dumps(rep,indent=2))
101if __name__=='__main__': main()