1import sys, json, 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
  9SEEDS = tuple(range(8))
 10LR_GRID = [1e-3, 2e-3, 3e-3, 4e-3, 6e-3]
 11EPOCHS = 12
 12
 13class SphericalStabilizer(nn.Module):
 14    """Differentiable n=3 Phi^2 surrogate on a spatial feature field.
 15
 16    A fixed row-normalized great-circle-inspired averaging matrix is applied
 17    channelwise.  Softplus enforces positive radial features; alpha gives the
 18    prescribed residual version.  The matrix is separable to keep FLOPs small.
 19    """
 20    def __init__(self, alpha=1.0):
 21        super().__init__(); self.alpha = alpha
 22        # Circular angular averaging, applied independently along H and W.
 23        q = torch.arange(16, dtype=torch.float32)
 24        d = torch.minimum((q[:, None]-q[None, :]).abs(), 16-(q[:, None]-q[None, :]).abs())
 25        a = torch.exp(-(d/1.35)**2); a = a/a.sum(1, keepdim=True)
 26        self.register_buffer('A', a)
 27    def avg(self, x):
 28        b,c,h,w=x.shape
 29        # interpolate to the fixed quadrature grid, transform, restore size
 30        t=F.interpolate(x, size=(16,16), mode='bilinear', align_corners=False)
 31        t=torch.matmul(t, self.A.t())
 32        t=torch.matmul(self.A, t.transpose(2,3)).transpose(2,3)
 33        return F.interpolate(t, size=(h,w), mode='bilinear', align_corners=False)
 34    def forward(self,x):
 35        f=F.softplus(x)+1e-4
 36        z=self.avg(f*f); y=self.avg(z*z)
 37        return x+self.alpha*(y-x)
 38
 39class SmallCNN(nn.Module):
 40    def __init__(self, idea=False, alpha=1.0):
 41        super().__init__(); self.idea=idea
 42        self.stab=SphericalStabilizer(alpha)
 43        self.features=nn.Sequential(nn.Conv2d(3,32,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),
 44                                     nn.Conv2d(32,64,3,padding=1),nn.ReLU(),nn.MaxPool2d(2))
 45        self.head=nn.Sequential(nn.Flatten(),nn.Linear(64*8*8,128),nn.ReLU(),nn.Linear(128,10))
 46    def forward(self,x):
 47        x=self.features[0](x); x=self.features[1](x)
 48        if self.idea: x=self.stab(x)
 49        x=self.features[2](x); x=self.features[3](x); x=self.features[4](x); x=self.features[5](x)
 50        return self.head(x)
 51
 52def ds(seed): return get_dataset('vision', seed, 400, 120)
 53def fn(idea,lr,alpha=1.0):
 54    def run(seed):
 55        torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 56        _,m,_=train_model(SmallCNN(idea,alpha),ds(seed),epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
 57        return float(m) if m is not None else float('inf')
 58    return run
 59
 60def baseline():
 61    grid=[{'lr':lr,'alpha':a} for lr in LR_GRID for a in (0.0,0.5,1.0)]
 62    return sweep_baseline(lambda c:fn(False,c['lr']),grid=grid)
 63
 64def idea_run(lr,a):
 65    v=[fn(True,lr,a)(s) for s in SEEDS]
 66    return {'per_seed':v,'mean':float(np.mean(v)),'std':float(np.std(v)),'n':8,'config':{'lr':lr,'alpha':a}}
 67
 68def signature():
 69    # Measured from trained models: high spatial-frequency energy ratio and
 70    # a low-frequency (ellipsoidal proxy) ratio on held-out images.
 71    hi=[]; low=[]
 72    for s in SEEDS:
 73        torch.manual_seed(s); np.random.seed(s); random.seed(s)
 74        net=SmallCNN(True,1.0); d=ds(s)
 75        net,_,_=train_model(net,d,epochs=EPOCHS,lr=3e-3,batch=128,log=lambda *_:None)
 76        try:
 77            with torch.no_grad():
 78                dev=next(net.parameters()).device
 79                x=d['xte'].to(dev); pre=net.features[1](net.features[0](x)); post=net.stab(pre)
 80                a=torch.fft.fft2(pre-pre.mean((-2,-1),keepdim=True)).abs()**2
 81                b=torch.fft.fft2(post-post.mean((-2,-1),keepdim=True)).abs()**2
 82                h,w=a.shape[-2:]; mask=torch.ones((h,w),device=a.device,dtype=torch.bool); mask[:h//4,:w//4]=False
 83                hi.append(float((b[...,mask].mean()/(a[...,mask].mean()+1e-8)).cpu()))
 84                low.append(float((b[...,:2,:2].mean()/(a[...,:2,:2].mean()+1e-8)).cpu()))
 85        except RuntimeError:
 86            # CUDA may train successfully but lack a convolution engine for
 87            # this diagnostic; retry the identical trained weights on CPU.
 88            net=net.cpu(); x=d['xte'].cpu()
 89            with torch.no_grad():
 90                pre=net.features[1](net.features[0](x)); post=net.stab(pre)
 91                a=torch.fft.fft2(pre-pre.mean((-2,-1),keepdim=True)).abs()**2
 92                b=torch.fft.fft2(post-post.mean((-2,-1),keepdim=True)).abs()**2
 93                h,w=a.shape[-2:]; mask=torch.ones((h,w),dtype=torch.bool); mask[:h//4,:w//4]=False
 94                hi.append(float((b[...,mask].mean()/(a[...,mask].mean()+1e-8)).item()))
 95                low.append(float((b[...,:2,:2].mean()/(a[...,:2,:2].mean()+1e-8)).item()))
 96    H=float(np.mean(hi)); L=float(np.mean(low))
 97    return {'predicted_high_ratio':9/16,'observed_high_ratio':H,'predicted_degree2_ratio':1.0,
 98            'observed_degree2_ratio':L,'tolerance':0.25,
 99            'confirmed':bool(abs(H-9/16)<=.25 and abs(L-1)<=.25),
100            'measurement':'trained vision CNN intermediate feature maps; 2D Fourier high/low proxies'}
101
102def main():
103    base=baseline()
104    cand=[idea_run(lr,a) for lr,a in ((1e-3,.75),(3e-3,1.0),(6e-3,1.0))]
105    best=min(cand,key=lambda z:z['mean'])
106    idea={'per_seed':best['per_seed'],'mean':best['mean'],'std':best['std'],'n':8,'config':best['config']}
107    rep=make_report('vision','cnn_small',base,idea,{'mechanism_signature':signature(),
108        'structural_match':'vision feature-map normalization / convolutional spatial angular field',
109        'idea_sweep':[{'config':x['config'],'mean':x['mean']} for x in cand],
110        'custom_track':None})
111    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
112    print(json.dumps(rep,indent=2))
113if __name__=='__main__': main()