1import os, sys, json, time
  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 train_model, evaluate, sweep_baseline, make_report
  8
  9# Local custom track: convolution is the relevant mathematical structure.
 10META = {'name':'long_1d_convolution', 'domain':'cnn',
 11        'description':'Synthetic long 1D signal classification requiring learned local convolution.'}
 12
 13def get_dataset(seed, n_train=400, n_test=200):
 14    rng=np.random.RandomState(seed)
 15    def make(n):
 16        x=rng.normal(0,0.35,(n,1,64)).astype('float32')
 17        y=rng.randint(0,2,n).astype('int64')
 18        # Class-dependent local motifs at random positions, with nuisance low frequency signal.
 19        for i in range(n):
 20            pos=rng.randint(4,59)
 21            motif=np.array([1.,1.,-1.,-1.,1.,1.],dtype='float32') if y[i] else np.array([-1.,1.,1.,-1.,-1.,1.],dtype='float32')
 22            x[i,0,pos:pos+6] += 0.9*motif
 23            x[i,0] += (0.15*np.sin(np.arange(64)*rng.uniform(.08,.22)+rng.uniform(0,6.28))).astype('float32')
 24        return x,y
 25    xtr,ytr=make(n_train); xte,yte=make(n_test)
 26    return {'xtr':torch.tensor(xtr), 'ytr':torch.tensor(ytr),
 27            'xte':torch.tensor(xte), 'yte':torch.tensor(yte),
 28            'task':'classification','metric':'err','input_shape':(1,64),'out_dim':2,
 29            'track':'long_1d_convolution','meta':META,'custom':True}
 30
 31def implicit_dft(x, M, m):
 32    # Positive-sign DFT, with zero injection only in bounded tiles (no length-M input).
 33    L=x.shape[-1]; p=(L+m-1)//m; q=M//m
 34    tile=F.pad(x, (0,p*m-L)).reshape(*x.shape[:-1],p,m).to(torch.float32).to(torch.complex64)
 35    r=torch.arange(q,device=x.device)[:,None]; t=torch.arange(p,device=x.device)[None,:]
 36    s=torch.arange(m,device=x.device)[None,:]
 37    zq=torch.exp(2j*torch.pi*(r*t)/q); zqm=torch.exp(2j*torch.pi*(r*s)/(q*m))
 38    U=torch.einsum('rt,...ts->...rs',zq,tile)*zqm
 39    ell=torch.arange(m,device=x.device)[:,None]
 40    zm=torch.exp(2j*torch.pi*(ell*s)/m)
 41    out=torch.einsum('ls,...rs->...rl',zm,U)
 42    return out.transpose(-2,-1).reshape(*x.shape[:-1],M)
 43
 44def fft_conv(x, w, implicit=False, m=8):
 45    # Cross-correlation, matching Conv1d(no padding): convolution with reversed kernel.
 46    L=x.shape[-1]; K=w.shape[-1]; M=1
 47    while M < L+K-1: M*=2
 48    wr=w.flip(-1)
 49    if implicit:
 50        X=implicit_dft(x,M,m); W=implicit_dft(wr,M,m)
 51    else:
 52        X=torch.fft.fft(F.pad(x,(0,M-L)).to(torch.complex64))
 53        W=torch.fft.fft(F.pad(wr,(0,M-K)).to(torch.complex64))
 54    full=torch.fft.ifft(X.unsqueeze(1)*W.unsqueeze(0), n=M).real
 55    return full[...,K-1:K-1+L-K+1]
 56
 57class FFTConv1d(nn.Module):
 58    def __init__(self, cin, cout, k, implicit=False, m=8):
 59        super().__init__(); self.weight=nn.Parameter(torch.empty(cout,cin,k)); self.bias=nn.Parameter(torch.zeros(cout))
 60        nn.init.kaiming_uniform_(self.weight,a=np.sqrt(5)); self.implicit=implicit; self.m=m
 61    def forward(self,x):
 62        # x B,C,L; einsum gives B,O,Lout and sums input channels.
 63        ys=[]
 64        for o in range(self.weight.shape[0]):
 65            yo=0
 66            for c in range(self.weight.shape[1]): yo=yo+fft_conv(x[:,c,:],self.weight[o:o+1,c,:],self.implicit,self.m)
 67            ys.append(yo)
 68        return torch.cat(ys,dim=1)+self.bias[None,:,None]
 69
 70class LongCNN(nn.Module):
 71    def __init__(self, implicit=False, m=8):
 72        super().__init__(); self.conv=FFTConv1d(1,8,15,implicit,m); self.fc=nn.Linear(8,2)
 73    def forward(self,x): return self.fc(F.relu(self.conv(x)).mean(-1))
 74
 75def run_one(seed, lr, implicit, m=8, epochs=12, return_model=False):
 76    torch.manual_seed(seed); np.random.seed(seed)
 77    d=get_dataset(seed)
 78    net=LongCNN(implicit,m)
 79    net,metric,_=train_model(net,d,epochs=epochs,lr=lr,batch=128)
 80    return (metric,net,d) if return_model else metric
 81
 82def main():
 83    # Shared union: every idea learning rate is also evaluated by baseline.
 84    lrs=[0.001,0.003,0.01]
 85    grid=[{'lr':v,'m':8} for v in lrs]
 86    base=sweep_baseline(lambda cfg: (lambda seed: run_one(seed,cfg['lr'],False,8)), grid)
 87    idea=evaluate(lambda seed: run_one(seed,base['best_cfg']['lr'],True,8))
 88    # Two nearby settings, same shared LR grid; select best idea configuration on 4 sweep seeds.
 89    idea_sweep=[]
 90    for cfg in grid:
 91        r=evaluate(lambda seed,cfg=cfg: run_one(seed,cfg['lr'],True,cfg['m']), seeds=(0,1,2,3))
 92        idea_sweep.append({'cfg':cfg,'mean':r['mean']})
 93    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']
 94    idea=evaluate(lambda seed: run_one(seed,best['lr'],True,best['m']))
 95    # Trained-model behavioral signature: compare each operator's prediction to direct torch convolution.
 96    _,ib,d=run_one(0,best['lr'],True,best['m'],return_model=True)
 97    _,bb,_=run_one(0,base['best_cfg']['lr'],False,8,return_model=True)
 98    ib.eval(); bb.eval(); x=d['xte'][:16].to(next(ib.parameters()).device)
 99    with torch.no_grad():
100        wi=ib.conv.weight; xb=x
101        direct=[]
102        for o in range(8):
103            yo=0
104            for c in range(wi.shape[1]): yo=yo+fft_conv(xb[:,c,:],wi[o:o+1,c,:],False,8)
105            direct.append(yo)
106        direct=torch.cat(direct,1)+ib.conv.bias[None,:,None]; observed=float((ib.conv(xb)-direct).norm()/direct.norm())
107        # same trained-system output comparison, not an analytic-only check
108        system_gap=float((ib(xb)-bb(xb)).norm()/bb(xb).norm())
109    sig={'prediction':'implicit convolution equals explicit convolution; tile storage ratio M/(ceil(L/m)*m)',
110         'predicted_relative_operator_error':0.0,'observed_relative_operator_error':observed,
111         'predicted_storage_ratio':128/( ((64+8-1)//8)*8 ), 'trained_system_output_gap':system_gap,
112         'confirmed': bool(observed < 1e-5)}
113    rep=make_report('long_1d_convolution','cnn_small',base,idea,{'mechanism_signature':sig,'custom_track':{'name':'long_1d_convolution','file':'stage2_bench.py','domain':'cnn'}})
114    rep['idea_sweep']=idea_sweep; rep['selection_note']='Baseline and idea share lr union [0.001, 0.003, 0.01]; idea m fixed a priori at 8.'
115    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
116    print(json.dumps(rep,indent=2))
117if __name__=='__main__': main()