1import itertools, math, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5import torch.nn.functional as F
  6
  7SEED=17
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
 11
 12def supports(N,S):
 13    out=[]
 14    for s in range(1,S+1):
 15        out += [(tuple(c), tuple(k)) for c in itertools.combinations(range(N),s)
 16                for k in itertools.product((1,), repeat=s)]
 17    return out
 18
 19def interaction_features(x,S):
 20    # x is {-1,+1}; q=2 Fourier characters, one per nonzero support
 21    cols=[x]
 22    N=x.shape[1]
 23    for s in range(2,S+1):
 24        for c in itertools.combinations(range(N),s): cols.append(torch.prod(x[:,c],dim=1,keepdim=True))
 25    return torch.cat(cols,1)
 26
 27def bh_penalty(w,N,S,eps=1e-8):
 28    # w includes one coefficient per feature and one output bias
 29    pos=0; val=0.
 30    for s in range(1,S+1):
 31        n=math.comb(N,s); p=2*s/(s+1)
 32        a=w[pos:pos+n]; pos += n
 33        val=val+(a.square()+eps).pow(p/2).sum().pow(1/p)
 34    return val
 35
 36class FourierModel(nn.Module):
 37    def __init__(self,N,S):
 38        super().__init__(); self.N=N; self.S=S
 39        self.w=nn.Parameter(torch.zeros(sum(math.comb(N,s) for s in range(1,S+1))))
 40        self.b=nn.Parameter(torch.zeros(1))
 41    def forward(self,x): return interaction_features(x,self.S)@self.w+self.b
 42
 43class MLP(nn.Module):
 44    def __init__(self,N,h):
 45        super().__init__(); self.net=nn.Sequential(nn.Linear(N,h),nn.Tanh(),nn.Linear(h,1))
 46    def forward(self,x): return self.net(x).squeeze(1)
 47
 48def train(model,xtr,ytr,xte,yte, penalty=False, epochs=350):
 49    opt=torch.optim.Adam(model.parameters(),lr=.03,weight_decay=0)
 50    for _ in range(epochs):
 51        opt.zero_grad(); z=model(xtr); loss=F.binary_cross_entropy_with_logits(z,ytr.float())
 52        if penalty: loss=loss+0.003*bh_penalty(model.w,model.N,model.S)
 53        loss.backward(); opt.step()
 54    with torch.no_grad():
 55        pred=(model(xte)>0).long(); acc=(pred==yte).float().mean().item()
 56        tr=(model(xtr)>0).long(); tracc=(tr==ytr).float().mean().item()
 57    return tracc,acc
 58
 59def math_check():
 60    # Exact Fourier coefficients and infinity norm on C_2^N for homogeneous order s.
 61    rows=[]
 62    for N in [3,4,5,6,7,8]:
 63        s=2; terms=list(itertools.combinations(range(N),s)); X=torch.tensor(list(itertools.product([-1.,1.],repeat=N)))
 64        # same coefficients across dimensions: normalized random draws, repeated trials
 65        ratios=[]
 66        for t in range(30):
 67            a=torch.randn(len(terms)); f=torch.zeros(len(X))
 68            for j,c in enumerate(terms): f += a[j]*torch.prod(X[:,c],1)
 69            p=4/3; q=(a.abs().pow(p).sum())**(1/p); ratios.append((q/f.abs().max()).item())
 70        rows.append({'N':N,'max_ratio':max(ratios),'median_ratio':float(np.median(ratios))})
 71    return rows
 72
 73def benchmark_one(N, S, order, ntrain, seed):
 74    X=np.array(list(itertools.product([-1.,1.],repeat=N)),dtype=np.float32)
 75    y=((np.prod(X[:,:order],axis=1))<0).astype(np.int64)
 76    rng=np.random.RandomState(seed); perm=rng.permutation(len(X)); tr,te=perm[:ntrain],perm[ntrain:]
 77    xtr=torch.tensor(X[tr],device=DEVICE); xte=torch.tensor(X[te],device=DEVICE)
 78    ytr=torch.tensor(y[tr],device=DEVICE); yte=torch.tensor(y[te],device=DEVICE)
 79    fm=FourierModel(N,S).to(DEVICE); fu=FourierModel(N,S).to(DEVICE)
 80    # width chosen to be close to the Fourier parameter count.
 81    target=sum(math.comb(N,s) for s in range(1,S+1))+1
 82    h=max(1,round((target-1-N-1)/(N+1)))
 83    mm=MLP(N,h).to(DEVICE)
 84    rp=train(fm,xtr,ytr,xte,yte,True,epochs=300)
 85    ru=train(fu,xtr,ytr,xte,yte,False,epochs=300)
 86    rm=train(mm,xtr,ytr,xte,yte,False,epochs=300)
 87    with torch.no_grad():
 88        w=fm.w.detach(); pos=0; norms={}
 89        for s in range(1,S+1):
 90            n=math.comb(N,s); p=2*s/(s+1)
 91            norms[str(s)]=float((w[pos:pos+n].abs().pow(p).sum().pow(1/p)).cpu()); pos+=n
 92    return {'order':order,'N':N,'ntrain':ntrain,'fourier_budget_params':sum(p.numel() for p in fm.parameters()),
 93            'mlp_params':sum(p.numel() for p in mm.parameters()),'mlp_width':h,
 94            'budgeted_fourier_test':rp[1],'unpenalized_fourier_test':ru[1],'mlp_test':rm[1],
 95            'budgeted_fourier_train':rp[0],'unpenalized_fourier_train':ru[0], 'mlp_train':rm[0],
 96            'budgeted_quasinorm_by_order':norms}
 97
 98def benchmark():
 99    return {'task_2way_low_order': [benchmark_one(8,2,2,64,SEED+i) for i in range(3)],
100            'task_3way_higher_order': [benchmark_one(8,3,3,64,SEED+10+i) for i in range(3)]}
101
102if __name__=='__main__':
103    # CUDA can fail on shared hardware; retry whole benchmark on CPU.
104    try:
105        result={'math_check':math_check(),'benchmark':benchmark()}
106    except Exception as e:
107        DEVICE='cpu'; result={'error_fallback':repr(e),'math_check':math_check(),'benchmark':benchmark()}
108    print(json.dumps(result,indent=2))