1import sys, json, math, 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, evaluate, sweep_baseline, make_report
 8
 9EPOCHS=10; NPOP=4; BATCH=128
10LRS=[1e-3,3e-3,1e-2]
11BASE_GRID=[{'lr':lr,'weight_decay':wd} for lr in LRS for wd in (0.0,1e-4)]
12# Same lr union; interaction is the only method difference.
13IDEA_GRID=[{'lr':lr,'alpha':a,'normalized':norm} for lr in LRS for a,norm in ((0.0,True),(0.1,True),(0.1,False))]
14
15def seed_all(s):
16    random.seed(s); np.random.seed(s); torch.manual_seed(s)
17    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
18
19def flat(m): return torch.cat([p.detach().reshape(-1) for p in m.parameters()])
20def put(m,v):
21    k=0
22    with torch.no_grad():
23        for p in m.parameters(): p.copy_(v[k:k+p.numel()].view_as(p)); k+=p.numel()
24def grad(m): return torch.cat([(p.grad if p.grad is not None else torch.zeros_like(p)).reshape(-1) for p in m.parameters()])
25
26def population_train(seed, cfg, n_particles=NPOP, epochs=EPOCHS, weight_decay=0.0, want_sig=False, device=None):
27    seed_all(seed); d=get_dataset('tabular',seed,n_train=800,n_test=400)
28    device=device or ('cuda' if torch.cuda.is_available() else 'cpu')
29    try:
30        ms=[make_model('mlp_tiny',d['input_shape'],d['out_dim']).to(device) for _ in range(n_particles)]
31        x,y=d['xtr'].to(device),d['ytr'].to(device); lossf=nn.MSELoss()
32        # independent random starts are essential for an extreme-search population
33        opt=[torch.optim.Adam(m.parameters(),lr=cfg['lr'],weight_decay=weight_decay) for m in ms]
34        with torch.no_grad():
35            init=float(torch.stack([lossf(m(x),y) for m in ms]).mean())
36        target=init*0.70; first=None; ratios=[]; hist=[]
37        for ep in range(epochs):
38            perm=torch.randperm(len(x),device=device); losses=[]
39            for q in range(0,len(x),BATCH):
40                ix=perm[q:q+BATCH]; gs=[]
41                for m,o in zip(ms,opt):
42                    m.train(); o.zero_grad(set_to_none=True); z=lossf(m(x[ix]),y[ix]); z.backward(); gs.append(grad(m))
43                th=torch.stack([flat(m) for m in ms]); disp=th[None]-th[:,None]
44                psi=disp/(disp.norm(dim=2,keepdim=True).clamp_min(1e-8)); ii=torch.arange(n_particles,device=device); psi[ii,ii]=0
45                force=cfg.get('alpha',0.0)*psi.sum(1)
46                if cfg.get('normalized',True) and n_particles>1: force/=n_particles-1
47                gg=torch.stack(gs); update=[-g+f for g,f in zip(gg,force)]
48                if cfg.get('alpha',0.0): ratios.append(float(force.norm(dim=1).mean()/(gg.norm(dim=1).mean()+1e-12)))
49                # Adam update with interaction added to the gradient direction.
50                for m,o,u in zip(ms,opt,update):
51                    o.zero_grad(set_to_none=True)
52                    put(m,flat(m)+cfg['lr']*u/(u.square().sqrt()+1e-8))
53                losses.extend([float(lossf(m(x[ix]),y[ix]).detach()) for m in ms])
54            with torch.no_grad(): full=torch.stack([lossf(m(x),y) for m in ms])
55            hist.append(float(full.min()))
56            if first is None and float(full.min())<=target: first=ep+1
57        with torch.no_grad():
58            xt,yt=d['xte'].to(device),d['yte'].to(device)
59            test=[float(lossf(m(xt),yt)) for m in ms]
60        return {'metric':min(test),'first_hit_epoch':first or epochs+1,'force_ratio':float(np.mean(ratios)) if ratios else 0.0,'history':hist}
61    except Exception:
62        if device=='cuda':
63            torch.cuda.empty_cache(); return population_train(seed,cfg,n_particles,epochs,weight_decay,want_sig,'cpu')
64        raise
65
66def baseline_fn(c):
67    return lambda s: population_train(s,{'lr':c['lr'],'alpha':0.0,'normalized':True},weight_decay=c['weight_decay'])['metric']
68def idea_fn(c): return lambda s: population_train(s,c)['metric']
69
70def main():
71    base=sweep_baseline(baseline_fn,BASE_GRID)
72    trials=[]
73    for c in IDEA_GRID: trials.append((c,evaluate(idea_fn(c))))
74    best_cfg,best=min(trials,key=lambda z:z[1]['mean'])
75    # Re-test the NN-scale prediction on trained populations; target is fixed relative to each run's initial loss.
76    sig=[]
77    for n in (2,4,8):
78        oo=[population_train(s,best_cfg,n_particles=n,want_sig=True) for s in (0,1,2,3)]
79        sig.append({'N':n,'mean_test_metric':float(np.mean([o['metric'] for o in oo])),'mean_first_hit_epoch':float(np.mean([o['first_hit_epoch'] for o in oo])),'mean_force_to_gradient':float(np.mean([o['force_ratio'] for o in oo]))})
80    slope=float(np.polyfit(np.log([z['N'] for z in sig]),np.log([z['mean_first_hit_epoch'] for z in sig]),1)[0])
81    mechanism={'prediction':'unnormalized coherent force should yield algebraic first-hit acceleration (ideal slope -1), unlike normalized bounded force','observed_NN_scale':sig,'observed_loglog_first_hit_slope':slope,'confirmed':bool(best_cfg['alpha']>0 and not best_cfg['normalized'] and slope< -0.3)}
82    rep=make_report('tabular','mlp_tiny',base,best,{'prediction':'unnormalized coherent interaction','best_cfg':best_cfg,'mechanism_signature':mechanism})
83    rep['idea_sweep']=[{'cfg':c,'mean':r['mean'],'per_seed':r['per_seed']} for c,r in trials]
84    rep['equal_budget']={'baseline':'4 independent Adam particles','idea':'4 interacting Adam particles','epochs_per_particle':EPOCHS,'batch':BATCH}
85    Path('bench_report.json').write_text(json.dumps(rep,indent=2)); print(json.dumps(rep,indent=2))
86if __name__=='__main__': main()