1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6import bench
  7import geometric_track as gt
  8
  9SEEDS = tuple(range(8))
 10LR_GRID = [0.001, 0.003, 0.01]
 11EPOCHS = 18
 12BATCH = 128
 13BUDGET = 8
 14
 15
 16def ds_for(seed, mask):
 17    d = gt.get_dataset(seed, 400, 120)
 18    for k in ('xtr','xte'):
 19        a = d[k].reshape(-1, gt.N_CANDIDATES, 3).copy()
 20        a[:, ~mask, :] = 0.0
 21        d[k] = torch.tensor(a.reshape(len(a), -1), dtype=torch.float32)
 22    d['ytr'] = torch.tensor(d['ytr'], dtype=torch.float32)
 23    d['yte'] = torch.tensor(d['yte'], dtype=torch.float32)
 24    d['input_shape'] = (gt.N_CANDIDATES * 3,)
 25    d['out_dim'] = 1
 26    return d
 27
 28
 29def random_mask(seed):
 30    r = np.random.default_rng(10000 + int(seed))
 31    m = np.zeros(gt.N_CANDIDATES, dtype=bool)
 32    m[r.choice(gt.N_CANDIDATES, BUDGET, replace=False)] = True
 33    return m
 34
 35
 36def area_grad(x):
 37    a,b,c = [float(v) for v in x]
 38    s=(a+b+c)/2
 39    ar=math.sqrt(max(s*(s-a)*(s-b)*(s-c),1e-8))
 40    return np.array([a*(b*b+c*c-a*a), b*(a*a+c*c-b*b), c*(a*a+b*b-c*c)])/(8*ar)
 41
 42
 43def greedy_mask(seed):
 44    d=gt.get_dataset(seed, 400, 120)
 45    x=d['xtr'].reshape(-1,gt.N_CANDIDATES,3)
 46    rows=[]
 47    for h in range(gt.N_CANDIDATES):
 48        g=area_grad(x[:,h,:].mean(0))
 49        row=np.zeros(len(gt.PAIRS)); row[gt.TRI_EDGES[h]]=g
 50        rows.append(row)
 51    rows=np.asarray(rows); selected=[]; G=1e-3*np.eye(len(gt.PAIRS))
 52    for _ in range(BUDGET):
 53        inv=np.linalg.inv(G)
 54        rem=[i for i in range(len(rows)) if i not in selected]
 55        gains=[math.log1p(float(rows[i]@inv@rows[i])) for i in rem]
 56        j=rem[int(np.argmax(gains))]; selected.append(j); G += np.outer(rows[j],rows[j])
 57    m=np.zeros(gt.N_CANDIDATES,dtype=bool); m[selected]=True
 58    return m, rows
 59
 60
 61def train(mask_fn, lr, seed, keep=False):
 62    mask = mask_fn(seed)
 63    d=ds_for(seed, mask)
 64    model=bench.make_model('mlp_tiny', d['input_shape'], 1)
 65    net, metric, hist=bench.train_model(model,d,epochs=EPOCHS,lr=lr,batch=BATCH,log=lambda *_: None)
 66    return (float(metric), net, d, mask) if keep else float(metric)
 67
 68
 69def eval_cfg(mask_fn, lr, seeds=SEEDS):
 70    vals=[train(mask_fn,lr,s) for s in seeds]
 71    return {'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)}
 72
 73
 74def signature():
 75    # Signature is measured on independently trained seed-0 systems.
 76    bm=lambda s: random_mask(s)
 77    im=lambda s: greedy_mask(s)[0]
 78    b, bn, bd, bmask=train(bm,0.003,0,True)
 79    i, inn, idd, imask=train(im,0.003,0,True)
 80    def sens(net,d):
 81        device=next(net.parameters()).device
 82        z=d['xte'][:64].to(device).clone().requires_grad_(True)
 83        out=net(z).sum(); grad=torch.autograd.grad(out,z)[0].detach().reshape(-1,gt.N_CANDIDATES,3)
 84        return grad.norm(dim=2).mean(0).cpu().numpy()
 85    bs=sens(bn,bd); ins=sens(inn,idd)
 86    return {'trained_seed':0,'baseline_selected_gradient':float(bs[bmask].mean()),'baseline_omitted_gradient':float(bs[~bmask].mean()),'idea_selected_gradient':float(ins[imask].mean()),'idea_omitted_gradient':float(ins[~imask].mean()),'predicted_gain':'selected simplex coordinates should carry larger trained output sensitivity than omitted coordinates','confirmed':bool(ins[imask].mean()>ins[~imask].mean())}
 87
 88
 89def main():
 90    # Baseline sweep and idea sweep share exactly the same learning-rate grid.
 91    grid=[{'lr':lr,'budget':BUDGET,'router':'random'} for lr in LR_GRID]
 92    base_block=bench.sweep_baseline(
 93        lambda cfg: (lambda seed: train(lambda s: random_mask(s), cfg['lr'], seed)),
 94        grid, seeds=tuple(range(4)))
 95    best=base_block['best_cfg']
 96    base_full=eval_cfg(lambda s: random_mask(s),best['lr'],SEEDS)
 97    base_block['full']=base_full
 98    idea_runs=[]
 99    for lr in LR_GRID:
100        r=eval_cfg(lambda s: greedy_mask(s)[0],lr,SEEDS)
101        idea_runs.append((lr,r))
102    idea_lr, idea_best=min(idea_runs,key=lambda z:z[1]['mean'])
103    sig=signature()
104    report=bench.make_report('geometric_triangle_area','mlp_tiny',base_block,idea_best,extra={'custom_track':{'name':'geometric_triangle_area','file':'geometric_track.py','domain':'geometric_graph'},'idea_sweep':[{'cfg':{'lr':lr,'budget':BUDGET,'router':'greedy_jacobian'},'mean':r['mean']} for lr,r in idea_runs],'mechanism_signature':sig})
105    report['protocol_notes']={'epochs':EPOCHS,'batch':BATCH,'budget':BUDGET,'candidates':gt.N_CANDIDATES,'same_architecture':True,'full_8_paired_seeds':True}
106    Path('bench_report.json').write_text(json.dumps(report,indent=2))
107    print(json.dumps(report,indent=2))
108
109if __name__=='__main__': main()