1import os, sys, json, time
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, evaluate, sweep_baseline, make_report
  8from bench.protocol import DEFAULT_SEEDS, SWEEP_SEEDS
  9
 10E = 16
 11
 12def seed_all(seed):
 13    np.random.seed(seed); torch.manual_seed(seed)
 14    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 15
 16def active_set_simplex(Z, b, lam=1e-3, tol=1e-9, max_pivots=100):
 17    """Exact-ish active-set solution of min .5*x'Ax-b'x, x>=0, 1'x=1."""
 18    Z = np.asarray(Z, dtype=np.float64); b = np.asarray(b, dtype=np.float64)
 19    e = Z.shape[1]; A = Z.T @ Z + lam*np.eye(e); F = list(range(e)); pivots = 0
 20    while pivots <= max_pivots:
 21        ii = np.asarray(F, dtype=int); Af = A[np.ix_(ii, ii)]
 22        K = np.zeros((len(ii)+1, len(ii)+1)); K[:-1,:-1] = Af
 23        K[:-1,-1] = 1.; K[-1,:-1] = 1.
 24        rhs = np.r_[b[ii], 1.]
 25        sol = np.linalg.solve(K, rhs); xf, nu = sol[:-1], sol[-1]
 26        x = np.zeros(e); x[ii] = xf
 27        bad = [(v, i) for v,i in zip(xf,ii) if v < -tol]
 28        if bad:
 29            m = min(v for v,_ in bad); i = min(i for v,i in bad if v <= m+1e-12)
 30            F.remove(i); pivots += 1; continue
 31        g = A @ x - b + nu
 32        blocked = [i for i in range(e) if i not in F and g[i] < -tol]
 33        if blocked:
 34            m = min(g[i] for i in blocked); i = min(i for i in blocked if g[i] <= m+1e-12)
 35            F.append(i); F.sort(); pivots += 1; continue
 36        x[np.abs(x) < tol] = 0.
 37        return x, {'pivots': pivots, 'support': int((x > tol).sum()),
 38                   'simplex_error': float(abs(x.sum()-1)), 'min_x': float(x.min()),
 39                   'kkt': float(max(np.max(np.abs(g[x>tol])) if np.any(x>tol) else 0.,
 40                                    max(0., -np.min(g[x<=tol])) if np.any(x<=tol) else 0.))}
 41    raise RuntimeError('active-set pivot limit')
 42
 43class RouterMLP(nn.Module):
 44    def __init__(self, d):
 45        super().__init__()
 46        self.trunk = nn.Sequential(nn.Linear(d,64), nn.ReLU(), nn.Linear(64,64), nn.ReLU())
 47        self.experts = nn.Linear(64,E)
 48        self.logits = nn.Parameter(torch.zeros(E))
 49    def expert_values(self,x): return self.experts(self.trunk(x))
 50    def forward(self,x,weights=None):
 51        z=self.expert_values(x)
 52        if weights is None: weights=torch.softmax(self.logits,dim=0)
 53        return z @ weights, z
 54
 55def fit(seed, cfg, idea=False, collect=False):
 56    seed_all(seed)
 57    d=get_dataset('tabular', seed, n_train=400, n_test=400)
 58    xtr,ytr=d['xtr'],d['ytr']; xte,yte=d['xte'],d['yte']
 59    dev='cuda' if torch.cuda.is_available() else 'cpu'
 60    try:
 61        net=RouterMLP(int(np.prod(d['input_shape']))).to(dev)
 62        xtr=xtr.to(dev).float().reshape(len(xtr),-1); ytr=ytr.to(dev).float().reshape(-1)
 63        xte=xte.to(dev).float().reshape(len(xte),-1); yte=yte.to(dev).float().reshape(-1)
 64        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'], weight_decay=cfg.get('weight_decay', 0.0))
 65        weights=torch.full((E,),1./E,device=dev)
 66        epochs=cfg['epochs']; bs=cfg['batch']; stats=[]
 67        for ep in range(epochs):
 68            order=torch.randperm(len(xtr),device=dev)
 69            for start in range(0,len(xtr),bs):
 70                ix=order[start:start+bs]; pred,z=net(xtr[ix], None if not idea else weights)
 71                loss=((pred-ytr[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 72            if idea:
 73                with torch.no_grad():
 74                    z=net.expert_values(xtr).detach().cpu().numpy(); yy=ytr.detach().cpu().numpy()
 75                # Ridge least-squares router calibration, then simplex projection via active set.
 76                weights_np, met=active_set_simplex(z, z.T@yy, lam=cfg['lam'])
 77                weights=torch.tensor(weights_np,dtype=torch.float32,device=dev)
 78                stats.append(met)
 79        with torch.no_grad():
 80            pred,_=net(xte, weights if idea else None)
 81            metric=float(((pred-yte)**2).mean().cpu())
 82        if collect:
 83            last=stats[-1] if stats else {'pivots':0,'support':E,'simplex_error':0.,'min_x':0.,'kkt':0.}
 84            return metric,last
 85        return metric
 86    except RuntimeError:
 87        # Robust CPU fallback for constrained shared GPU slots.
 88        if dev=='cuda':
 89            torch.cuda.empty_cache()
 90            torch.set_default_device('cpu')
 91            return fit(seed,cfg,idea,collect)
 92        raise
 93
 94def main():
 95    # Union parity: all lrs occur in both grids; baseline also sweeps its central
 96    # router knob (Adam optimizer) through weight decay values.
 97    grid=[{'lr':lr,'weight_decay':wd,'epochs':12,'batch':64,'lam':lam}
 98          for lr in (1e-3,3e-3,1e-2) for wd in (0.,1e-4) for lam in (1e-3,)]
 99    def base_fn(c): return lambda s: fit(s,c,False)
100    base=sweep_baseline(base_fn,grid,seeds=SWEEP_SEEDS)
101    best=base['best_cfg']; idea_grid=[dict(best,lam=lam) for lam in (3e-4,1e-3,3e-3)]
102    # Evaluate every idea setting on all paired seeds; report the best by sweep-seed mean.
103    tried=[]
104    for c in idea_grid:
105        r=evaluate(lambda s,c=c: fit(s,c,True), seeds=SWEEP_SEEDS)
106        tried.append((r['mean'],c))
107    idea_cfg=min(tried,key=lambda x:x[0])[1]
108    idea=evaluate(lambda s: fit(s,idea_cfg,True),seeds=DEFAULT_SEEDS)
109    base['idea_grid']= [{'cfg':c,'mean':m} for m,c in tried]
110    sigs=[fit(s,idea_cfg,True,True)[1] for s in DEFAULT_SEEDS]
111    sig={'track_match':'tabular optimizer/training-procedure structure',
112         'prediction':'simplex router remains feasible and active-set identifies sparse support',
113         'predicted':{'simplex_error':0.0,'min_x':0.0,'sparse_support':True},
114         'observed':{'mean_simplex_error':float(np.mean([q['simplex_error'] for q in sigs])),
115                     'min_x':float(min(q['min_x'] for q in sigs)),
116                     'mean_support':float(np.mean([q['support'] for q in sigs])),
117                     'mean_pivots':float(np.mean([q['pivots'] for q in sigs]))},
118         'confirmed':bool(max(q['simplex_error'] for q in sigs)<1e-6 and min(q['min_x'] for q in sigs)>=-1e-7)}
119    rep=make_report('tabular','mlp_tiny',base,idea,sig)
120    rep['idea_cfg']=idea_cfg
121    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
122    print(json.dumps(rep,indent=2))
123
124if __name__=='__main__': main()