1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from scipy.optimize import linprog
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  9
 10TRACK, MODEL = 'tabular', 'mlp_tiny'
 11SEEDS = tuple(range(8))
 12SWEEP_SEEDS = tuple(range(4))
 13EPOCHS, BATCH = 12, 128
 14# The union is shared by baseline and idea. Fixed ratios are the standard
 15# alternatives to equal weighting and expose the method's replacement knob.
 16LRS = [1e-3, 3e-3, 1e-2]
 17RATIOS = [[1., 1., 1.], [10., 3., 1.], [100., 10., 1.]]
 18GRID = [{'lr': lr, 'ratio': ratio} for lr in LRS for ratio in RATIOS]
 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
 25
 26def certificate(q, S, epsilon=0.02, wmax=30.):
 27    q, S = np.asarray(q, float), np.asarray(S, float)
 28    d, tiers = len(q), S.shape[1]
 29    # variables: w_1..w_L, delta; maximize delta
 30    A, b = [], []
 31    for i in range(tiers):
 32        row = np.zeros(tiers + 1); row[i] = -1; row[-1] = 1
 33        A.append(row); b.append(0.)
 34    for k in range(d):
 35        row = np.zeros(tiers + 1); row[:tiers] = S[k]
 36        A.append(row); b.append(epsilon - q[k])
 37        row = np.zeros(tiers + 1); row[:tiers] = -S[k]
 38        A.append(row); b.append(epsilon + q[k])
 39    r = linprog(np.r_[np.zeros(tiers), -1.], A_ub=np.asarray(A), b_ub=np.asarray(b),
 40                bounds=[(0., wmax)] * tiers + [(0., wmax)], method='highs')
 41    if not r.success:
 42        return np.ones(tiers), 0., float('inf'), False
 43    w = r.x[:tiers]
 44    return w, float(r.x[-1]), float(np.max(np.abs(q + S @ w))), True
 45
 46
 47def tier_values(pred, target, scale):
 48    # Three ordered accuracy requirements: tier 1 is the strictest.
 49    err = (pred - target).abs() / scale
 50    return [torch.relu(err - t).mean() for t in (0.25, 0.50, 0.80)]
 51
 52
 53def train_idea(model, ds, lr, ratio, epsilon=0.02, refresh=4):
 54    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 55    try:
 56        net = model.to(device); x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 57        scale = y.std().clamp_min(1e-4)
 58        opt = torch.optim.Adam(net.parameters(), lr=lr)
 59        weights = torch.tensor(ratio, device=device); stats = []
 60        for ep in range(EPOCHS):
 61            net.train(); perm = torch.randperm(len(x), device=device)
 62            for start in range(0, len(x), BATCH):
 63                idx = perm[start:start+BATCH]; pred = net(x[idx])
 64                tiers = tier_values(pred, y[idx], scale)
 65                perf = ((pred-y[idx])**2).mean()
 66                if (ep * ((len(x)+BATCH-1)//BATCH) + start//BATCH) % refresh == 0:
 67                    params = [p for p in net.parameters() if p.requires_grad]
 68                    def flat_grads(loss):
 69                        gs = torch.autograd.grad(loss, params, retain_graph=True, allow_unused=True)
 70                        return torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1) for p,g in zip(params,gs)])
 71                    q = flat_grads(perf).detach().cpu().numpy()
 72                    S = np.stack([flat_grads(v).detach().cpu().numpy() for v in tiers], axis=1)
 73                    nw, delta, residual, ok = certificate(q, S, epsilon=epsilon, wmax=30.)
 74                    if ok and delta > 1e-5:
 75                        weights = torch.tensor(nw, dtype=torch.float32, device=device)
 76                    stats.append((delta, residual, float(weights.min())))
 77                loss = perf + sum(w*v for w,v in zip(weights, tiers))
 78                opt.zero_grad(); loss.backward(); opt.step()
 79        net.eval()
 80        with torch.no_grad():
 81            out = net(ds['xte'].to(device)); tv = tier_values(out, ds['yte'].to(device), scale)
 82            metric = float(((out-ds['yte'].to(device))**2).mean())
 83            violations = [float(v) for v in tv]
 84        return net, metric, {'violations': violations, 'cert_stats': stats}
 85    except RuntimeError:
 86        if torch.cuda.is_available(): torch.cuda.empty_cache()
 87        return train_idea_cpu(model.to('cpu'), ds, lr, ratio, epsilon, refresh)
 88
 89
 90def train_idea_cpu(model, ds, lr, ratio, epsilon, refresh):
 91    # Same intervention and hyperparameters, explicit CPU retry after CUDA errors.
 92    x,y=ds['xtr'],ds['ytr']; scale=y.std().clamp_min(1e-4); net=model
 93    opt=torch.optim.Adam(net.parameters(),lr=lr); weights=torch.tensor(ratio); stats=[]
 94    for ep in range(EPOCHS):
 95        perm=torch.randperm(len(x))
 96        for st in range(0,len(x),BATCH):
 97            idx=perm[st:st+BATCH]; pred=net(x[idx]); tiers=tier_values(pred,y[idx],scale); perf=((pred-y[idx])**2).mean()
 98            if (ep*((len(x)+BATCH-1)//BATCH)+st//BATCH)%refresh==0:
 99                ps=[p for p in net.parameters() if p.requires_grad]
100                def fg(z):
101                    gs=torch.autograd.grad(z,ps,retain_graph=True,allow_unused=True)
102                    return torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1) for p,g in zip(ps,gs)])
103                q=fg(perf).detach().numpy(); S=np.stack([fg(v).detach().numpy() for v in tiers],1)
104                nw,de,re,ok=certificate(q,S,epsilon,30.)
105                if ok and de>1e-5: weights=torch.tensor(nw,dtype=torch.float32)
106                stats.append((de,re,float(weights.min())))
107            opt.zero_grad(); (perf+sum(w*v for w,v in zip(weights,tiers))).backward(); opt.step()
108    with torch.no_grad():
109        out=net(ds['xte']); vals=tier_values(out,ds['yte'],scale)
110        return net,float(((out-ds['yte'])**2).mean()),{'violations':[float(v) for v in vals],'cert_stats':stats}
111
112
113def baseline_fn(cfg):
114    def run(seed):
115        seed_all(seed); ds=get_dataset(TRACK, seed, n_train=400, n_test=200)
116        # Standard weighted multi-objective loss, implemented locally so tier
117        # objectives and architecture are identical to the certificate system.
118        net=make_model(MODEL,ds['input_shape'],ds['out_dim']); device='cuda' if torch.cuda.is_available() else 'cpu'
119        try:
120            net=net.to(device); x,y=ds['xtr'].to(device),ds['ytr'].to(device); scale=y.std().clamp_min(1e-4)
121            opt=torch.optim.Adam(net.parameters(),lr=cfg['lr']); ws=torch.tensor(cfg['ratio'],device=device)
122            for _ in range(EPOCHS):
123                for st in range(0,len(x),BATCH):
124                    ix=torch.randperm(len(x),device=device)[st:st+BATCH]; p=net(x[ix]); tv=tier_values(p,y[ix],scale); loss=((p-y[ix])**2).mean()+sum(w*v for w,v in zip(ws,tv)); opt.zero_grad(); loss.backward(); opt.step()
125            with torch.no_grad(): return float(((net(ds['xte'].to(device))-ds['yte'].to(device))**2).mean())
126        except RuntimeError:
127            return float('inf')
128    return run
129
130
131def idea_fn(cfg):
132    def run(seed):
133        seed_all(seed); ds=get_dataset(TRACK,seed,n_train=400,n_test=200)
134        _,m,_=train_idea(make_model(MODEL,ds['input_shape'],ds['out_dim']),ds,cfg['lr'],cfg['ratio'])
135        return m
136    return run
137
138
139def signature(cfg):
140    seed=0; seed_all(seed); ds=get_dataset(TRACK,seed,n_train=400,n_test=200)
141    net,m,info=train_idea(make_model(MODEL,ds['input_shape'],ds['out_dim']),ds,cfg['lr'],cfg['ratio'])
142    vals=np.asarray(info['cert_stats'],float); finite=vals[np.isfinite(vals[:,1])] if len(vals) else np.empty((0,3))
143    observed=float(np.mean(finite[:,1])) if len(finite) else float('inf')
144    predicted=float(0.02); return {'quantity':'parameter-gradient infinity stationarity residual at refreshes','predicted_epsilon':predicted,'observed_mean_residual':observed,'positive_margin_frequency':float(np.mean(vals[:,0]>1e-5)) if len(vals) else 0.,'confirmed':bool(observed <= predicted*1.5)}
145
146
147def main():
148    # Cheap analytical LP sanity check before NN training.
149    w,d,r,ok=certificate(np.array([-3.,-1.]),np.eye(2),epsilon=1e-9,wmax=100.)
150    math_check={'predicted_weights':[3.,1.],'observed_weights':w.tolist(),'residual':r,'passed':bool(np.max(np.abs(w-[3.,1.]))<1e-5)}
151    base=sweep_baseline(baseline_fn,GRID,seeds=SWEEP_SEEDS)
152    best_cfg=base['best_cfg']; idea_trials=[]
153    for cfg in GRID:
154        rr=evaluate(idea_fn(cfg),seeds=SEEDS); idea_trials.append({'cfg':cfg,'result':rr})
155    best=min(idea_trials,key=lambda z:z['result']['mean'])
156    extra={'mechanism_signature':signature(best['cfg']),'math_check':math_check,'idea_sweep':idea_trials,'protocol_notes':'Matched tabular Friedman regression and mlp_tiny; baseline and controller share all lr and ratio configurations, Adam, epochs, batch, and data seeds.'}
157    report=make_report(TRACK,MODEL,base,best['result'],extra)
158    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
159
160if __name__=='__main__': main()