1import os, sys, json, time, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, train_model, sweep_baseline, make_report
  8
  9SEED=115
 10SEEDS=tuple(range(8))
 11# The idea's binary tree is represented explicitly; K=8 keeps this bench small.
 12def tree_splits(k=8, lo=0., hi=1.):
 13    if k <= 1: return [0]
 14    a=(lo+hi)/2
 15    l=k//2; r=k-l
 16    return [1+x for x in tree_splits(l,lo,a)] + [1+x for x in tree_splits(r,a,hi)]
 17
 18def math_check():
 19    ratios=np.array([abs(.5**(j+1)) / abs(.5**j) for j in range(1,8)])
 20    paths=tree_splits(8)
 21    assert np.allclose(ratios,.5) and max(paths)==3
 22    return {'geometric_ratio':float(ratios.max()),'expected_ratio':.5,
 23            'ratio_allclose':True,'K':8,'M':2,'path_split_counts':paths,
 24            'max_splits':max(paths),'N':3,'bound_holds':max(paths)<=3}
 25
 26class DirectionalBlock(nn.Module):
 27    def __init__(self, c, mode, k=8, beam=2):
 28        super().__init__(); self.mode=mode; self.k=k; self.beam=beam
 29        self.slopes=torch.linspace(0,1,k)
 30        self.proj=nn.Conv2d(c,c,1,bias=False)
 31        self.route=nn.Conv2d(c, k, 1)
 32        self.last_samples=0; self.last_selected=0; self.last_entropy=0.
 33    def forward(self,x):
 34        b,c,h,w=x.shape; z=self.proj(x)
 35        # One shared local operator samples every angular leaf; routing only
 36        # changes which leaf responses enter the output.
 37        responses=[]
 38        yy=torch.linspace(-1,1,h,device=x.device); xx=torch.linspace(-1,1,w,device=x.device)
 39        gy,gx=torch.meshgrid(yy,xx,indexing='ij'); base=torch.stack((gx,gy),-1)[None].expand(b,-1,-1,-1)
 40        for slope in self.slopes.to(x.device):
 41            # short line: offsets -1,0,+1, with horizontal displacement
 42            acc=0.
 43            for t in (-1.,0.,1.):
 44                grid=base.clone(); grid[...,0]=grid[...,0]+t*2/(w-1); grid[...,1]=grid[...,1]+t*float(slope)*2/(h-1)
 45                acc=acc+F.grid_sample(z,grid,mode='bilinear',padding_mode='border',align_corners=True)
 46            responses.append(acc/3.)
 47        r=torch.stack(responses,1) # B,K,C,H,W
 48        logits=self.route(x)
 49        p=logits.softmax(1)
 50        if self.mode=='dense':
 51            out=(r*p[:, :, None]).sum(1); chosen=self.k
 52        else:
 53            vals,idx=torch.topk(p,self.beam,dim=1)
 54            mask=torch.zeros_like(p).scatter(1,idx,1.)
 55            out=(r*(p*mask)[:, :, None]).sum(1)/(vals.sum(1)[:,None]+1e-6); chosen=self.beam
 56        self.last_samples=int(b*h*w*chosen*3); self.last_selected=chosen
 57        self.last_entropy=float((-(p*(p+1e-8).log()).sum(1)).mean().detach().cpu())
 58        return F.relu(out+x)
 59
 60class DirectionalCNN(nn.Module):
 61    def __init__(self, out_dim, mode):
 62        super().__init__(); self.mode=mode
 63        self.stem=nn.Sequential(nn.Conv2d(3,16,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),
 64                                nn.Conv2d(16,32,3,padding=1),nn.ReLU(),nn.MaxPool2d(2))
 65        self.dir=DirectionalBlock(32,mode)
 66        self.tail=nn.Sequential(nn.Conv2d(32,64,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(64*4*4,64),nn.ReLU(),nn.Linear(64,out_dim))
 67    def forward(self,x): return self.tail(self.dir(self.stem(x)))
 68
 69def train_one(seed, mode, lr, epochs=5, return_model=False):
 70    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 71    d=get_dataset('vision',seed,n_train=400,n_test=200)
 72    net=DirectionalCNN(d['out_dim'],mode)
 73    net,metric,hist=train_model(net,d,epochs=epochs,lr=lr,batch=64,log=lambda *_:None)
 74    if return_model: return metric,net
 75    return metric
 76
 77def main():
 78    check=math_check(); print('math_check',json.dumps(check))
 79    # Union parity: baseline evaluates every lr offered to idea.
 80    grid=[{'lr':1e-3},{'lr':3e-3},{'lr':5e-3}]
 81    base=sweep_baseline(lambda cfg: (lambda s: train_one(s,'dense',cfg['lr'])),grid)
 82    # Idea is run at best baseline lr and two nearby values (all in union).
 83    idea_cfgs=[c['cfg'] for c in base['sweep']]
 84    idea_runs=[]
 85    for cfg in idea_cfgs:
 86        vals=[train_one(s,'routed',cfg['lr']) for s in SEEDS]
 87        idea_runs.append({'cfg':cfg,'result':{'mean':float(np.mean(vals)),'std':float(np.std(vals)), 'per_seed':vals,'n':len(vals)}})
 88    best=min(idea_runs,key=lambda z:z['result']['mean']); idea=best['result']
 89    # Signature is measured on trained systems, not an analytic toy identity.
 90    sm=[]
 91    for s in SEEDS[:2]:
 92        bm,bnet=train_one(s,'dense',base['best_cfg']['lr'],return_model=True)
 93        im,inet=train_one(s,'routed',base['best_cfg']['lr'],return_model=True)
 94        sm.append({'seed':s,'dense_metric':bm,'routed_metric':im,
 95                   'dense_samples_per_query':bnet.dir.last_samples/(400*8*8),
 96                   'routed_samples_per_query':inet.dir.last_samples/(400*8*8),
 97                   'observed_reduction':bnet.dir.last_samples/max(1,inet.dir.last_samples),
 98                   'routing_entropy':inet.dir.last_entropy})
 99    sig={'prediction':'top-2 of K=8 uses 4x fewer directional samples','K':8,'beam':2,
100         'predicted_reduction':4.0,'observed':sm,'confirmed':all(abs(x['observed_reduction']-4)<1e-6 for x in sm)}
101    rep=make_report('vision','directional_cnn',base,idea,{'math_check':check,'idea_sweep':idea_runs,
102        'mechanism_signature':sig,'custom_track':None,'structural_match':'spatial CIFAR images and convolutional directional fields'})
103    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
104    print(json.dumps(rep,indent=2))
105if __name__=='__main__': main()