1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report
  9
 10TRACK = "tabular"
 11MODEL = "mlp_tiny"
 12EPOCHS = 18
 13BATCH = 128
 14WEIGHT_DECAY = 0.0
 15# Shared union: baseline and idea both run every learning rate below.
 16LRS = [1e-3, 3e-3, 1e-2]
 17IDEA_SETTINGS = [
 18    {"lr": 1e-3, "r0": 0.02},
 19    {"lr": 3e-3, "r0": 0.02},
 20    {"lr": 1e-2, "r0": 0.02},
 21]
 22SEEDS = tuple(range(8))
 23SWEEP_SEEDS = tuple(range(4))
 24
 25
 26def seed_all(seed):
 27    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 28    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 29
 30
 31def global_norm(xs):
 32    return math.sqrt(sum(float(x.detach().pow(2).sum().cpu()) for x in xs))
 33
 34
 35def train_baseline(cfg, seed, return_model=False):
 36    seed_all(seed)
 37    ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=400)
 38    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
 39    device = "cuda" if torch.cuda.is_available() else "cpu"
 40    try:
 41        net = net.to(device)
 42        opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"], weight_decay=WEIGHT_DECAY)
 43        lossf = nn.MSELoss()
 44        x, y = ds["xtr"].to(device), ds["ytr"].to(device)
 45        for _ in range(EPOCHS):
 46            net.train(); perm = torch.randperm(len(x), device=device)
 47            for i in range(0, len(x), BATCH):
 48                ix = perm[i:i+BATCH]; loss = lossf(net(x[ix]), y[ix])
 49                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 50        net.eval()
 51        with torch.no_grad(): metric = float(lossf(net(ds["xte"].to(device)), ds["yte"].to(device)).cpu())
 52        return (metric, net, ds, device) if return_model else metric
 53    except (RuntimeError, torch.cuda.CudaError):
 54        if device == "cuda":
 55            seed_all(seed); net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
 56            return train_baseline(cfg, seed, return_model) if False else _train_cpu(net, ds, cfg, return_model)
 57        raise
 58
 59
 60def _train_cpu(net, ds, cfg, return_model=False):
 61    net = net.cpu(); opt = torch.optim.Adam(net.parameters(), lr=cfg["lr"], weight_decay=WEIGHT_DECAY)
 62    lossf = nn.MSELoss(); x, y = ds["xtr"], ds["ytr"]
 63    for _ in range(EPOCHS):
 64        perm = torch.randperm(len(x))
 65        for i in range(0, len(x), BATCH):
 66            ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]), y[ix]); opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 67    with torch.no_grad(): metric=float(lossf(net(ds["xte"]), ds["yte"]))
 68    return (metric, net, ds, "cpu") if return_model else metric
 69
 70
 71def train_idea(cfg, seed, return_model=False):
 72    seed_all(seed)
 73    ds = get_dataset(TRACK, seed=seed, n_train=400, n_test=400)
 74    net = make_model(MODEL, ds["input_shape"], ds["out_dim"])
 75    device = "cuda" if torch.cuda.is_available() else "cpu"
 76    try:
 77        net=net.to(device); opt=torch.optim.Adam(net.parameters(), lr=cfg["lr"], weight_decay=WEIGHT_DECAY)
 78        lossf=nn.MSELoss(); x,y=ds["xtr"].to(device),ds["ytr"].to(device)
 79        # Fixed cap is deliberately identical to initial adaptive radius.
 80        r=float(cfg["r0"]); rmin=.1*cfg["r0"]; rmax=10*cfg["r0"]
 81        qs=[]; radii=[]; scales=[]; proposal_norms=[]
 82        for _ in range(EPOCHS):
 83            net.train(); perm=torch.randperm(len(x),device=device)
 84            for i in range(0,len(x),BATCH):
 85                ix=perm[i:i+BATCH]; loss=lossf(net(x[ix]),y[ix]); opt.zero_grad(set_to_none=True); loss.backward()
 86                # Adam's actual proposal is reconstructed from its state after a harmless step.
 87                # We instead use the standard Adam preconditioned proposal explicitly.
 88                props=[]
 89                beta1,beta2=opt.param_groups[0]["betas"]; eps=opt.param_groups[0]["eps"]; step=int(opt.state[next(net.parameters())].get("step",0))+1
 90                for p in net.parameters():
 91                    if p.grad is None: continue
 92                    st=opt.state[p]
 93                    if not st: st["step"]=torch.tensor(0.,device=p.device); st["exp_avg"]=torch.zeros_like(p); st["exp_avg_sq"]=torch.zeros_like(p)
 94                    m=st["exp_avg"]; v=st["exp_avg_sq"]; m.mul_(beta1).add_(p.grad,alpha=1-beta1); v.mul_(beta2).addcmul_(p.grad,p.grad,value=1-beta2)
 95                    bc1=1-beta1**step; bc2=1-beta2**step
 96                    props.append(-cfg["lr"]*(m/bc1)/(v.sqrt()/math.sqrt(bc2)+eps))
 97                pn=global_norm(props); q=min(1.,pn/(r+1e-12)); scale=min(1.,r/(pn+1e-12))
 98                with torch.no_grad():
 99                    for p,prop in zip([p for p in net.parameters() if p.grad is not None],props): p.add_(prop,alpha=scale)
100                r=float(np.clip(r*math.exp(.02*q-.01*(1-q)),rmin,rmax)); qs.append(q); radii.append(r); scales.append(scale); proposal_norms.append(pn)
101        net.eval();
102        with torch.no_grad(): metric=float(lossf(net(ds["xte"].to(device)),ds["yte"].to(device)).cpu())
103        if return_model: return metric,net,ds,device,{"mean_q":float(np.mean(qs)),"clip_fraction":float(np.mean(np.array(qs)>=1-1e-9)),"radius_final":r,"radius_initial":cfg["r0"],"mean_scale":float(np.mean(scales)),"mean_proposal_norm":float(np.mean(proposal_norms))}
104        return metric
105    except (RuntimeError, torch.cuda.CudaError):
106        raise RuntimeError("CUDA failure; rerun with CUDA_VISIBLE_DEVICES='' for CPU fallback")
107
108
109def baseline_factory(cfg): return lambda seed: train_baseline(cfg, seed)
110def idea_factory(cfg): return lambda seed: train_idea(cfg, seed)
111
112
113def main():
114    # Official baseline sweep, then full-seed baseline at selected cfg.
115    base=sweep_baseline(baseline_factory, [{"lr":lr} for lr in LRS], seeds=SWEEP_SEEDS)
116    idea_cfgs=IDEA_SETTINGS
117    idea_sweep=[]
118    for cfg in idea_cfgs: idea_sweep.append({"cfg":cfg,"mean":evaluate(idea_factory(cfg),SWEEP_SEEDS)["mean"]})
119    best_idea_cfg=min(idea_sweep,key=lambda z:z["mean"])["cfg"]
120    base_full=base["full"]
121    idea_full=evaluate(idea_factory(best_idea_cfg),SEEDS)
122    # Trained-model signature: inspect one paired seed's actual proposals/updates.
123    _,_,_,_,sig=train_idea(best_idea_cfg,0,True)
124    sig.update({"prediction":"sustained saturation expands radius and later unsaturation contracts it","observed_radius_increase":sig["radius_final"]>sig["radius_initial"],"observed_clipping_fraction":sig["clip_fraction"],"confirmed":bool(sig["radius_final"]>sig["radius_initial"] and sig["clip_fraction"]>0)})
125    report=make_report(TRACK,MODEL,{"best_cfg":base["best_cfg"],"sweep":base["sweep"],"full":base_full},idea_full,{"idea_sweep":idea_sweep,"mechanism_signature":sig})
126    Path("bench_report.json").write_text(json.dumps(report,indent=2))
127    print(json.dumps(report,indent=2))
128
129if __name__=="__main__": main()