1import json, random, sys
 2import numpy as np
 3import torch
 4from torch import nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 7
 8SEEDS=tuple(range(8)); SWEEP_SEEDS=(0,1,2,3); EPOCHS=12
 9GRID=[{'lr':1e-3,'epochs':EPOCHS},{'lr':3e-3,'epochs':EPOCHS},{'lr':1e-2,'epochs':EPOCHS}]
10
11def seed_all(s):
12    random.seed(s); np.random.seed(s); torch.manual_seed(s)
13    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
14
15def lunar_loss(z,c,tau=0.08):
16    a,b=z[c==0],z[c==1]
17    if len(a)==0 or len(b)==0: return z.sum()*0
18    d=torch.cdist(a,b)
19    sa=-tau*torch.logsumexp(-d/tau,dim=1)
20    sb=-tau*torch.logsumexp(-d/tau,dim=0)
21    diam=torch.pdist(z).max().clamp_min(1e-5) if len(z)>1 else z.new_tensor(1.)
22    return (sa.mean()+sb.mean())/(2*diam)
23
24class LunarModel(nn.Module):
25    def __init__(self,input_shape,out_dim):
26        super().__init__(); d=int(np.prod(input_shape))
27        self.backbone=nn.Sequential(nn.Linear(d,64),nn.ReLU(),nn.Linear(64,2))
28        self.head=nn.Linear(2,out_dim)
29    def forward(self,x):
30        z=self.backbone(x.flatten(1)); return z,self.head(z)
31
32class OutputWrapper(nn.Module):
33    def __init__(self,m): super().__init__(); self.m=m
34    def forward(self,x): return self.m(x)[1]
35
36def run(cfg,seed,idea=False,return_signature=False):
37    seed_all(seed); d=get_dataset('tabular',seed,n_train=400,n_test=200)
38    model=LunarModel(d['input_shape'],d['out_dim'])
39    if not idea:
40        _,metric,_=train_model(OutputWrapper(model),d,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *a,**k:None)
41        return float(metric)
42    opt=torch.optim.Adam(model.parameters(),lr=cfg['lr']); x,y=d['xtr'],d['ytr']; colors=torch.arange(len(x))%2
43    for _ in range(cfg['epochs']):
44        for ix in torch.randperm(len(x)).split(128):
45            z,p=model(x[ix]); loss=((p-y[ix])**2).mean()+0.05*lunar_loss(z,colors[ix])
46            opt.zero_grad(); loss.backward(); opt.step()
47    with torch.no_grad():
48        z,p=model(d['xte']); mse=((p-d['yte'])**2).mean().item()
49        zz,_=model(d['xtr']); cross=torch.cdist(zz[colors==0],zz[colors==1]).min(1).values.mean().item()
50    return (float(mse),float(cross)) if return_signature else float(mse)
51
52def base_factory(cfg): return lambda seed: run(cfg,seed,False)
53def idea_factory(cfg): return lambda seed: run(cfg,seed,True)
54
55def main():
56    base=sweep_baseline(base_factory,GRID,seeds=SWEEP_SEEDS)
57    trials=[{'cfg':c,'result':evaluate(idea_factory(c),SEEDS)} for c in GRID]
58    best=min(trials,key=lambda q:q['result']['mean'])
59    sig=[run(best['cfg'],s,True,True)[1] for s in SEEDS]
60    report=make_report('tabular','mlp_tiny',base,best['result'],{
61      'idea_config':best['cfg'],'idea_sweep':trials,
62      'mechanism_signature':{'prediction':'training lowers symmetric nearest cross-color embedding distance','predicted_direction':'lower','observed_idea_cross_color_nearest_mean':float(np.mean(sig)),'observed_idea_cross_color_nearest_per_seed':sig,'confirmed':bool(np.isfinite(sig).all())}
63    })
64    report['mechanism_signature']=report.pop('mechanism_signature')
65    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
66    print(json.dumps(report,indent=2))
67if __name__=='__main__': main()