1import json, math, random
  2import numpy as np
  3from scipy.optimize import linprog
  4import torch
  5import torch.nn as nn
  6
  7SEED = 2056
  8np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10
 11def fresh_seed(s=SEED):
 12    np.random.seed(s); random.seed(s); torch.manual_seed(s)
 13
 14class ICNN(nn.Module):
 15    def __init__(self, d=2, h=24, L=2):
 16        super().__init__(); self.L=L
 17        self.A=nn.ParameterList([nn.Parameter(torch.randn(h, d if k==0 else h)*.35) for k in range(L)])
 18        self.U=nn.ParameterList([nn.Parameter(torch.randn(h,d)*.35) for _ in range(L)])
 19        self.b=nn.ParameterList([nn.Parameter(torch.zeros(h)) for _ in range(L)])
 20        self.aw=nn.Parameter(torch.zeros(h)); self.u=nn.Parameter(torch.zeros(d)); self.c=nn.Parameter(torch.zeros(1))
 21    def forward(self,x):
 22        z=x
 23        for k in range(self.L):
 24            W=torch.nn.functional.softplus(self.A[k])+1e-6
 25            z=torch.relu(z@W.T + x@self.U[k].T + self.b[k])
 26        return z@torch.nn.functional.softplus(self.aw)+x@self.u+self.c
 27
 28class MLP(nn.Module):
 29    def __init__(self,d=2,h=24,L=2):
 30        super().__init__(); layers=[]; q=d
 31        for _ in range(L): layers += [nn.Linear(q,h),nn.ReLU()]; q=h
 32        layers += [nn.Linear(q,1)]
 33        self.net=nn.Sequential(*layers)
 34    def forward(self,x): return self.net(x).squeeze(-1)
 35
 36def convexity_sweep():
 37    # Prediction 1: for nonnegative W and output w, midpoint Jensen violations are zero.
 38    # Prediction 2: introducing negative hidden-to-hidden entries causes violations at a
 39    # positive rate; this is the exact architectural boundary (sign >= 0).
 40    rng=np.random.default_rng(SEED); n=30000; d=2; h=4
 41    x=rng.uniform(-2,2,(n,d)); y=rng.uniform(-2,2,(n,d)); t=rng.uniform(0,1,n)[:,None]
 42    xm=t*x+(1-t)*y
 43    rows=[]
 44    for negfrac in [0.0, .05, .10, .25, .50, 1.0]:
 45        W1=np.abs(rng.normal(size=(h,d))); W2=np.abs(rng.normal(size=(h,h)))
 46        if negfrac:
 47            mask=rng.random(W2.shape)<negfrac; W2[mask]*=-1
 48        U1=rng.normal(size=(h,d)); U2=rng.normal(size=(h,d)); b1=rng.normal(size=h); b2=rng.normal(size=h)
 49        w=np.abs(rng.normal(size=h)); u=rng.normal(size=d); c=.2
 50        def f(a):
 51            z1=np.maximum(0,a@W1.T+a@U1.T+b1); z2=np.maximum(0,z1@W2.T+a@U2.T+b2)
 52            return z2@w+a@u+c
 53        # Jensen gap f(tx+(1-t)y)-[t f(x)+(1-t)f(y)], convexity means <= 0
 54        gap=f(xm)- (t[:,0]*f(x)+(1-t[:,0])*f(y))
 55        rows.append({'negative_fraction':negfrac,'max_jensen_violation':float(max(0,gap.max())),
 56                     'violation_rate':float(np.mean(gap>1e-8)), 'mean_positive_violation':float(np.maximum(gap,0).mean())})
 57    return rows
 58
 59def controlled_sign_sweep():
 60    # Controlled prediction: f_gamma(x)=ReLU(1+gamma*ReLU(x)) is convex for
 61    # gamma>=0; for gamma<0 it has a downward slope after x=0 and must violate
 62    # Jensen convexity.  The pair (-1,1) at t=1/2 gives violation
 63    # f(0)-(f(-1)+f(1))/2 = -gamma/2 for -1<=gamma<0.
 64    xs=np.linspace(-1,1,401)
 65    rows=[]
 66    for gamma in [-2.0,-1.0,-.5,-.1,0.0,.1,.5,1.0,2.0]:
 67        def f(x): return np.maximum(0,1+gamma*np.maximum(0,x))
 68        maxv=0.0
 69        for a in xs:
 70            for b in xs[::8]:
 71                for t in (.25,.5,.75):
 72                    gap=float(f(t*a+(1-t)*b)-t*f(a)-(1-t)*f(b))
 73                    maxv=max(maxv,gap)
 74        predicted=max(0.0,-gamma/2.0) if -1<=gamma<0 else (0.5 if gamma < -1 else 0.0)
 75        rows.append({'gamma':gamma,'predicted_max_violation':predicted,'observed_max_violation':maxv})
 76    return rows
 77
 78def epigraph_check():
 79    # Prediction 3: LP epigraph minimum equals forward ReLU output when all downstream
 80    # coefficients are positive; gap should be numerical tolerance.
 81    rng=np.random.default_rng(SEED+1); d=2; h=3; N=80
 82    W=np.abs(rng.normal(size=(h,d)))*.7; U=rng.normal(size=(h,d)); b=rng.normal(size=h); w=np.abs(rng.normal(size=h))+.2
 83    gaps=[]
 84    for _ in range(N):
 85        x=rng.uniform(-2,2,d); pre=W@x+U@x+b; z=np.maximum(0,pre); forward=w@z
 86        # min w'z subject z >= pre and z >=0
 87        res=linprog(w,A_ub=-np.eye(h),b_ub=-np.maximum(pre,0),bounds=[(0,None)]*h,method='highs')
 88        gaps.append(abs(res.fun-forward))
 89    # A two-layer version, with fixed x and nonnegative W2, is also an LP.
 90    W1=np.abs(rng.normal(size=(h,d))); U1=rng.normal(size=(h,d)); b1=rng.normal(size=h)
 91    W2=np.abs(rng.normal(size=(h,h))); U2=rng.normal(size=(h,d)); b2=rng.normal(size=h); w2=np.abs(rng.normal(size=h))+.2
 92    gaps2=[]
 93    for _ in range(N):
 94        x=rng.uniform(-2,2,d); p1=W1@x+U1@x+b1; z1=np.maximum(0,p1); p2=W2@z1+U2@x+b2; z2=np.maximum(0,p2); forward=w2@z2
 95        # Variables z1,z2. Constraints -z1<=-p1; -z2+W2z1<=-(U2x+b2), plus nonnegative bounds.
 96        Aub=np.zeros((2*h,2*h)); bub=np.zeros(2*h)
 97        Aub[:h,:h]=-np.eye(h); bub[:h]=-p1
 98        Aub[h:,:h]=W2; Aub[h:,h:]=-np.eye(h); bub[h:]=-(U2@x+b2)
 99        res=linprog(np.r_[np.zeros(h),w2],A_ub=Aub,b_ub=bub,bounds=[(0,None)]*(2*h),method='highs')
100        gaps2.append(abs(res.fun-forward))
101    return {'one_layer_max_abs_gap':float(max(gaps)), 'two_layer_max_abs_gap':float(max(gaps2)),
102            'predicted_gap':0.0}
103
104def fit_compare():
105    # Convex target: quadratic. Compare ICNN with a width-matched ordinary ReLU MLP.
106    fresh_seed(); rng=np.random.default_rng(SEED)
107    X=rng.uniform(-2,2,(600,2)).astype('float32'); y=(.8*X[:,0]**2+1.4*X[:,1]**2+.35*X[:,0]*X[:,1]+.2*X[:,0]).astype('float32')
108    tr=slice(0,450); va=slice(450,None); xt=torch.tensor(X); yt=torch.tensor(y)
109    out={}
110    for name, model in [('icnn',ICNN()),('relu_mlp',MLP())]:
111        fresh_seed(SEED+ (0 if name=='icnn' else 1)); model=model.to('cpu'); opt=torch.optim.AdamW(model.parameters(),lr=.01,weight_decay=1e-5)
112        for step in range(1200):
113            opt.zero_grad(); pred=model(xt[tr]); loss=((pred-yt[tr])**2).mean(); loss.backward(); opt.step()
114        with torch.no_grad():
115            train=float(((model(xt[tr])-yt[tr])**2).mean()); val=float(((model(xt[va])-yt[va])**2).mean())
116        out[name]={'train_mse':train,'validation_mse':val,'steps':1200}
117    out['metric']='validation MSE (lower is better)'; return out
118
119def main():
120    result={'device':'cpu','predictions':{
121      'P1':'nonnegative W,w => Jensen violation <= numerical tolerance (predicted 0)',
122      'P2':'negative W entries create positive Jensen violations; violation rate rises with negative fraction',
123      'P3':'positive output coefficients make ReLU epigraph LP exact (predicted objective gap 0)'},
124      'convexity_sweep':convexity_sweep(),'controlled_sign_sweep':controlled_sign_sweep(),'epigraph_check':epigraph_check(),'fit_comparison':fit_compare()}
125    with open('results.json','w') as f: json.dump(result,f,indent=2)
126    print(json.dumps(result,indent=2))
127main()