1import json, random, sys
  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, sweep_baseline, make_report, evaluate
  7
  8SEEDS = tuple(range(8))
  9# Union-parity: every lr tried by the idea is also evaluated by baseline.
 10GRID = [
 11    {'lr': 0.0003, 'beta2': 0.99},
 12    {'lr': 0.0010, 'beta2': 0.99},
 13    {'lr': 0.0030, 'beta2': 0.999},
 14]
 15IDEA_GRID = [
 16    {'lr': 0.0003, 'beta2': 0.99, 'metric_beta': -0.5},
 17    {'lr': 0.0010, 'beta2': 0.99, 'metric_beta': -0.5},
 18    {'lr': 0.0030, 'beta2': 0.999, 'metric_beta': -0.5},
 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
 25def train_one(seed, cfg, idea=False, collect=False):
 26    seed_all(seed)
 27    ds = get_dataset('tabular', seed, n_train=400, n_test=200)
 28    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 29    use_cuda = torch.cuda.is_available()
 30    device = 'cuda' if use_cuda else 'cpu'
 31    try:
 32        net.to(device)
 33    except Exception:
 34        device = 'cpu'; net.to(device)
 35    x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 36    params = list(net.parameters())
 37    mtrace=[]
 38    vtrace=[]
 39    if idea:
 40        v = [torch.zeros_like(p) for p in params]
 41    else:
 42        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], betas=(0.9, cfg['beta2']))
 43    loss_fn = nn.MSELoss(); violations = 0.0; hist=[]
 44    for ep in range(12):
 45        perm = torch.randperm(len(x), device=device)
 46        total = 0.0
 47        for start in range(0, len(x), 128):
 48            idx = perm[start:start+128]
 49            loss = loss_fn(net(x[idx]), y[idx])
 50            net.zero_grad(set_to_none=True); loss.backward()
 51            if not idea:
 52                opt.step()
 53            else:
 54                # State-dependent diagonal metric M=diag(exp(beta log(vhat))).
 55                with torch.no_grad():
 56                    for j,p in enumerate(params):
 57                        g = p.grad
 58                        v[j].mul_(cfg['beta2']).addcmul_(g, g, value=1-cfg['beta2'])
 59                    step = ep * ((len(x)+127)//128) + start//128 + 1
 60                    for j,p in enumerate(params):
 61                        g=p.grad
 62                        vhat=v[j] / (1.0-cfg['beta2']**step)
 63                        logm=torch.clamp(cfg['metric_beta']*torch.log(vhat+1e-8), np.log(1e-3), np.log(10.0))
 64                        m=torch.exp(logm)
 65                        z=p - cfg['lr']*m*g
 66                        proposal=torch.clamp(z, -2.0, 2.0)
 67                        p.add_(proposal-p)
 68                        mtrace.append(float(m.mean().cpu()))
 69                        vtrace.append(float(vhat.mean().cpu()))
 70                        violations=max(violations, float(torch.relu(-2.0-p).max().cpu()), float(torch.relu(p-2.0).max().cpu()))
 71            total += float(loss.detach().cpu()) * len(idx)
 72        hist.append(total/len(x))
 73    net.eval()
 74    with torch.no_grad():
 75        pred=net(ds['xte'].to(device)); metric=float(((pred-ds['yte'].to(device))**2).mean().cpu())
 76    result={'metric':metric, 'history':hist, 'max_constraint_violation':violations}
 77    if collect:
 78        result['model']=net; result['signature_metric_means']=mtrace; result['signature_variances']=vtrace
 79    return result
 80
 81def baseline_factory(cfg):
 82    return lambda seed: train_one(seed, cfg, idea=False)['metric']
 83
 84def idea_eval(cfg):
 85    vals=[]
 86    for s in SEEDS:
 87        vals.append(train_one(s,cfg,idea=True)['metric'])
 88    return {'per_seed':vals, 'mean':float(np.mean(vals))}
 89
 90def signature(seed, base_cfg, idea_cfg):
 91    b=train_one(seed,base_cfg,False,True); q=train_one(seed,idea_cfg,True,True)
 92    # Prediction tested on trained NN: inverse-variance metric gives larger m
 93    # to coordinates with smaller observed second moments.
 94    a=np.asarray(q['signature_metric_means']); vv=np.asarray(q['signature_variances'])
 95    ok=len(a)>2 and np.isfinite(a).all() and np.isfinite(vv).all()
 96    corr=float(np.corrcoef(np.log(a+1e-12), -np.log(vv+1e-12))[0,1]) if ok else float('nan')
 97    confirmed=bool(ok and corr > 0.8)
 98    return {'predicted':'m decreases as observed v increases (inverse-variance metric)',
 99            'observed_metric_mean':float(np.mean(a)) if len(a) else None,
100            'observed_log_metric_vs_negative_log_variance_corr':corr,
101            'baseline_final_train_loss':b['history'][-1],
102            'idea_final_train_loss':q['history'][-1],
103            'max_constraint_violation':q['max_constraint_violation'],
104            'confirmed':confirmed}
105
106def main():
107    base=sweep_baseline(baseline_factory, GRID, seeds=(0,1,2,3))
108    best=base['best_cfg']
109    idea_cfgs=[c for c in IDEA_GRID if c['lr'] in {x['lr'] for x in GRID}]
110    idea_runs=[(c,idea_eval(c)) for c in idea_cfgs]
111    best_idea_cfg,best_idea=max(idea_runs,key=lambda z:z[1]['mean']) if False else min(idea_runs,key=lambda z:z[1]['mean'])
112    extra=signature(0,best,best_idea_cfg)
113    rep=make_report('tabular','mlp_tiny',base,best_idea,extra)
114    rep['idea']['best_cfg']=best_idea_cfg
115    rep['idea']['nearby_settings']=[{'cfg':c,'mean':r['mean']} for c,r in idea_runs]
116    rep['protocol_note']='Baseline and idea use identical tabular task, mlp_tiny architecture, 12 epochs, batch 128, and the same lr union; baseline Adam beta2 was swept.'
117    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
118    print(json.dumps(rep,indent=2))
119
120if __name__=='__main__': main()